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

#!/usr/bin/env python3

from collections import namedtuple
from dataclasses import dataclass
from functools import partial, partialmethod
from pathlib import Path
import argparse
import json
import logging
import requests
import sqlite3
import sys
import typing as t


logging.TRACE = 5
logging.addLevelName(logging.TRACE, "TRACE")
logging.Logger.trace = partialmethod(logging.Logger.log, logging.TRACE)
logging.trace = partial(logging.log, logging.TRACE)


class LogFormatter(logging.Formatter):

    format = "%(name)s [%(levelname)s] %(message)s"
    FORMATS = {
        logging.TRACE: f"\x1b[30;20m{format}\x1b[0m",
        logging.DEBUG: f"\x1b[38;20m{format}\x1b[0m",
        logging.INFO: f"\x1b[34;20m{format}\x1b[0m",
        logging.WARNING: f"\x1b[33;20m{format}\x1b[0m",
        logging.ERROR: f"\x1b[31;20m{format}\x1b[0m",
        logging.CRITICAL: f"\x1b[31;1m{format}\x1b[0m",
    }

    def format(self, record):
        fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(fmt)
        return formatter.format(record)


APP = "hhirlstats"
LOG = logging.getLogger(APP)
_ch = logging.StreamHandler()
_ch.setFormatter(LogFormatter())
LOG.addHandler(_ch)


TEAMS = {}
_logged_teams = []


def team(player: str) -> str:
    """Maps a username to a defined team."""
    if player in TEAMS:
        return TEAMS[player]
    else:
        if not player in _logged_teams:
            LOG.warning(f"missing team mapping for {player}")
        _logged_teams.append(player)
        return player


class safelist(list):
    def get(self, index, default=None):
        try:
            return self.__getitem__(index)
        except IndexError:
            return default


def _init_db(conn: sqlite3.Connection):
    def namedtuple_factory(cursor, row):
        fields = [column[0] for column in cursor.description]
        cls = namedtuple("Row", fields)
        return cls._make(row)

    conn.row_factory = namedtuple_factory

    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS moves(
            game, turn, player, name, user, target,
            UNIQUE(game, turn, player, user)
        );
        CREATE TABLE IF NOT EXISTS switches(
            game, turn, player, name,
            UNIQUE(game, turn, player, name)
        );
        CREATE TABLE IF NOT EXISTS nicknames(
            game, player, name, specie,
            UNIQUE(game, player, specie)
        );
        CREATE TABLE IF NOT EXISTS knockouts(
            game, turn, player, name,
            UNIQUE(game, turn, player)
        );
        CREATE TABLE IF NOT EXISTS indirect_knockouts(
            game, turn, player, name, source, source_user, source_player,
            UNIQUE(game, turn, player)
        );
        CREATE TABLE IF NOT EXISTS games(
            id, p1, p2, format, uploadtime, winner,
            UNIQUE(id)
        )
        """
    )


def parse_log(game: str, log: str, into: sqlite3.Connection):
    conn = into

    turn = 0
    players = {}

    # ("p2a: Edward", "p1a: Meteo")
    # memorises the user of the move that causes environment setting or status,
    # and its target
    last_move: t.Optional[tuple[str, str]]

    # ("p1", "Spikes") => "p2a: Frosslas"
    last_env_set: dict[tuple[str, str], str] = {}

    # ("p1a: Meteo", "brn") => "p2a: Edward"
    last_status_set: dict[tuple[str, str], str] = {}

    def resolve_mon(user: str) -> tuple[str, str]:
        [player, name] = user.split(": ")
        return players[player.strip("ab")], name

    for line in log.split("\n"):
        chunks = line.split("|")[1:]
        if not chunks:
            continue

        LOG.trace(line)

        match chunks:
            case ["player", id, username, *rest]:
                players[id] = username

            case ["turn", turn]:
                turn = int(turn)

            case ["move", user, move, target]:
                last_move = (user, target)
                player, user = resolve_mon(user)
                _, target = resolve_mon(target)
                conn.execute(
                    """
                    INSERT INTO moves(game, turn, player, name, user, target)
                    VALUES (?, ?, ?, ?, ?, ?)
                    ON CONFLICT DO NOTHING
                    """,
                    (game, turn, team(player), move, user, target),
                )

            case ["switch", name, specie, *rest]:
                player, name = resolve_mon(name)
                conn.execute(
                    """
                    INSERT INTO switches(game, turn, player, name)
                    VALUES (?, ?, ?, ?)
                    ON CONFLICT DO NOTHING
                    """,
                    (game, turn, team(player), name),
                )
                conn.execute(
                    """
                    INSERT INTO nicknames(game, player, name, specie)
                    VALUES(?, ?, ?, ?)
                    ON CONFLICT DO NOTHING
                    """,
                    (game, team(player), name, specie.split(", ")[0]),
                )

            case ["faint", mon]:
                player, mon = resolve_mon(mon)
                conn.execute(
                    """
                    INSERT INTO knockouts(game, turn, player, name)
                    VALUES(?, ?, ?, ?)
                    ON CONFLICT DO NOTHING
                    """,
                    (game, turn, team(player), mon),
                )

            case ["win", player]:
                conn.execute(
                    """
                    UPDATE games
                    SET winner = ?
                    WHERE id = ?
                    """,
                    (team(player), game),
                )

            case ["-sidestart", side, env]:
                if not last_move:
                    LOG.warning(f"missing previous move for {line}")
                    continue
                LOG.debug(f"{line} <- {last_move}")
                last_env_set[(side[0:1], env.replace("move: ", ""))] = last_move[0]

            case ["-status", mon, cond]:
                if not last_move or last_move[1] != mon:
                    LOG.warning(f"missing previous move for {line}")
                    continue
                LOG.debug(f"{line} <- {last_move}")
                last_status_set[(mon, cond)] = last_move[0]

            case ["-damage", mon, "0 fnt", from_]:
                # mon has fainted from an indirect damage source
                LOG.debug(f"tracing source for {line}")
                source = from_.replace("[from] ", "")
                source_user = None

                test_hazard = last_env_set.get((mon[0:1], source))
                if test_hazard:
                    source_user = test_hazard
                    LOG.debug(f"identified hazard source {source_user}")

                test_status = last_status_set.get((mon, source))
                if test_status:
                    source_user = test_status
                    LOG.debug(f"identified move source {source_user}")

                if source == "Recoil" or source.startswith("item: "):
                    LOG.debug(f"identified special source {source}")
                    source = source.replace("item: ", "")
                    source_user = "self"

                if not source_user:
                    LOG.error(f"missing source for {line}")
                    continue

                player, mon = resolve_mon(mon)
                if source_user.startswith("p1") or source_user.startswith("p2"):
                    source_player, source_user = resolve_mon(source_user)
                else:
                    source_player = None

                conn.execute(
                    """
                    INSERT INTO indirect_knockouts(game, turn, player, name, source, source_user, source_player)
                    VALUES(?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT DO NOTHING
                    """,
                    (
                        game,
                        turn,
                        team(player),
                        mon,
                        source,
                        source_user,
                        team(source_player),
                    ),
                )

            case _:
                # LOG.debug(f"unhandled message {chunks[0]}")
                pass


@dataclass(frozen=True)
class Replay:
    id: str
    p1: str
    p2: str
    format: str
    log: str
    uploadtime: int
    views: int
    p1id: str
    p2id: str
    formatid: str
    rating: int
    private: int
    password: t.Optional[str]


def fetch(replay: str, cache: bool = True) -> Replay:
    replay = replay.replace("https://replay.pokemonshowdown.com/", "")
    replay_file = Path.cwd() / "cache" / f"{replay}.json"

    if cache and replay_file.exists():
        with replay_file.open() as f:
            return Replay(**json.load(f))

    data = requests.get(f"https://replay.pokemonshowdown.com/{replay}.json")
    if data.status_code != 200:
        raise Exception(data.text)
    data = data.json()

    if cache:
        replay_file.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
        with replay_file.open(mode="w") as f:
            json.dump(data, f)

    return Replay(**data)


def main(args):
    parser = argparse.ArgumentParser(
        prog=APP,
        description="extracts stats from a Showdown replay",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("-v", "--verbose", action="count", help="add debugging info")
    parser.add_argument(
        "-C",
        "--no-cache",
        action="store_true",
        help="fetch replays instead of using cache",
    )
    parser.add_argument(
        "-t",
        "--teams",
        action="store",
        metavar="FILE",
        default="teams.json",
        help="JSON file defining players to teams",
    )
    parser.add_argument(
        "-o",
        "--output",
        action="store",
        metavar="FILE",
        default="data.db",
        help="output data file",
    )
    parser.add_argument("replay", nargs="+", help="replay ID or URL")

    args = parser.parse_args(args)
    if args.verbose and args.verbose > 1:
        LOG.setLevel(logging.TRACE)
    elif args.verbose:
        LOG.setLevel(logging.DEBUG)

    if args.teams:
        with open(args.teams) as f:
            global TEAMS
            TEAMS = json.load(f)

    try:
        db = sqlite3.connect(args.output)
        _init_db(db)

        for r in args.replay:
            try:
                replay = fetch(r, cache=not args.no_cache)
            except Exception as e:
                LOG.error(f"bad replay {r}")
                continue

            LOG.info(f"indexing game {replay.id}")
            db.execute(
                """
                INSERT INTO games(id, p1, p2, format, uploadtime)
                VALUES (?, ?, ?, ?, ?)
                ON CONFLICT DO NOTHING
                """,
                (
                    replay.id,
                    team(replay.p1),
                    team(replay.p2),
                    replay.format,
                    replay.uploadtime,
                ),
            )

            parse_log(replay.id, replay.log, into=db)
            db.commit()

    finally:
        db.close()


if __name__ == "__main__":
    main(sys.argv[1:])