"""Bounded public game queue; private outbound B300 worker connection."""
import asyncio
from contextlib import contextmanager
import hashlib
import hmac
import json
import os
import re
import secrets
import sqlite3
import time
from pathlib import Path
from typing import Literal

from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import FileResponse
from pydantic import BaseModel, ConfigDict, Field, model_validator

router = APIRouter()
ROOT = Path(__file__).resolve().parent
DB_PATH = os.getenv("HF_GAMES_DB", "/tmp/hf-games.sqlite3")
lock = asyncio.Lock()
worker_seen = 0.0
GAMES = ("mario", "cube", "tetris")
ACTIVE = ("queued", "running")
FINAL = ("completed", "failed", "cancelled")
MAX_BYTES = 100 * 1024 * 1024


@contextmanager
def db():
    con = sqlite3.connect(DB_PATH)
    con.row_factory = sqlite3.Row
    con.execute("PRAGMA journal_mode=WAL")
    con.executescript("""
    CREATE TABLE IF NOT EXISTS jobs (
      id TEXT PRIMARY KEY, ticket_hash TEXT, client_id TEXT, game TEXT,
      seed INTEGER, status TEXT, created REAL, updated REAL,
      cancel INTEGER DEFAULT 0, summary TEXT DEFAULT '{}',
      last_seq INTEGER DEFAULT 0, bytes INTEGER DEFAULT 0);
    CREATE TABLE IF NOT EXISTS events (
      job_id TEXT, seq INTEGER, body TEXT, PRIMARY KEY(job_id,seq));
    CREATE INDEX IF NOT EXISTS jobs_age ON jobs(created);
    CREATE TABLE IF NOT EXISTS starts (client_id TEXT, created REAL);
    """)
    try:
        with con:
            yield con
    finally:
        con.close()


def reap(con):
    now = time.time()
    con.execute("DELETE FROM starts WHERE created < ?", (now - 3600,))
    # Bound event retention; original run traces live on the B300 worker.
    old = [r[0] for r in con.execute(
        "SELECT id FROM jobs WHERE updated < ? AND status NOT IN ('queued','running')",
        (now - 3600,))]
    for job_id in old:
        con.execute("DELETE FROM events WHERE job_id=?", (job_id,))
        con.execute("DELETE FROM jobs WHERE id=?", (job_id,))
    con.execute(
        "UPDATE jobs SET status='failed',cancel=1,summary=?,updated=? "
        "WHERE status='running' AND (updated < ? OR created < ?)",
        (json.dumps({"message": "Worker connection ended. Start a new run."}),
         now, now - 45, now - 360))
    con.execute(
        "UPDATE jobs SET status='cancelled',cancel=1,summary=?,updated=? "
        "WHERE status='queued' AND created < ?",
        (json.dumps({"message": "Queue expired. Start a new run."}), now, now - 360))


def worker_token():
    dedicated = os.getenv("HF_GAMES_WORKER_TOKEN", "").strip()
    if dedicated:
        return dedicated
    existing = os.getenv("JEV_API_KEY", "").strip()
    if not existing:
        return ""
    # Domain-separated key: a worker credential cannot authenticate to model inference.
    return hmac.new(existing.encode(), b"autotrust/JEV-27B-Demo:game-worker:v1",
                    hashlib.sha256).hexdigest()


def worker_auth(request):
    expected = worker_token()
    received = request.headers.get("authorization", "")
    if not expected or not hmac.compare_digest(received, "Bearer " + expected):
        raise HTTPException(401, "Worker authentication required.")


def ticket_auth(request, row):
    value = request.headers.get("x-job-ticket", "")
    digest = hashlib.sha256(value.encode()).hexdigest()
    if not row or not hmac.compare_digest(digest, row["ticket_hash"]):
        raise HTTPException(404, "This run is unavailable. Start a new run.")


class Start(BaseModel):
    model_config = ConfigDict(extra="forbid")
    game: Literal["mario", "cube", "tetris"]
    seed: int = Field(default=0, ge=0, le=2147483647, strict=True)
    client_id: str = Field(min_length=20, max_length=80, pattern=r"^[A-Za-z0-9_-]+$")


    @model_validator(mode="after")
    def mario_seed(self):
        if self.game == "mario" and self.seed != 0:
            raise ValueError("Mario World 1-1 uses seed zero")
        return self


class Claim(BaseModel):
    model_config = ConfigDict(extra="forbid")
    games: list[Literal["mario", "cube", "tetris"]] = Field(max_length=3)


class Push(BaseModel):
    model_config = ConfigDict(extra="forbid")
    events: list[dict] = Field(default_factory=list, max_length=10)
    status: Literal["queued", "running", "completed", "failed", "cancelled"] | None = None
    summary: dict | None = None


@router.get("/arena")
@router.get("/arena.html")
async def arena_page():
    return FileResponse(ROOT / "arena.html")


@router.get("/api/arena/status")
async def status():
    async with lock:
        with db() as con:
            reap(con)
            games = {}
            for game in GAMES:
                counts = dict(con.execute(
                    "SELECT status,COUNT(*) FROM jobs WHERE game=? AND status IN ('queued','running') GROUP BY status",
                    (game,)).fetchall())
                games[game] = {"busy": bool(counts.get("running", 0)), "queued": counts.get("queued", 0)}
    return {"online": bool(worker_token()) and time.time()-worker_seen < 20,
            "model": "jev-27b-vl", "games": games}


@router.post("/api/arena/jobs", status_code=201)
async def start(payload: Start):
    if not worker_token() or time.time()-worker_seen >= 20:
        raise HTTPException(503, "Game worker is reconnecting. Please retry shortly.")
    async with lock:
        with db() as con:
            reap(con)
            if con.execute("SELECT 1 FROM jobs WHERE client_id=? AND game=? AND status IN ('queued','running')",
                           (payload.client_id, payload.game)).fetchone():
                raise HTTPException(409, "You already have an active run for this game.")
            if con.execute("SELECT COUNT(*) FROM jobs WHERE game=? AND status='queued'",
                           (payload.game,)).fetchone()[0] >= 3:
                raise HTTPException(429, "This game's queue is full. Try again shortly.", headers={"Retry-After":"30"})
            now = time.time()
            global_count = con.execute("SELECT COUNT(*) FROM starts WHERE created > ?", (now-60,)).fetchone()[0]
            user_count = con.execute("SELECT COUNT(*) FROM starts WHERE client_id=?", (payload.client_id,)).fetchone()[0]
            if global_count >= 15 or user_count >= 20:
                raise HTTPException(429, "Please wait before starting another run.", headers={"Retry-After":"60"})
            if con.execute("SELECT COALESCE(SUM(bytes),0) FROM jobs").fetchone()[0] > 600*1024*1024:
                raise HTTPException(429, "The replay store is busy. Try again later.")
            job_id, ticket = secrets.token_urlsafe(18), secrets.token_urlsafe(32)
            con.execute("INSERT INTO jobs(id,ticket_hash,client_id,game,seed,status,created,updated) VALUES (?,?,?,?,?,'queued',?,?)",
                        (job_id, hashlib.sha256(ticket.encode()).hexdigest(), payload.client_id,
                         payload.game, payload.seed, now, now))
            con.execute("INSERT INTO starts VALUES (?,?)", (payload.client_id, now))
    return {"id": job_id, "ticket": ticket, "status":"queued", "game":payload.game, "seed":payload.seed}


@router.get("/api/arena/jobs/{job_id}")
async def get_job(job_id: str, request: Request, after: int = 0, limit: int = 30):
    after, limit = max(0, after), max(1, min(30, limit))
    async with lock:
        with db() as con:
            reap(con)
            row = con.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
            ticket_auth(request, row)
            events = []
            size = 0
            for event in con.execute("SELECT body FROM events WHERE job_id=? AND seq>? ORDER BY seq LIMIT ?",
                                     (job_id, after, limit)):
                if events and size+len(event[0]) > 2*1024*1024:
                    break
                events.append(json.loads(event[0]))
                size += len(event[0])
            last = events[-1]["seq"] if events else after
            return {"id":row["id"], "game":row["game"], "seed":row["seed"], "status":row["status"],
                    "events":events, "last_seq":row["last_seq"], "has_more":last < row["last_seq"],
                    "summary":json.loads(row["summary"]), "cancelled":bool(row["cancel"])}


@router.post("/api/arena/jobs/{job_id}/cancel")
async def cancel(job_id: str, request: Request):
    async with lock:
        with db() as con:
            row = con.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
            ticket_auth(request, row)
            if row["status"] in ACTIVE:
                next_status = "cancelled" if row["status"] == "queued" else "running"
                con.execute("UPDATE jobs SET cancel=1,status=?,updated=? WHERE id=?",
                            (next_status, time.time(), job_id))
    return {"cancelled":True}


@router.post("/api/arena/worker/claim")
async def claim(payload: Claim, request: Request):
    global worker_seen
    worker_auth(request)
    worker_seen = time.time()
    jobs = []
    async with lock:
        with db() as con:
            reap(con)
            for game in dict.fromkeys(payload.games):
                if con.execute("SELECT 1 FROM jobs WHERE game=? AND status='running'", (game,)).fetchone():
                    continue
                row = con.execute("SELECT * FROM jobs WHERE game=? AND status='queued' AND cancel=0 ORDER BY created LIMIT 1",
                                  (game,)).fetchone()
                if row:
                    con.execute("UPDATE jobs SET status='running',updated=? WHERE id=?", (time.time(), row["id"]))
                    jobs.append({"id":row["id"], "game":game, "seed":row["seed"]})
    return {"jobs":jobs}


@router.post("/api/arena/worker/jobs/{job_id}/events")
async def push(job_id: str, payload: Push, request: Request):
    global worker_seen
    worker_auth(request)
    worker_seen = time.time()
    async with lock:
        with db() as con:
            row = con.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
            if row is None:
                raise HTTPException(404, "Run expired.")
            if row["status"] in FINAL:
                return {"cancelled": True, "last_seq": row["last_seq"]}
            total, seq = row["bytes"], row["last_seq"]
            for event in payload.events:
                n = event.get("seq")
                if type(n) is not int or n < 1:
                    raise HTTPException(422, "Invalid event sequence.")
                if n <= seq:
                    continue
                if n != seq + 1:
                    raise HTTPException(409, "Event sequence has a gap.")
                if event.get("kind") not in ("initial", "step", "complete"):
                    raise HTTPException(422, "Invalid event kind.")
                body = json.dumps(event, separators=(",", ":"), allow_nan=False)
                if len(body) > 2*1024*1024 or total+len(body)>MAX_BYTES or n>6000:
                    con.execute("UPDATE jobs SET cancel=1 WHERE id=?", (job_id,))
                    return {"cancelled":True,"last_seq":seq}
                con.execute("INSERT INTO events VALUES (?,?,?)", (job_id, n, body))
                total += len(body)
                seq = n
            state = payload.status or row["status"]
            if state == "queued":
                state = "running"
            summary = json.dumps(payload.summary, allow_nan=False) if payload.summary is not None else row["summary"]
            if len(summary)>20000:
                raise HTTPException(422,"Summary is too large.")
            if row["cancel"] and state in FINAL:
                state = "cancelled"
            con.execute("UPDATE jobs SET status=?,summary=?,updated=?,last_seq=?,bytes=? WHERE id=?",
                        (state,summary,time.time(),seq,total,job_id))
    return {"cancelled":bool(row["cancel"]),"last_seq":seq}