TSSBOT/tss_ws.py

#!/usr/bin/env python3
"""
tss_ws.py

Standalone TSS WebSocket listener for TSSBOT.

Each received game is written to:
  STORAGE_VOL_PATH/REPLAYS/TSS/<session_id>/replay_data.json.gz

Run modes:
  python tss_ws.py          — listen forever, write every game
  python tss_ws.py --one    — capture exactly one game, write it, then exit

Reads SPECTRA_WS_TSS_URL, SPECTRA_API_KEY, and STORAGE_VOL_PATH from TSSBOT/.env.
"""
from __future__ import annotations

import asyncio
import gzip
import json
import logging
import os
import sys
import time
from pathlib import Path
from typing import Any, Callable, Awaitable, List, Dict, Optional

from dotenv import load_dotenv
from websockets.asyncio.client import connect as wsconnect

from BOT.storage import insert_match, insert_player_games, upsert_tss_teams
from BOT.autologging import process_game as autolog_process_game
from BOT.receiver_bridge import publish_replay_batch
from spectra_replay_normalize import normalize_spectra_replay
from spectra_ws_payload import SpectraPayloadError, decode_spectra_ws_payload

_HERE = Path(__file__).resolve().parent
load_dotenv(dotenv_path=_HERE / ".env")

logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s] [%(levelname)s] [tss-ws] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("tss-ws")

WS_URL  = os.getenv("SPECTRA_WS_TSS_URL", "").strip()
API_KEY = os.getenv("SPECTRA_API_KEY", "").strip()

_storage_raw = os.getenv("STORAGE_VOL_PATH", "").strip()
if not _storage_raw:
    log.error("STORAGE_VOL_PATH is not set in TSSBOT/.env — aborting")
    sys.exit(1)

REPLAYS_DIR = Path(_storage_raw) / "REPLAYS" / "TSS"
REPLAYS_DIR.mkdir(parents=True, exist_ok=True)


def _auth_header() -> str:
    return API_KEY if API_KEY.startswith("Bearer ") else f"Bearer {API_KEY}"


def _session_id(game: Dict[str, Any]) -> str:
    """Return a hex session ID, converting decimal string/int IDs from Spectra."""
    raw = game.get("_id") or game.get("id")
    if raw is None:
        return f"unknown_{int(time.time())}"
    try:
        return hex(int(raw))[2:].lower()
    except (ValueError, TypeError):
        s = str(raw).strip().lower()
        return s[2:] if s.startswith("0x") else s


def _write_game(game: Dict[str, Any]) -> Path:
    """Normalize _id to hex, then write to REPLAYS/TSS/<session_id>/replay_data.json.gz."""
    game.update(normalize_spectra_replay(game))
    sid = _session_id(game)
    game["_id"] = sid  # hex from this point forward
    session_dir = REPLAYS_DIR / sid
    session_dir.mkdir(parents=True, exist_ok=True)
    out = session_dir / "replay_data.json.gz"
    raw = json.dumps(game, ensure_ascii=False).encode("utf-8")
    with gzip.open(out, "wb") as fh:
        fh.write(raw)
    return out


def normalize(data: Any) -> Optional[List[Dict[str, Any]]]:
    """Coerce whatever the WS sends into a list of game dicts."""
    if isinstance(data, list) and data and isinstance(data[0], dict):
        return data
    if isinstance(data, dict):
        if "_id" in data or "id" in data:
            return [data]
        if isinstance(data.get("data"), dict):
            return [data["data"]]
        if isinstance(data.get("data"), list):
            return data["data"]
        if "completed" in data:
            return data["completed"]
    log.warning("Unknown WS frame shape: %s", type(data))
    return None


async def listen(
    on_game: Callable[[Dict[str, Any]], Awaitable[None]],
    *,
    stop_after: int = 0,
) -> None:
    """
    Maintain a persistent WSS connection to the TSS endpoint.

    Args:
        on_game:    Async callback invoked once per game dict received.
        stop_after: If > 0, stop after this many games have been delivered.
    """
    if not WS_URL:
        log.error("SPECTRA_WS_TSS_URL is not set in TSSBOT/.env — aborting")
        return

    headers = {"Authorization": _auth_header()}
    reconnect_delay = 1
    received = 0

    while True:
        try:
            log.info("Connecting → %s", WS_URL)
            async with wsconnect(
                WS_URL,
                additional_headers=headers,
                max_size=32 * 1024 * 1024,
            ) as ws:
                log.info("Connected.")
                reconnect_delay = 1  # reset on successful connect
                async for message in ws:
                    try:
                        data = decode_spectra_ws_payload(message)
                    except SpectraPayloadError as exc:
                        log.warning("Bad Spectra WS frame, skipping: %s", exc)
                        continue

                    games = normalize(data)
                    if not games:
                        continue

                    for game in games:
                        received += 1
                        await on_game(game)
                        if stop_after and received >= stop_after:
                            log.info("Captured %d game(s), done.", received)
                            return

        except Exception as exc:
            log.error("WS error: %s", exc)

        log.info("Reconnecting in %ds...", reconnect_delay)
        await asyncio.sleep(reconnect_delay)
        reconnect_delay = min(reconnect_delay * 2, 30)


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

async def _handle_game(game: Dict[str, Any]) -> None:
    out = _write_game(game)
    sid = game["_id"]  # hex, normalized by _write_game
    log.info("Saved game %s → %s", sid, out)
    try:
        await insert_match(game)
        await insert_player_games(game)
        await upsert_tss_teams(game)
        from BOT.match_logs import build_event_log, build_match_logs, upsert_match_logs
        from BOT.storage import TSS_BATTLES_DB_PATH
        chat_log, battle_log = build_match_logs(game)
        await upsert_match_logs(TSS_BATTLES_DB_PATH, sid, chat_log, battle_log, build_event_log(game))
        log.info("Stored game %s in DB", sid)
    except Exception as exc:
        log.error("DB insert failed for %s: %s", sid, exc)
    # Forward the processed game to the relay gateway (tss channel). Best-effort:
    # a bridge failure must never break ingest.
    try:
        await publish_replay_batch([game])
    except Exception as exc:
        log.warning("[TSS-BRIDGE] Failed to forward replay batch for %s: %s", sid, exc)
    # Link this game's replay to its bracket slot from tss.match_id (zero TSS
    # calls) and request a debounced schedule refresh; unknown matches fall back
    # to a full scan inside maybe_link_battle. Never blocks ingest.
    try:
        from BOT.tss_tournaments import maybe_link_battle
        await maybe_link_battle(game)
    except Exception as exc:
        log.error("tournament link trigger failed for %s: %s", sid, exc)
    # Autolog match/dispatch (no-ops in standalone mode where no bot is set).
    try:
        await autolog_process_game(game)
    except Exception as exc:
        log.error("autolog failed for %s: %s", sid, exc)
    print(json.dumps(game, indent=2, ensure_ascii=False))
    print("-" * 80)


async def main() -> None:
    one_shot = "--one" in sys.argv
    if one_shot:
        log.info("--one mode: will exit after the first game arrives")
        await listen(_handle_game, stop_after=1)
    else:
        log.info("Listening forever (Ctrl-C to stop)")
        await listen(_handle_game)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        log.info("Stopped.")