wow-token-app-bot/token_bot/tracker.py

import copy
import logging
from typing import Type, Dict, List

import aiohttp
from interactions import Extension, SlashContext, component_callback, \
    ComponentContext, StringSelectMenu, Message, Embed
from interactions import Task, IntervalTrigger
from interactions import slash_command, listen
from interactions.api.events import Component
from interactions.api.events import Startup

from token_bot.controller.alerts import AlertsController
from token_bot.controller.users import UsersController
from token_bot.history_manager.history_manager import HistoryManager
from token_bot.persistant_database.alert_category import AlertCategory
from token_bot.persistant_database.alert_schema import Alert
from token_bot.persistant_database.alert_type import AlertType
from token_bot.persistant_database.user_schema import User
from token_bot.token_database import database as tdb
from token_bot.token_database.flavor import Flavor
from token_bot.token_database.region import Region
from token_bot.ui.action_row.tracker import ALERT_TYPE_ROW
from token_bot.ui.select_menus.alert_menu import HIGH_ALERT_MENU, LOW_ALERT_MENU
from token_bot.ui.select_menus.flavor_menu import FLAVOR_MENU
from token_bot.ui.select_menus.region_menu import REGION_MENU


class Tracker(Extension):
    def __init__(self, bot):
        self._users: UsersController | None = None
        self._alerts: AlertsController | None = None
        self._tdb: tdb.Database | None = None
        self._history_manager: HistoryManager | None = None


    ###################################
    # Task Functions                  #
    ###################################

    @Task.create(IntervalTrigger(minutes=1))
    async def update_data(self):
        self.bot.logger.log(logging.INFO, "TokenBot Tracker: Updating Price")
        users_alerts: Dict[User, List[Alert]] = {}
        for flavor in Flavor:
            for region in Region:
                alerts = await self._history_manager.update_data(flavor, Region(region))
                for alert in alerts:
                    users = await self._alerts.get_users(alert, consistent=True)
                    for user in users:
                        if user not in users_alerts:
                            users_alerts[user] = [alert]
                        else:
                            users_alerts[user].append(alert)
        for user in users_alerts:
            discord_user = await self.bot.fetch_user(user.user_id)
            alert_message = str()
            for alert in users_alerts[user]:
                if alert.alert_type != AlertType.SPECIFIC_PRICE:
                    alert_message += f"{alert.to_human_string()}"
            embed = Embed(
                title="TokenBot Tracker Alert Triggered",
                color=0xb10000,
                description=f"Hello, you requested to be sent an alert when the price of the World of Warcraft "
                                    f"token reaches a certain value.\n\n"
                                    f"As a reminder, you can remove an alert via /remove-alert\n"
                                    f"or you can remove all registrations via /remove-registration\n\n"
                                    f"The following alerts have been triggered: {alert_message}",
            )
            await discord_user.send(embed=embed)


    ###################################
    # Slash Commands                  #
    ###################################

    @listen(Startup)
    async def on_start(self):
        self.bot.logger.log(logging.INFO, "TokenBot Tracker: Initializing")
        self._users = UsersController(aiohttp.ClientSession())
        self._alerts = AlertsController(aiohttp.ClientSession())
        self._tdb = tdb.Database(aiohttp.ClientSession())
        self._history_manager = HistoryManager(self._tdb)
        self.bot.logger.log(logging.INFO, "TokenBot Tracker: Initialized")
        self.bot.logger.log(logging.INFO, "TokenBot Tracker: Loading Historical Data")
        await self._history_manager.load_data()
        self.bot.logger.log(logging.INFO, "TokenBot Tracker: Loading Historical Data Finished")
        self.bot.logger.log(logging.INFO, "TokenBot Tracker: Started")
        self.update_data.start()


    @slash_command()
    async def register(self, ctx: SlashContext):
        text = ("## Select a region to register with \n\n"
                "Please note: \n"
                "* You can only be registered with one region at a time \n"
                "* Changing your region will remove all previous alerts you have signed up for \n"
                "* You can remove all alerts and registration using /remove-registration")
        await self._users.add(ctx.user.id)
        await ctx.send(text, components=REGION_MENU, ephemeral=True)

    @slash_command()
    async def remove_registration(self, ctx: SlashContext):
        await self._users.delete(ctx.user.id)
        await ctx.send("All alert subscriptions and user registration deleted", ephemeral=True)

    @slash_command()
    async def exists(self, ctx: SlashContext):
        await ctx.send(str(await self._users.exists(ctx.user.id)), ephemeral=True)

    @slash_command()
    async def add_alert(self, ctx: SlashContext):
        if not await self._users.exists(ctx.user.id):
            await ctx.send("You are not registered with any region\n"
                           "Please register with /register before adding alerts",
                           ephemeral=True)
            return
        user = await self._users.get(ctx.user.id)

        try:
            flavor = await self.flavor_select_menu(ctx)
            alert_category = await self.alert_category_select_menu(ctx)
            match alert_category:
                case AlertCategory.LOW:
                    alert_type = await self.low_alert_select_menu(ctx)
                case AlertCategory.HIGH:
                    alert_type = await self.high_alert_select_menu(ctx)
                case _:
                    raise NotImplementedError

        except TimeoutError:
            return

        else:
            alert = Alert(alert_type, flavor, user.region)
            if not await self._users.is_subscribed(user, alert):
                await self._users.add_alert(user, alert)
                await self._alerts.add_user(user, alert)
                await ctx.send("Successfully added alert", ephemeral=True)

            else:
                await ctx.send("You are already subscribed to this alert", ephemeral=True)


    @slash_command()
    async def remove_alert(self, ctx: SlashContext):
        await ctx.send("This is not implemented yet, use /remove-registration for the time being", ephemeral=True)

    @slash_command()
    async def list_alerts(self, ctx: SlashContext):
        alerts = await self._users.list_alerts(ctx.user.id)
        alerts_str = str()
        for alert in alerts:
            alerts_str += f"{alert.to_human_string()}\n"
        await ctx.send(str(alerts), ephemeral=True)

    ###################################
    # Callbacks Commands              #
    ###################################

    # TODO: export to a separate file if possible

    @component_callback('flavor_menu')
    async def flavor_menu(self, ctx: ComponentContext):
        await ctx.send(f"Selected Flavor: {ctx.values[0]}", ephemeral=True)

    @component_callback('high_alert_menu')
    async def alert_menu(self, ctx: ComponentContext):
        await ctx.send(f"Selected Alert: {ctx.values[0]}", ephemeral=True)

    @component_callback('low_alert_menu')
    async def alert_menu(self, ctx: ComponentContext):
        await ctx.send(f"Selected Alert: {ctx.values[0]}", ephemeral=True)

    @component_callback('region_menu')
    async def region_menu(self, ctx: ComponentContext):
        user = User(ctx.user.id, Region(ctx.values[0].lower()), subscribed_alerts=[])
        await self._users.add(user)
        await ctx.send(f"Successfully registered with the {ctx.values[0]} region", ephemeral=True)

    @component_callback('high_alert_button')
    async def high_alert_button(self, ctx: ComponentContext):
        await ctx.send("You selected to add a High Price Alert", ephemeral=True)

    @component_callback('low_alert_button')
    async def low_alert_button(self, ctx: ComponentContext):
        await ctx.send("You selected to add a Low Price Alert", ephemeral=True)

    @component_callback('custom_alert_button')
    async def custom_alert_button(self, ctx: ComponentContext):
        await ctx.send("You selected to add a Custom Price Alert", ephemeral=True)

    ###################################
    # Helper Functions                #
    ###################################

    async def flavor_select_menu(self, ctx: SlashContext) -> Type[Flavor]:
        flavor_menu = copy.deepcopy(FLAVOR_MENU)

        flavor_message = await ctx.send(
            "Select a flavor to add alerts for",
            components=flavor_menu,
            ephemeral=True)
        try:
            flavor_component: Component = await self.bot.wait_for_component(messages=flavor_message,
                                                                            components=flavor_menu, timeout=30)
        except TimeoutError:
            flavor_menu.disabled = True
            await flavor_message.edit(context=ctx, components=flavor_menu, content="Timed out")
            raise TimeoutError
        else:
            flavor = Flavor[flavor_component.ctx.values[0].upper()]
            flavor_menu.disabled = True
            await flavor_message.edit(context=ctx, components=flavor_menu)
            return flavor

    async def alert_category_select_menu(self, ctx: SlashContext) -> AlertCategory:
        alert_type_button = copy.deepcopy(ALERT_TYPE_ROW)
        alert_type_message = await ctx.send(
            "Select an alert type to add",
            components=alert_type_button,
            ephemeral=True)
        try:
            alert_type_component: Component = await self.bot.wait_for_component(messages=alert_type_message,
                                                                                components=alert_type_button,
                                                                                timeout=30)
        except TimeoutError:
            for button in alert_type_button[0].components:
                button.disabled = True
            await alert_type_message.edit(context=ctx, components=alert_type_button, content="Timed out")
            raise TimeoutError
        else:
            alert_type = AlertCategory.from_str(alert_type_component.ctx.custom_id)
            for button in alert_type_button[0].components:
                button.disabled = True
            await alert_type_message.edit(context=ctx, components=alert_type_button)
            return alert_type


    async def _alert_select_menu_handler(self, ctx: SlashContext, menu: StringSelectMenu, message: Message) -> AlertType:
        try:
            component: Component = await self.bot.wait_for_component(messages=message, components=menu, timeout=30)
        except TimeoutError:
            menu.disabled = True
            await message.edit(context=ctx, components=menu, content="Timed out")
            raise TimeoutError
        else:
            menu.disabled = True
            await message.edit(context=ctx, components=menu)
            return AlertType.from_str(component.ctx.values[0])

    async def high_alert_select_menu(self, ctx: SlashContext) -> AlertType:
        high_menu = copy.deepcopy(HIGH_ALERT_MENU)
        high_message = await ctx.send(
            "Select a time range to add a High Alert for",
            components=high_menu,
            ephemeral=True)
        return await self._alert_select_menu_handler(ctx, high_menu, high_message)

    async def low_alert_select_menu(self, ctx: SlashContext) -> AlertType:
        low_menu = copy.deepcopy(LOW_ALERT_MENU)
        low_message = await ctx.send(
            "Select a time range to add a Low Alert for",
            components=low_menu,
            ephemeral=True)
        return await self._alert_select_menu_handler(ctx, low_menu, low_message)