import copy import json import logging import os from enum import Enum from typing import Any, Type, Dict, List import aiohttp from interactions import Extension, Permissions, SlashContext, OptionType, slash_option, component_callback, \ ComponentContext, StringSelectMenu, Message from interactions import check, is_owner, slash_command, slash_default_member_permission, listen from interactions.api.events import Startup, MessageCreate from interactions.api.events import Component from interactions import Task, IntervalTrigger from token_bot.history_manager.history_manager import HistoryManager from token_bot.persistant_database import database as pdb 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.persistant_database.user_schema import User from token_bot.persistant_database.alert_schema import Alert from token_bot.persistant_database.alert_type import AlertType from token_bot.persistant_database.alert_category import AlertCategory from token_bot.controller.users import UsersController from token_bot.controller.alerts import AlertsController from token_bot.ui.action_row.tracker import ALERT_TYPE_ROW from token_bot.ui.select_menus.region_menu import REGION_MENU from token_bot.ui.select_menus.flavor_menu import FLAVOR_MENU from token_bot.ui.select_menus.alert_menu import HIGH_ALERT_MENU, LOW_ALERT_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 = self.bot.get_user(user.user_id) alert_message = str() for alert in users_alerts[user]: alert_message += f"{alert.to_human_string()}" await discord_user.send(f"Hello, you requested to be sent an alert when the price of the World of Warcraft" f"token reaches a certain value. The following alerts have been triggered: {alert_message}") ################################### # 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") await self._users.add(ctx.member.id) await ctx.send(text, components=REGION_MENU, ephemeral=True) @slash_command() async def remove_registration(self, ctx: SlashContext): await self._users.delete(ctx.member.id) await ctx.send("All alert subscriptions and user registration deleted", ephemeral=True) @slash_command() async def delete(self, ctx: SlashContext): await self._users.delete(ctx.member.id) await ctx.send("Deletion Successful", ephemeral=True) @slash_command() async def exists(self, ctx: SlashContext): await ctx.send(str(await self._users.exists(ctx.member.id)), ephemeral=True) @slash_command() async def add_alert(self, ctx: SlashContext): if not await self._users.exists(ctx.member.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.member.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): pass @slash_command() async def list_alerts(self, ctx: SlashContext): await ctx.send(str(await self._users.list_alerts(ctx.member.id)), 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.member.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)