diff --git a/bot/config.py b/bot/config.py index 89e36eda..02f706fb 100644 --- a/bot/config.py +++ b/bot/config.py @@ -97,6 +97,13 @@ class YouTube(BaseModel): role_id: int +class Chess(BaseModel): + channel_id: int + role_id: int + team_id: str + access_token: str + + class Settings(BaseSettings): aoc: AoC bot: Bot @@ -111,6 +118,7 @@ class Settings(BaseSettings): errors: ErrorHandling custom_roles: CustomRoles youtube: YouTube + chess: Chess class Config: env_file = ".env" diff --git a/bot/extensions/chess/__init__.py b/bot/extensions/chess/__init__.py new file mode 100644 index 00000000..aa656bb7 --- /dev/null +++ b/bot/extensions/chess/__init__.py @@ -0,0 +1,9 @@ +from bot.core import DiscordBot + +from .events import ChessEvents +from .tasks import ChessTasks + + +async def setup(bot: DiscordBot) -> None: + await bot.add_cog(ChessTasks(bot=bot)) + await bot.add_cog(ChessEvents(bot=bot)) diff --git a/bot/extensions/chess/events.py b/bot/extensions/chess/events.py new file mode 100644 index 00000000..606a0599 --- /dev/null +++ b/bot/extensions/chess/events.py @@ -0,0 +1,66 @@ +import asyncio +from collections import defaultdict + +import discord +from discord.ext import commands + +from bot import core +from bot.config import settings +from bot.extensions.chess.tasks import ChessTasks + + +class ChessEvents(commands.Cog): + """Events for Chess functions""" + + def __init__(self, bot: core.DiscordBot): + self.bot = bot + self.tasks: dict[int, list[tuple[asyncio.Task, discord.ScheduledEvent]]] = defaultdict(list) + + def cog_load(self) -> None: + for guild in self.bot.guilds: + for event in guild.scheduled_events: + if event.location != f"<#{settings.chess.channel_id}>": + continue + + if any(e.id == event.id for task, e in self.tasks[guild.id]): + continue + + self.tasks[guild.id].append((asyncio.create_task(ChessTasks(self.bot).run_task(event)), event)) + + @commands.Cog.listener() + async def on_scheduled_event_create(self, event: discord.ScheduledEvent): + if event.location != f"<#{settings.chess.channel_id}>": + return + + self.tasks[event.guild.id].append((asyncio.create_task(ChessTasks(self.bot).run_task(event)), event)) + + @commands.Cog.listener() + async def on_scheduled_event_delete(self, event: discord.ScheduledEvent): + if event.location != f"<#{settings.chess.channel_id}>": + return + + for task, e in self.tasks[event.guild.id]: + if e.id == event.id: + task.cancel() + self.tasks[event.guild.id].remove((task, e)) + break + + @commands.Cog.listener() + async def on_scheduled_event_update(self, before: discord.ScheduledEvent, after: discord.ScheduledEvent): + if before.location != f"<#{settings.chess.channel_id}>": + return + + for task, e in self.tasks[before.guild.id]: + if e.id == before.id: + if after.location.startswith("https://lichess.org/swiss/"): + self.tasks[before.guild.id].remove((task, e)) + break + + task.cancel() + self.tasks[before.guild.id].remove((task, e)) + + if after.location != f"<#{settings.chess.channel_id}>": + return + + self.tasks[before.guild.id].append((asyncio.create_task(ChessTasks(self.bot).run_task(after)), after)) + break diff --git a/bot/extensions/chess/tasks.py b/bot/extensions/chess/tasks.py new file mode 100644 index 00000000..29a6c4fc --- /dev/null +++ b/bot/extensions/chess/tasks.py @@ -0,0 +1,81 @@ +import asyncio +import datetime + +import discord +from discord.ext import commands + +from bot import core +from bot.config import settings +from bot.extensions.chess.utils import lichess + + +class ChessTasks(commands.Cog): + """Tasks for Chess functions""" + + def __init__(self, bot: core.DiscordBot): + self.bot = bot + self.medals = ["🥇", "🥈", "🥉"] + self.time_controls = { + "60+0": ("Blitz", "180+2", 5), + "180+2": ("Rapid", "600+0", 5), + "600+0": ("Bullet", "60+0", 7), + } + + @property + def channel(self) -> discord.TextChannel | None: + return self.bot.get_channel(settings.chess.channel_id) + + async def run_task(self, event: discord.ScheduledEvent): + """Creates the chess tournament when it's time based on the event start time.""" + + until = (event.start_time - datetime.timedelta(hours=1)).replace(tzinfo=None) + if until > datetime.datetime.utcnow(): + await asyncio.sleep((until - datetime.datetime.utcnow()).total_seconds()) + else: + return + + tournament = await lichess.get_team_tournament() + time_control = self.time_controls.get( + f"{tournament.clock.limit}+{tournament.clock.increment}", ("Blitz", "180+2", 5) + ) + time = time_control[1].split("+") + + tournament = await lichess.create_tournament( + f"Weekly {time_control[0]} Tournament", + int(event.start_time.timestamp() * 1000), + int(time[1]), + int(time[0]), + time_control[2], + ) + + await self.channel.send( + f"<@&{settings.chess.role_id}> {time_control[0]} chess tourney in less than an hour! " + "Click the link below to join!\n" + f"****\n\n" + "Prizes:\n" + "🥇 2500 pancakes 🥞\n" + "🥈 1000 pancakes 🥞\n" + "🥉 500 pancakes 🥞", + allowed_mentions=discord.AllowedMentions(roles=True), + ) + + await event.edit( + name=f"Weekly {time_control[0]} Tournament", + location=f"https://lichess.org/swiss/{tournament.id}", + ) + + await asyncio.sleep((event.start_time.replace(tzinfo=None) - datetime.datetime.utcnow()).total_seconds()) + while True: + tournament = await lichess.get_tournament(tournament.id) + if tournament.status == "finished": + break + + await asyncio.sleep(10) + + players = await lichess.get_tournament_results(tournament.id) + await self.channel.send( + "The tournament has ended! See you next time!\n\n" + "Winners:\n" + + "\n".join(f"{self.medals[i]} **{player['username']}**" for i, player in enumerate(players[:3])) + + "\n\n(If you are one of the winners, please let us know by verifying your lichess username.)" + ) diff --git a/bot/extensions/chess/utils.py b/bot/extensions/chess/utils.py new file mode 100644 index 00000000..9af08599 --- /dev/null +++ b/bot/extensions/chess/utils.py @@ -0,0 +1,74 @@ +import json + +import pydantic +from pydantic import BaseModel + +from bot.config import settings +from bot.services import http + + +class Clock(BaseModel): + increment: int + limit: int + + +class Tournament(BaseModel): + id: str + status: str + clock: Clock + + +class LichessAPI: + def __init__(self, token: str, team_id: str): + self.headers = {"Authorization": f"Bearer {token}"} + self.team = team_id + + async def get_team_tournament(self) -> Tournament: + async with http.session.get( + f"https://lichess.org/api/team/{self.team}/swiss", + headers=self.headers, + params={"max": 1}, + raise_for_status=True, + ) as resp: + try: + return Tournament.parse_obj(await resp.json(content_type="application/x-ndjson")) + except pydantic.error_wrappers.ValidationError: + return Tournament(id="0", status="finished", clock={"limit": 60, "increment": 0}) + + async def get_tournament(self, tournament_id: str) -> Tournament: + async with http.session.get( + f"https://lichess.org/api/swiss/{tournament_id}", + headers=self.headers, + raise_for_status=True, + ) as resp: + return Tournament.parse_obj(await resp.json()) + + async def get_tournament_results(self, tournament_id: str) -> list[dict]: + async with http.session.get( + f"https://lichess.org/api/swiss/{tournament_id}/results", + headers=self.headers, + params={"nb": 3}, + raise_for_status=True, + ) as resp: + results = await resp.text() + return [json.loads(result) for result in results.split("\n") if result] + + async def create_tournament( + self, name: str, start_time: int, increment: int, limit: int, rounds: int + ) -> Tournament: + async with http.session.post( + f"https://lichess.org/api/swiss/new/{self.team}", + headers=self.headers, + data={ + "name": name, + "clock.limit": limit, + "clock.increment": increment, + "nbRounds": rounds, + "startsAt": start_time, + }, + raise_for_status=True, + ) as resp: + return Tournament.parse_obj(await resp.json()) + + +lichess = LichessAPI(settings.chess.access_token, settings.chess.team_id) diff --git a/cli.py b/cli.py index 39d0b581..cb122bae 100644 --- a/cli.py +++ b/cli.py @@ -132,6 +132,7 @@ async def main(ctx): "jishaku", "bot.extensions.adventofcode", "bot.extensions.challenges", + "bot.extensions.chess", "bot.extensions.github", "bot.extensions.levelling", "bot.extensions.persistent_roles", diff --git a/example.env b/example.env index b165b9f5..c722f7ff 100644 --- a/example.env +++ b/example.env @@ -65,3 +65,9 @@ CUSTOM_ROLES__DIVIDER_ROLE_ID= YOUTUBE__CHANNEL_ID=UC4JX40jDee_tINbkjycV4Sg YOUTUBE__TEXT_CHANNEL_ID= YOUTUBE__ROLE_ID= + +# --- Chess +CHESS__CHANNEL_ID= +CHESS__ROLE_ID= +CHESS__TEAM_ID=twt-chess-club +CHESS__ACCESS_TOKEN= # https://lichess.org/account/oauth/token/create?scopes[]=tournament:write&description=TWT+Discord+Bot