holy-heck-i-really-like-stats/bot.py

#!/usr/bin/env python3

from collections import deque
from discord.utils import setup_logging
from typing import Optional
import argparse
import discord
import logging
import os.path
import random
import re
import shutil
import subprocess as sp

discord.utils.setup_logging()
_log = logging.getLogger(__name__)

_GAMES = "games.txt"
_DB = "holy-heck.db"
_DB_DEST = f"/var/lib/grafana/{_DB}"


def _write_game(content: str):
    try:
        with open(_GAMES, "a") as f:
            f.write(content)
            f.write("\n")
    except:
        _log.exception(f"failed writing game {content}")


def _update_db():
    try:
        games = []
        with open(_GAMES) as f:
            for line in f:
                games.append(line.strip())
        sp.run(["./index.py", "-o", _DB] + games)
        shutil.move(_DB, _DB_DEST)
        _log.info("updated db")
    except:
        _log.exception(f"failed updating db")


class BotClient(discord.Client):
    def __init__(
        self,
        intents: discord.Intents,
        leaguefacts: Optional[list[str]] = None,
    ):
        super().__init__(intents=intents)
        self._recentfacts = deque(maxlen=min(len(leaguefacts) - 1, 5))
        self._leaguefacts = leaguefacts or []
        self._leaguefacts.append(
            f"I know {len(self._leaguefacts)} leaguefacts but you'll never see most of them."
        )
        random.seed()

    async def on_ready(self):
        _log.info(f"ready as {self.user}")

    async def on_message(self, message: discord.Message):
        content = message.content
        if self.is_command(message):
            await self.on_command(message)
        elif self.is_replay(message):
            await self.on_replay(message)
        elif self.is_leaguefact(message):
            await self.on_leaguefact(message)

    def is_command(self, message: discord.Message) -> bool:
        return message.content.startswith("%")

    async def on_command(self, message: discord.Message):
        match re.split(" +", message.content):
            case ["%calc?"], ["%calc", "?"]:
                await message.reply(content=self._help_calculate())
            case ["%calc", *args]:
                try:
                    calc = await self._calculate(args)
                    await message.reply(content=calc.strip())
                except Exception as e:
                    _log.exception("running calculation")
                    await message.reply(
                        content="```\n" + str(e) + "\n```", delete_after=10
                    )
            case _:
                _log.info(f"Unrecognised command {command}")

    async def _help_calculate(self) -> str:
        return """
            Run Showdown damage calculations.
            Example format:
            `  %calc -2 8 SpA Choice Specs Torkoal Overheat vs. 252 HP / 4+ SpD Assault Vest Abomasnow in Sun through Light Screen`
            Supported: attacker/defender boosts, EVs, item, species; attacker ability; weather/terrain; screens.
        """

    async def _calculate(self, args: list[str]) -> str:
        proc = sp.run(
            ["node", "calc_main.js", "--"] + args, stdout=sp.PIPE, stderr=sp.PIPE
        )
        if proc.returncode != 0:
            raise Exception(proc.stderr.decode())
        return proc.stdout.decode()

    def is_replay(self, message: discord.Message) -> bool:
        if re.match("https://replay.pokemonshowdown.com/dl-.*", message.content):
            return True
        return False

    async def on_replay(self, message: discord.Message):
        _log.info(f"Recognised {message.content} as a League game")
        _write_game(message.content)
        _update_db()

    def is_leaguefact(self, message: discord.Message) -> bool:
        return message.content.lower() in ["leaguefact", "leaguefacts"]

    async def on_leaguefact(self, message: discord.Message):
        _log.info("leaguefact requested")
        fact = self._select_leaguefact()
        if fact:
            await message.channel.send(f"Did you know? {fact}")
        else:
            await message.channel.send("There are no league facts.")

    def _select_leaguefact(self) -> Optional[str]:
        if not self._leaguefacts:
            return None
        choice = None
        while True:
            choice = random.choice(self._leaguefacts)
            if choice not in self._recentfacts:
                break
        self._recentfacts.append(choice)
        return choice


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-t",
        "--token-file",
        metavar="FILE",
        default="token",
        help="file containing Discord API token",
    )
    parser.add_argument(
        "-f",
        "--facts",
        metavar="FILE",
        default="facts.txt",
        help="file containing leagefacts",
    )

    args = parser.parse_args()

    facts = []
    if os.path.exists(args.facts):
        with open(args.facts) as f:
            for line in f:
                facts.append(line)

    intents = discord.Intents.default()
    intents.message_content = True
    client = BotClient(leaguefacts=facts, intents=intents)

    with open(args.token_file) as f:
        token = f.read().strip()
        client.run(token, log_handler=None)


if __name__ == "__main__":
    main()