app.py

"""⚔️ Rift Chronicles — Build Small Hackathon 2026 Space backend.

Gradio gr.Server backend: hosts the PWA frontend (dist/) and exposes, via native
routes, the exact same /api protocol as the Cloudflare Worker (/api/chat SSE,
/api/image, /api/health), so the frontend runs on the Space with zero changes.

Compliance notes (the three hard Build Small constraints):
- Total parameters ≤ 32B: text uses Qwen3-14B (14B dense, self-hosted vLLM endpoint
  on Modal, OpenAI-compatible protocol); images use Z-Image-Turbo (6B, via HF
  Inference Providers). 14B + 6B = 20B total, under the 32B cap.
- Gradio app: gr.Server with a custom frontend (Off-Brand badge route, same as tokenwood).
- Judges can play instantly: no passphrase gate; abuse protection is an in-memory
  "global + per-IP" daily quota instead.

Secrets: text inference auth uses CARTBRAWL_API_KEY (configured in Space Settings →
Secrets; local development falls back to a .cartbrawl_api_key file in this directory,
which is never uploaded). Images still require HF_TOKEN.
"""
from __future__ import annotations

import io
import json
import os
import re
import threading
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

import httpx
from fastapi import Request
from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse
from fastapi.staticfiles import StaticFiles
from gradio import Server

ROOT = Path(__file__).resolve().parent
DIST = ROOT / "dist"

def _load_cartbrawl_key() -> str:
    """Prefer the env var (Space Secrets); local dev falls back to a key file in this directory (never uploaded)."""
    key = os.environ.get("CARTBRAWL_API_KEY", "")
    if key:
        return key
    try:
        return (ROOT / ".cartbrawl_api_key").read_text().strip()
    except OSError:
        return ""


# Image generation + tale-library writes. Note: the literal secret name "HF_TOKEN"
# can be silently reserved/dropped by the Hub, so we prefer HF_API_TOKEN.
HF_TOKEN = os.environ.get("HF_API_TOKEN", "") or os.environ.get("HF_TOKEN", "")
CARTBRAWL_API_KEY = _load_cartbrawl_key()
# Text inference: dedicated self-hosted vLLM on Modal (own endpoint + key for this project)
LLM_BASE_URL = os.environ.get(
    "CARTBRAWL_LLM_URL", "https://wuzhenze--rift-llm-serve.modal.run/v1"
)
LLM_MODEL = os.environ.get("CARTBRAWL_MODEL", "Qwen/Qwen3-14B")
# A single 14B endpoint is both smart and fast enough: all three tiers point to it (tier protocol kept, zero frontend changes)
MODELS = {"cheap": LLM_MODEL, "standard": LLM_MODEL, "smart": LLM_MODEL}
IMAGE_MODEL = os.environ.get("IMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")
# For Qwen creative writing, 0.9 on the usual temperature scale is plenty loose
TEMPERATURE = float(os.environ.get("LLM_TEMPERATURE", "0.9"))
DAILY_LIMIT = int(os.environ.get("DAILY_LIMIT", "1500"))
IP_DAILY_LIMIT = int(os.environ.get("IP_DAILY_LIMIT", "200"))

_budget_lock = threading.Lock()
_budget_day = ""
_budget_counts: dict[str, int] = {}


def _client_ip(request: Request) -> str:
    fwd = request.headers.get("x-forwarded-for", "")
    if fwd:
        return fwd.split(",")[0].strip()
    return request.client.host if request.client else "unknown"


def check_budget(request: Request) -> tuple[bool, int]:
    """In-memory daily quota (resets on Space restart — good enough for the hackathon judging week)."""
    global _budget_day
    day = datetime.now(timezone.utc).date().isoformat()
    ip = _client_ip(request)
    with _budget_lock:
        if day != _budget_day:
            _budget_day = day
            _budget_counts.clear()
        total = _budget_counts.get("total", 0)
        per_ip = _budget_counts.get(ip, 0)
        if total >= DAILY_LIMIT or per_ip >= IP_DAILY_LIMIT:
            return False, total
        _budget_counts["total"] = total + 1
        _budget_counts[ip] = per_ip + 1
        return True, total + 1


def today_calls() -> int:
    with _budget_lock:
        return _budget_counts.get("total", 0)


def _err(message: str, status: int, detail: str = "") -> JSONResponse:
    body: dict[str, str] = {"error": message}
    if detail:
        body["detail"] = detail
    return JSONResponse(body, status_code=status)


app = Server()


@app.get("/api/health")
async def health() -> JSONResponse:
    return JSONResponse({"ok": True, "todayCalls": today_calls()})


def _chat_payload(body: dict, with_extras: bool) -> dict:
    tier = body.get("tier")
    model = MODELS.get(tier if tier in MODELS else "standard", MODELS["standard"])
    payload: dict = {
        "model": model,
        "messages": body["messages"],
        "stream": True,
        "max_tokens": min(int(body.get("maxTokens") or 2048), 4096),
        "temperature": TEMPERATURE,
    }
    if with_extras:
        # Qwen3 models have an extended thinking mode that wraps output in <think>
        # blocks, which breaks JSON streaming. Disable it explicitly via vLLM's
        # chat_template_kwargs (the caller retries without extras on 400/422).
        payload["chat_template_kwargs"] = {"enable_thinking": False}
        if body.get("json"):
            payload["response_format"] = {"type": "json_object"}
    return payload


async def _open_upstream(client: httpx.AsyncClient, payload: dict) -> httpx.Response:
    req = client.build_request(
        "POST",
        f"{LLM_BASE_URL}/chat/completions",
        json=payload,
        headers={"Authorization": f"Bearer {CARTBRAWL_API_KEY}"},
    )
    return await client.send(req, stream=True)


@app.post("/api/chat")
async def chat(request: Request):
    if not CARTBRAWL_API_KEY:
        return _err("CARTBRAWL_API_KEY not configured on server", 501)
    try:
        body = await request.json()
    except Exception:
        return _err("Request body is not valid JSON", 400)
    if not isinstance(body.get("messages"), list) or not body["messages"]:
        return _err("messages must not be empty", 400)

    ok, _ = check_budget(request)
    if not ok:
        return _err("Today's story quota is used up — come back tomorrow.", 429)

    client = httpx.AsyncClient(timeout=httpx.Timeout(180, connect=30))
    try:
        upstream = await _open_upstream(client, _chat_payload(body, True))
        # Some providers reject chat_template_kwargs / response_format:
        # retry once without the extra params (the frontend already has its own
        # "extract JSON + re-ask on validation failure" fallback)
        if upstream.status_code in (400, 422):
            await upstream.aclose()
            upstream = await _open_upstream(client, _chat_payload(body, False))
        if upstream.status_code != 200:
            detail = (await upstream.aread()).decode("utf-8", "replace")[:500]
            await upstream.aclose()
            await client.aclose()
            return _err(f"Upstream model error {upstream.status_code}", 502, detail)
    except httpx.HTTPError as exc:
        await client.aclose()
        return _err("Failed to connect to upstream model", 502, str(exc)[:300])

    async def relay():
        try:
            async for chunk in upstream.aiter_raw():
                yield chunk
        finally:
            await upstream.aclose()
            await client.aclose()

    return StreamingResponse(
        relay(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache"},
    )


@app.post("/api/image")
def image(request: Request, body: dict):
    """Text-to-image (Z-Image-Turbo 6B). Sync endpoint: FastAPI runs it on the thread
    pool automatically, so the event loop is not blocked. Any failure returns non-200 —
    the frontend silently falls back to a procedurally generated image."""
    if not HF_TOKEN:
        return _err("Image service not configured on server", 501)
    prompt = body.get("prompt") or ""
    if not prompt or len(prompt) > 600:
        return _err("prompt is missing or too long", 400)
    ok, _ = check_budget(request)
    if not ok:
        return _err("Today's quota is used up", 429)

    try:
        width, height = (int(x) for x in (body.get("size") or "1024x1024").split("x"))
    except ValueError:
        width, height = 1024, 1024

    try:
        from huggingface_hub import InferenceClient

        img = InferenceClient(token=HF_TOKEN).text_to_image(
            prompt, model=IMAGE_MODEL, width=width, height=height
        )
        buf = io.BytesIO()
        img.save(buf, format="PNG")
    except Exception as exc:  # too many provider-side failure modes; return a uniform 502 and let the frontend fall back
        return _err("Image generation failed", 502, str(exc)[:300])

    return Response(
        content=buf.getvalue(),
        media_type="image/png",
        headers={"Cache-Control": "no-store"},
    )


# ---------------- Community tale library (storage = public HF Dataset repo) ----------------

TALES_REPO = os.environ.get("TALES_REPO", "Zhenzewu/rift-chronicles-tales")
_tales_lock = threading.Lock()
_tales_index: list[dict] | None = None
_tales_index_at = 0.0
_tales_cache: dict[str, dict] = {}
_TALE_ID_RE = re.compile(r"^[0-9]{6,}-[0-9a-f]{6}$")


def _tales_api():
    """Return an HfApi handle, or None when the library is not enabled."""
    if not HF_TOKEN:
        return None
    from huggingface_hub import HfApi

    return HfApi(token=HF_TOKEN)


def _sanitize(text: str, limit: int) -> str:
    return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", str(text))[:limit].strip()


def _tale_meta(tale: dict) -> dict:
    return {
        "id": tale["id"],
        "title": tale["title"],
        "hero": tale["hero"],
        "lang": tale.get("lang", "zh"),
        "chapterCount": len(tale.get("chapters", [])),
        "createdAt": tale.get("createdAt", 0),
    }


def _fetch_tale(api, tale_id: str) -> dict | None:
    if tale_id in _tales_cache:
        return _tales_cache[tale_id]
    from huggingface_hub import hf_hub_download

    try:
        path = hf_hub_download(
            repo_id=TALES_REPO,
            repo_type="dataset",
            filename=f"tales/{tale_id}.json",
            token=HF_TOKEN,
        )
        tale = json.loads(Path(path).read_text("utf-8"))
        tale["id"] = tale_id
        _tales_cache[tale_id] = tale
        return tale
    except Exception:
        return None


@app.get("/api/tales")
def tales_list() -> JSONResponse:
    api = _tales_api()
    if api is None:
        return _err("Tale library not enabled on this server", 501)
    global _tales_index, _tales_index_at
    with _tales_lock:
        if _tales_index is not None and time.time() - _tales_index_at < 60:
            return JSONResponse({"tales": _tales_index})
    try:
        files = api.list_repo_files(TALES_REPO, repo_type="dataset")
    except Exception as exc:
        return _err("Failed to read the tale library", 502, str(exc)[:200])
    ids = sorted(
        (f.removeprefix("tales/").removesuffix(".json") for f in files if f.startswith("tales/")),
        reverse=True,
    )[:30]
    metas = []
    for tale_id in ids:
        tale = _fetch_tale(api, tale_id)
        if tale:
            metas.append(_tale_meta(tale))
    with _tales_lock:
        _tales_index = metas
        _tales_index_at = time.time()
    return JSONResponse({"tales": metas})


@app.get("/api/tales/{tale_id}")
def tales_get(tale_id: str) -> JSONResponse:
    api = _tales_api()
    if api is None:
        return _err("Tale library not enabled on this server", 501)
    if not _TALE_ID_RE.match(tale_id):
        return _err("Invalid tale id", 400)
    tale = _fetch_tale(api, tale_id)
    if tale is None:
        return _err("Tale not found", 404)
    return JSONResponse(tale)


@app.post("/api/tales")
def tales_publish(request: Request, body: dict) -> JSONResponse:
    api = _tales_api()
    if api is None:
        return _err("Tale library not enabled on this server", 501)
    ok, _ = check_budget(request)
    if not ok:
        return _err("Today's quota is used up", 429)

    chapters_in = body.get("chapters")
    if not isinstance(chapters_in, list) or not (1 <= len(chapters_in) <= 20):
        return _err("chapters must contain 1–20 entries", 400)
    chapters = [
        {"title": _sanitize(c.get("title", ""), 80), "text": _sanitize(c.get("text", ""), 12000)}
        for c in chapters_in
        if isinstance(c, dict)
    ]
    if any(not c["title"] or len(c["text"]) < 50 for c in chapters):
        return _err("Every chapter needs a title and at least 50 characters of text", 400)
    tale = {
        "title": _sanitize(body.get("title", ""), 80) or "Untitled",
        "hero": _sanitize(body.get("hero", ""), 24) or "Anonymous",
        "lang": "en" if body.get("lang") == "en" else "zh",
        "createdAt": int(time.time() * 1000),
        "chapters": chapters,
    }
    payload = json.dumps(tale, ensure_ascii=False).encode("utf-8")
    if len(payload) > 120_000:
        return _err("Tale too large (120KB max)", 400)

    tale_id = f"{int(time.time())}-{uuid.uuid4().hex[:6]}"
    try:
        api.create_repo(TALES_REPO, repo_type="dataset", exist_ok=True)
        api.upload_file(
            path_or_fileobj=payload,
            path_in_repo=f"tales/{tale_id}.json",
            repo_id=TALES_REPO,
            repo_type="dataset",
            commit_message=f"Add tale {tale_id}: {tale['title'][:40]}",
        )
    except Exception as exc:
        return _err("Failed to publish the tale", 502, str(exc)[:200])

    tale["id"] = tale_id
    global _tales_index
    with _tales_lock:
        _tales_cache[tale_id] = tale
        _tales_index = None  # invalidate the list cache
    return JSONResponse({"id": tale_id})


@app.api(name="stats")
def stats() -> dict:
    """Gradio API endpoint: shows judges/curious folks today's usage and the model list (all ≤32B)."""
    return {"todayCalls": today_calls(), "models": MODELS, "imageModel": IMAGE_MODEL}


# Static frontend: cannot be mounted wholesale at "/" (it would shadow the
# /gradio_api/* routes registered by Gradio launch). The root path serves index.html,
# directories are mounted at sub-paths, and loose root-level PWA files
# (sw.js / manifest, etc.) are registered one by one.
@app.get("/")
async def homepage() -> FileResponse:
    return FileResponse(DIST / "index.html")


def _serve_file(path: Path):
    async def handler() -> FileResponse:
        return FileResponse(path)

    return handler


for _f in DIST.iterdir():
    if _f.is_file() and _f.name != "index.html":
        app.get(f"/{_f.name}")(_serve_file(_f))
for _d in ("assets", "icons"):
    app.mount(f"/{_d}", StaticFiles(directory=DIST / _d), name=_d)

if __name__ == "__main__":
    app.launch(show_error=True)