Emily Doherty
5db4e76de8
This is so interactions will work in a DM with the bot as well as in servers where the bot is located
270 lines
12 KiB
Python
270 lines
12 KiB
Python
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.\n\n"
|
|
f"As a reminder, you can always 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}")
|
|
|
|
|
|
###################################
|
|
# 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.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 delete(self, ctx: SlashContext):
|
|
await self._users.delete(ctx.user.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.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):
|
|
pass
|
|
|
|
@slash_command()
|
|
async def list_alerts(self, ctx: SlashContext):
|
|
alerts = await self._users.list_alerts(ctx.user.id)
|
|
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)
|