qrm2/exts/study.py

"""
Study extension for qrm
---
Copyright (C) 2019-2021 classabbyamp, 0x5c

SPDX-License-Identifier: LiLiQ-Rplus-1.1
"""


import random
import json
from datetime import datetime
import asyncio

import aiohttp

import discord.ext.commands as commands

import common as cmn
from resources import study


class StudyCog(commands.Cog):
    choices = {"A": cmn.emojis.a, "B": cmn.emojis.b, "C": cmn.emojis.c, "D": cmn.emojis.d, "E": cmn.emojis.e}
    choices_inv = {y: x for x, y in choices.items()}

    def __init__(self, bot: commands.Bot):
        self.bot = bot
        self.lastq = dict()
        self.source = "Data courtesy of [HamStudy.org](https://hamstudy.org/)"
        self.session = aiohttp.ClientSession(connector=bot.qrm.connector)

    @commands.command(name="hamstudy", aliases=["rq", "randomquestion", "randomq"], category=cmn.Cats.STUDY)
    async def _random_question(self, ctx: commands.Context, country: str = "", level: str = "", element: str = ""):
        """Gets a random question from [HamStudy's](https://hamstudy.org) question pools."""
        with ctx.typing():
            embed = cmn.embed_factory(ctx)

            country = country.lower()
            level = level.lower()
            element = element.upper()

            if country in study.pool_names.keys():
                if level in study.pool_names[country].keys():
                    pool_name = study.pool_names[country][level]

                elif level in ("random", "r"):
                    # select a random level in that country
                    pool_name = random.choice(list(study.pool_names[country].values()))

                else:
                    # show list of possible pools
                    embed.title = "Pool Not Found!"
                    embed.description = "Possible arguments are:"
                    embed.colour = cmn.colours.bad
                    for cty in study.pool_names:
                        levels = "`, `".join(study.pool_names[cty].keys())
                        embed.add_field(name=f"**Country: `{cty}` {study.pool_emojis[cty]}**",
                                        value=f"Levels: `{levels}`", inline=False)
                    embed.add_field(name="**Random**", value="To select a random pool or country, use `random` or `r`")
                    await ctx.send(embed=embed)
                    return

            elif country in ("random", "r"):
                # select a random country and level
                country = random.choice(list(study.pool_names.keys()))
                pool_name = random.choice(list(study.pool_names[country].values()))

            else:
                # show list of possible pools
                embed.title = "Pool Not Found!"
                embed.description = "Possible arguments are:"
                embed.colour = cmn.colours.bad
                for cty in study.pool_names:
                    levels = "`, `".join(study.pool_names[cty].keys())
                    embed.add_field(name=f"**Country: `{cty}` {study.pool_emojis[cty]}**",
                                    value=f"Levels: `{levels}`", inline=False)
                embed.add_field(name="**Random**", value="To select a random pool or country, use `random` or `r`")
                await ctx.send(embed=embed)
                return

            pools = await self.hamstudy_get_pools()

            pool_matches = [p for p in pools.keys() if "_".join(p.split("_")[:-1]) == pool_name]

            if len(pool_matches) > 0:
                if len(pool_matches) == 1:
                    pool = pool_matches[0]
                else:
                    # look at valid_from and expires dates to find the correct one
                    for p in pool_matches:
                        valid_from = datetime.fromisoformat(pools[p]["valid_from"][:-1])
                        expires = datetime.fromisoformat(pools[p]["expires"][:-1])

                        if valid_from < datetime.utcnow() < expires:
                            pool = p
                            break
            else:
                # show list of possible pools
                embed.title = "Pool Not Found!"
                embed.description = "Possible arguments are:"
                embed.colour = cmn.colours.bad
                for cty in study.pool_names:
                    levels = "`, `".join(study.pool_names[cty].keys())
                    embed.add_field(name=f"**Country: `{cty}` {study.pool_emojis[cty]}**",
                                    value=f"Levels: `{levels}`", inline=False)
                embed.add_field(name="**Random**", value="To select a random pool or country, use `random` or `r`")
                await ctx.send(embed=embed)
                return

            pool_meta = pools[pool]

            async with self.session.get(f"https://hamstudy.org/pools/{pool}") as resp:
                if resp.status != 200:
                    raise cmn.BotHTTPError(resp)
                pool = json.loads(await resp.read())["pool"]

            # Select a question
            if element:
                els = [el["id"] for el in pool]
                if element in els:
                    pool_section = pool[els.index(element)]["sections"]
                else:
                    embed.title = "Element Not Found!"
                    embed.description = f"Possible Elements for Country `{country}` and Level `{level}` are:"
                    embed.colour = cmn.colours.bad
                    embed.description += "\n\n" + "`" + "`, `".join(els) + "`"
                    await ctx.send(embed=embed)
                    return
            else:
                pool_section = random.choice(pool)["sections"]
            pool_questions = random.choice(pool_section)["questions"]
            question = random.choice(pool_questions)
            answers = question['answers']
            answers_str = ""
            answers_str_bolded = ""
            for letter, ans in answers.items():
                answers_str += f"{self.choices[letter]} {ans}\n"
                if letter == question["answer"]:
                    answers_str_bolded += f"{self.choices[letter]} **{ans}**\n"
                else:
                    answers_str_bolded += f"{self.choices[letter]} {ans}\n"

            embed.title = f"{study.pool_emojis[country]} {pool_meta['class']} {question['id']}"
            embed.description = self.source
            embed.add_field(name="Question", value=question["text"], inline=False)
            embed.add_field(name="Answers", value=answers_str, inline=False)
            embed.add_field(name="To Answer",
                            value=("Answer with reactions below. If not answered within 5 minutes,"
                                   " the answer will be revealed."),
                            inline=False)
            if "image" in question:
                image_url = f"https://hamstudy.org/images/{pool_meta['year']}/{question['image']}"
                embed.set_image(url=image_url)

        q_msg = await ctx.send(embed=embed)

        for i in range(len(answers)):
            await cmn.add_react(q_msg, list(self.choices.values())[i])
        await cmn.add_react(q_msg, cmn.emojis.question)

        def check(reaction, user):
            return (user.id != self.bot.user.id
                    and reaction.message.id == q_msg.id
                    and (str(reaction.emoji) in self.choices.values() or str(reaction.emoji) == cmn.emojis.question))

        try:
            reaction, user = await self.bot.wait_for("reaction_add", timeout=300.0, check=check)
        except asyncio.TimeoutError:
            embed.set_field_at(1, name="Answers", value=answers_str_bolded, inline=False)
            embed.set_field_at(2, name="Answer",
                               value=(f"{cmn.emojis.stopwatch} "
                                      f"**Timed out!** The correct answer was {self.choices[question['answer']]}"))
            embed.colour = cmn.colours.timeout
            await q_msg.edit(embed=embed)
        else:
            if str(reaction.emoji) == cmn.emojis.question:
                embed.set_field_at(1, name="Answers", value=answers_str_bolded, inline=False)
                embed.set_field_at(2, name="Answer",
                                   value=f"The correct answer was {self.choices[question['answer']]}", inline=False)
                embed.add_field(name="Answer Requested By", value=str(user), inline=False)
                embed.colour = cmn.colours.timeout
                await q_msg.edit(embed=embed)
            else:
                answers_str_checked = ""
                chosen_ans = self.choices_inv[str(reaction.emoji)]
                for letter, ans in answers.items():
                    answers_str_checked += f"{self.choices[letter]}"
                    if letter == question["answer"] == chosen_ans:
                        answers_str_checked += f"{cmn.emojis.check_mark} **{ans}**\n"
                    elif letter == question["answer"]:
                        answers_str_checked += f" **{ans}**\n"
                    elif letter == chosen_ans:
                        answers_str_checked += f"{cmn.emojis.x} {ans}\n"
                    else:
                        answers_str_checked += f" {ans}\n"

                if self.choices[question["answer"]] == str(reaction.emoji):
                    embed.set_field_at(1, name="Answers", value=answers_str_checked, inline=False)
                    embed.set_field_at(2, name="Answer", value=(f"{cmn.emojis.check_mark} "
                                                                f"**Correct!** The answer was {reaction.emoji}"))
                    embed.add_field(name="Answered By", value=str(user), inline=False)
                    embed.colour = cmn.colours.good
                    await q_msg.edit(embed=embed)
                else:
                    embed.set_field_at(1, name="Answers", value=answers_str_checked, inline=False)
                    embed.set_field_at(2, name="Answer",
                                       value=(f"{cmn.emojis.x} **Incorrect!** The correct answer was "
                                              f"{self.choices[question['answer']]}, not {reaction.emoji}"))
                    embed.add_field(name="Answered By", value=str(user), inline=False)
                    embed.colour = cmn.colours.bad
                    await q_msg.edit(embed=embed)

    async def hamstudy_get_pools(self):
        async with self.session.get("https://hamstudy.org/pools/") as resp:
            if resp.status != 200:
                raise cmn.BotHTTPError(resp)
            else:
                pools_dict = json.loads(await resp.read())

        pools = dict()
        for ls in pools_dict.values():
            for pool in ls:
                pools[pool["id"]] = pool

        return pools


def setup(bot: commands.Bot):
    bot.add_cog(StudyCog(bot))