#!/usr/bin/env python3
"""Show running games in a sandboxed Discord as Rich Presence, using each game's official Discord app ID."""

import argparse
import json
import logging
import os
import select
import signal
import socket
import struct
import sys
import threading
import time
import tomllib
import urllib.request
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any

__version__ = "0.2"

NAME = "discord-game-presence"
DETECTABLE_URL = "https://discord.com/api/v9/applications/detectable"
CACHE_MAX_AGE = 7 * 24 * 3600  # seconds before the game list is downloaded again
CACHE_CHECK = 3600  # seconds between cache age checks in the daemon
CACHE_RETRY = 600  # seconds between download attempts while there is no game list at all
SOCKET_DIRS = ("", "app/com.discordapp.Discord", "snap.discord")  # under XDG_RUNTIME_DIR: host/firejail, Flatpak, Snap

OP_HANDSHAKE, OP_FRAME, OP_CLOSE, OP_PING, OP_PONG = range(5)

log = logging.getLogger(NAME)


@dataclass(frozen=True)
class App:
    id: str
    name: str


@dataclass(frozen=True)
class Executable:
    name: str  # lowercase path suffix with '/' separators
    app: App
    arguments: str | None = None  # lowercase text that must appear in the arguments


@dataclass(frozen=True)
class Game:
    app: App
    pid: int
    start: int  # process start, unix seconds


@dataclass(frozen=True)
class Config:
    interval: float = 5.0
    linux_executables: bool = False
    ignore: frozenset[str] = frozenset()  # app IDs or lowercase app names
    executables: dict[str, str] = field(default_factory=dict)  # normalized path suffix -> app ID


def normalize(path: str) -> str:
    return path.lower().replace("\\", "/")


def load_config(path: Path) -> Config:
    """Read the TOML config; a missing file gives the defaults. Raises ValueError on bad content."""
    try:
        with path.open("rb") as f:
            data = tomllib.load(f)
    except FileNotFoundError:
        return Config()

    unknown = set(data) - {"interval", "linux_executables", "ignore", "executables"}
    if unknown:
        raise ValueError(f"unknown keys: {', '.join(sorted(unknown))}")
    interval = data.get("interval", Config.interval)
    if isinstance(interval, bool) or not isinstance(interval, int | float) or interval <= 0:
        raise ValueError("interval must be a positive number")
    linux = data.get("linux_executables", False)
    if not isinstance(linux, bool):
        raise ValueError("linux_executables must be true or false")
    ignore = data.get("ignore", [])
    if not isinstance(ignore, list) or not all(isinstance(x, str) for x in ignore):
        raise ValueError("ignore must be a list of app IDs or names")
    exes = data.get("executables", {})
    if not isinstance(exes, dict) or not all(isinstance(v, str) and v.isdigit() for v in exes.values()):
        raise ValueError('executables must map names to app IDs, e.g. "game.exe" = "123456789"')
    return Config(
        interval=float(interval),
        linux_executables=linux,
        ignore=frozenset(x.lower() for x in ignore),
        executables={normalize(k): v for k, v in exes.items()},
    )


class GameIndex:
    """Lookup tables built from Discord's detectable list."""

    def __init__(self, apps: list[dict[str, Any]], config: Config) -> None:
        self.steam: dict[str, App] = {}
        self.extra: dict[str, list[Executable]] = {}
        self.by_base: dict[str, list[Executable]] = {}
        systems = {"win32", "linux"} if config.linux_executables else {"win32"}
        by_id: dict[str, App] = {}

        for entry in apps:
            app_id, name = entry.get("id"), entry.get("name")
            if not app_id or not isinstance(name, str):
                continue
            app = App(str(app_id), name)
            by_id[app.id] = app
            if app.id in config.ignore or name.lower() in config.ignore:
                continue
            for sku in entry.get("third_party_skus") or []:
                if sku.get("distributor") == "steam" and sku.get("id"):
                    self.steam.setdefault(str(sku["id"]), app)
            for exe in entry.get("executables") or []:
                if exe.get("is_launcher") or exe.get("os") not in systems or not exe.get("name"):
                    continue
                args = exe.get("arguments")
                exe_name = normalize(exe["name"]).lstrip(">")
                self._add(self.by_base, Executable(exe_name, app, args.lower() if args else None))

        for exe_name, app_id in config.executables.items():
            self._add(self.extra, Executable(exe_name, by_id.get(app_id, App(app_id, app_id))))

    @staticmethod
    def _add(table: dict[str, list[Executable]], exe: Executable) -> None:
        table.setdefault(exe.name.rsplit("/", 1)[-1], []).append(exe)

    def match_steam(self, argv: list[str]) -> App | None:
        """Steam's `reaper SteamLaunch AppId=N -- <game>` lives as long as the game."""
        if normalize(argv[0]).rsplit("/", 1)[-1] != "reaper" or "SteamLaunch" not in argv:
            return None
        for arg in argv[1:]:
            if arg == "--":
                break
            if arg.startswith("AppId="):
                return self.steam.get(arg.removeprefix("AppId="))
        return None

    def match_exe(self, argv: list[str]) -> App | None:
        """Match argv[0] against path suffixes at a '/' boundary; the longest match wins, user mappings first."""
        path = normalize(argv[0])
        base = path.rsplit("/", 1)[-1]
        args = normalize(" ".join(argv[1:]))
        for table in (self.extra, self.by_base):
            best: Executable | None = None
            for exe in table.get(base, []):
                if path != exe.name and not path.endswith("/" + exe.name):
                    continue
                if exe.arguments is not None and exe.arguments not in args:
                    continue
                if best is None or len(exe.name) > len(best.name):
                    best = exe
            if best is not None:
                return best.app
        return None

    def match(self, argv: list[str]) -> App | None:
        return self.match_steam(argv) or self.match_exe(argv)


def read_argv(pid_dir: Path) -> list[str]:
    """Return a process's argv; Wine pads it with trailing NULs."""
    args = (pid_dir / "cmdline").read_bytes().decode(errors="replace").split("\0")
    while args and not args[-1]:
        args.pop()
    return args


def boot_time(proc: Path) -> int:
    for line in (proc / "stat").read_text().splitlines():
        if line.startswith("btime "):
            return int(line.split()[1])
    raise ValueError("no btime in /proc/stat")


def start_time(pid_dir: Path, btime: int) -> int:
    # comm (field 2) may contain spaces and ')', so split after the last ')'; starttime is field 22
    stat = (pid_dir / "stat").read_text()
    ticks = int(stat[stat.rindex(")") + 2 :].split()[19])
    return btime + ticks // os.sysconf("SC_CLK_TCK")


def scan(index: GameIndex, proc: Path = Path("/proc")) -> list[Game]:
    """Running games, one per app (its oldest process), oldest first."""
    btime = boot_time(proc)
    games: dict[str, Game] = {}
    for pid_dir in proc.iterdir():
        if not pid_dir.name.isdigit():
            continue
        try:
            argv = read_argv(pid_dir)
            app = index.match(argv) if argv else None
            if app is None:
                continue
            game = Game(app, int(pid_dir.name), start_time(pid_dir, btime))
        except (OSError, ValueError, IndexError):
            continue  # the process exited mid-read
        old = games.get(app.id)
        if old is None or (game.start, game.pid) < (old.start, old.pid):
            games[app.id] = game
    return sorted(games.values(), key=lambda g: (g.start, g.pid))


def choose(games: list[Game], current: Game | None) -> Game | None:
    """Keep showing the current game while it runs, timed from its first process, else the oldest one."""
    if current is not None:
        for game in games:
            if game.app.id == current.app.id:
                return replace(game, start=min(game.start, current.start))
    return games[0] if games else None


def download() -> bytes:
    req = urllib.request.Request(DETECTABLE_URL, headers={"User-Agent": f"{NAME}/{__version__}"})
    with urllib.request.urlopen(req, timeout=60) as resp:
        data: bytes = resp.read()
    return data


def refresh_cache(path: Path, fetch: Callable[[], bytes] = download, max_age: float = CACHE_MAX_AGE) -> bool:
    """Download the game list if the cache is missing or old; keep the old one on failure. True if replaced."""
    try:
        if time.time() - path.stat().st_mtime < max_age:
            return False
    except FileNotFoundError:
        pass
    try:
        data = fetch()
        if not isinstance(json.loads(data), list):
            raise ValueError("not a JSON list")
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(".tmp")
        tmp.write_bytes(data)
        tmp.replace(path)
    except (OSError, ValueError) as e:
        log.warning("Can't download the game list: %s", e)
        return False
    log.info("Downloaded the game list")
    return True


def read_cache(path: Path) -> list[dict[str, Any]] | None:
    try:
        apps = json.loads(path.read_bytes())
    except FileNotFoundError:
        return None
    except (OSError, ValueError) as e:
        log.warning("Can't read %s: %s", path, e)
        return None
    if not isinstance(apps, list):
        log.warning("%s is not a JSON list", path)
        return None
    return [a for a in apps if isinstance(a, dict)]


class IPCError(Exception):
    pass


class DiscordIPC:
    """One Rich Presence connection; the client (app) ID is fixed per connection."""

    def __init__(self, sock: socket.socket) -> None:
        self.sock = sock

    @classmethod
    def connect(cls, client_id: str, runtime_dir: Path, timeout: float = 5.0) -> "DiscordIPC":
        error: OSError = FileNotFoundError(f"no discord-ipc socket in {runtime_dir}")
        for path in (runtime_dir / d / f"discord-ipc-{i}" for d in SOCKET_DIRS for i in range(10)):
            if not path.exists():
                continue
            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            try:
                sock.connect(str(path))
            except OSError as e:
                sock.close()
                error = e
                continue
            ipc = cls(sock)
            try:
                ipc._handshake(client_id)
            except BaseException:
                ipc.close()
                raise
            return ipc
        raise error

    def _send(self, op: int, payload: dict[str, Any]) -> None:
        data = json.dumps(payload).encode()
        self.sock.sendall(struct.pack("<II", op, len(data)) + data)

    def _recv_exact(self, size: int) -> bytes:
        buf = b""
        while len(buf) < size:
            chunk = self.sock.recv(size - len(buf))
            if not chunk:
                raise ConnectionError("Discord closed the connection")
            buf += chunk
        return buf

    def _recv(self) -> tuple[int, dict[str, Any]]:
        op, size = struct.unpack("<II", self._recv_exact(8))
        payload = json.loads(self._recv_exact(size)) if size else {}
        if op == OP_CLOSE:
            raise IPCError(f"Discord closed the connection: {payload.get('message', payload)}")
        if op == OP_PING:
            self._send(OP_PONG, payload)
        return op, payload

    def _handshake(self, client_id: str) -> None:
        self._send(OP_HANDSHAKE, {"v": 1, "client_id": client_id})
        _, payload = self._recv()
        if payload.get("evt") != "READY":
            raise IPCError(f"unexpected handshake reply: {payload}")

    def set_activity(self, start: int) -> None:
        nonce = str(uuid.uuid4())
        args = {"pid": os.getpid(), "activity": {"timestamps": {"start": start}}}
        self._send(OP_FRAME, {"cmd": "SET_ACTIVITY", "args": args, "nonce": nonce})
        while True:
            op, payload = self._recv()
            if op == OP_FRAME and payload.get("nonce") == nonce:
                break
        if payload.get("evt") == "ERROR":
            raise IPCError(payload.get("data", {}).get("message", str(payload)))

    def alive(self) -> bool:
        """Handle pending frames without blocking; False once Discord has gone."""
        try:
            while select.select([self.sock], [], [], 0)[0]:
                self._recv()
        except (OSError, ValueError, IPCError):
            return False
        return True

    def close(self) -> None:
        self.sock.close()


class Presence:
    """Keeps Discord's activity in line with the chosen game; closing the connection clears it."""

    def __init__(self, runtime_dir: Path) -> None:
        self.runtime_dir = runtime_dir
        self.ipc: DiscordIPC | None = None
        self.shown: Game | None = None
        self.failing = False

    def update(self, game: Game | None) -> None:
        if self.ipc is not None and not self.ipc.alive():
            log.info("Lost the connection to Discord")
            self.close()
        if self.shown is not None and (game is None or game.app.id != self.shown.app.id):
            self.close()
        if game is None:
            self.failing = False
            return
        if self.shown is not None and game.start == self.shown.start:
            return  # same app and timer: only the process changed
        try:
            if self.ipc is None:
                self.ipc = DiscordIPC.connect(game.app.id, self.runtime_dir)
            self.ipc.set_activity(game.start)
        except (OSError, ValueError, IPCError) as e:
            if not self.failing:
                log.info("Can't reach Discord (%s); retrying every scan", e)
            self.failing = True
            self.close()
            return
        self.failing = False
        self.shown = game
        log.info("Showing %s in Discord", game.app.name)

    def close(self) -> None:
        if self.ipc is not None:
            self.ipc.close()
        self.ipc = None
        self.shown = None


def run(config: Config, cache: Path, runtime_dir: Path, stop: threading.Event, proc: Path = Path("/proc")) -> None:
    index: GameIndex | None = None
    next_check = 0.0
    current: Game | None = None
    presence = Presence(runtime_dir)
    try:
        while not stop.is_set():
            if time.monotonic() >= next_check:
                if refresh_cache(cache) or index is None:
                    apps = read_cache(cache)
                    if apps is not None:
                        index = GameIndex(apps, config)
                        log.info("Loaded %d games", len(apps))
                next_check = time.monotonic() + (CACHE_CHECK if index is not None else CACHE_RETRY)

            game = choose(scan(index, proc), current) if index is not None else None
            if current is not None and (game is None or game.app.id != current.app.id):
                log.info("%s exited", current.app.name)
            if game is not None and (current is None or game.app.id != current.app.id):
                log.info("Detected %s (pid %d)", game.app.name, game.pid)
            elif game is not None and current is not None and game.pid != current.pid:
                log.debug("%s now pid %d", game.app.name, game.pid)
            current = game
            presence.update(game)
            stop.wait(config.interval)
    finally:
        presence.close()


def main(argv: list[str] | None = None) -> int:
    config_home = Path(os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config")
    cache_home = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache")
    parser = argparse.ArgumentParser(prog=NAME, description=__doc__)
    parser.add_argument("--config", type=Path, default=config_home / NAME / "config.toml", help="config file")
    parser.add_argument("--scan", action="store_true", help="print the games running now and exit")
    parser.add_argument("-v", "--verbose", action="store_true", help="debug logging")
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    args = parser.parse_args(argv)
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(levelname)s: %(message)s")

    try:
        config = load_config(args.config)
    except (OSError, ValueError) as e:
        log.error("Bad config %s: %s", args.config, e)
        return 2
    cache = cache_home / NAME / "detectable.json"
    runtime_dir = Path(os.environ.get("XDG_RUNTIME_DIR") or f"/run/user/{os.getuid()}")

    if args.scan:
        refresh_cache(cache)
        apps = read_cache(cache)
        if apps is None:
            log.error("No game list")
            return 1
        for game in scan(GameIndex(apps, config)):
            started = time.strftime("%H:%M", time.localtime(game.start))
            print(f"{game.app.name}  (app {game.app.id}, pid {game.pid}, started {started})")
        return 0

    stop = threading.Event()
    signal.signal(signal.SIGTERM, lambda *_: stop.set())
    signal.signal(signal.SIGINT, lambda *_: stop.set())
    run(config, cache, runtime_dir, stop)
    return 0


if __name__ == "__main__":
    sys.exit(main())