app.py

"""
Sozai — hybrid room collaboration, served with gradio.Server (Server mode).

Gradio is used ONLY as the backend (FastAPI + Gradio's API engine). The whole
UI is custom and comes from these files:

    landing.html   the "Sozai" splash (make a room / join by code / skip)
    app.html       the collaborative album app (the React port from before)

What this adds on top of the previous version:
  * SQLite database (created locally on first run -> sozai_rooms.db) holding
    `rooms` and `participants`, matching the project spec.
  * Hybrid room model: secure room_id (uuid4) + short human room code whose
    LENGTH THE USER CHOOSES (4-12) on the splash screen.
  * REST endpoints: create room, get room state, resolve join-by-code.
  * A WebSocket endpoint /ws/{room_id} for real-time sync (presence, chat,
    shared selection, captions, live cursors, ...).
  * A background cleanup thread that expires idle (1h) and old (24h) rooms.
  * A "continue without a room" path: /app runs the app solo, no collaboration.

Run it:
    pip install gradio
    python app.py            # -> http://localhost:7860
Drop app.py, landing.html, app.html next to storybook_render.zip (or an
extracted public/ folder). Assets are read straight out of the zip.
"""

from __future__ import annotations

import os

# --------------------------------------------------------------------------- #
# ZeroGPU process-start setup. These MUST be set before NeMo / torch are
# imported, because on Hugging Face ZeroGPU the GPU is not attached at process
# start and CUDA must NOT be initialized in the main process. NeMo's RNNT path
# (and numba) can otherwise trigger a low-level CUDA init during import/load,
# which surfaces as: "Low-level CUDA init (torch._C._cuda_init) reached. This
# means ZeroGPU's PyTorch CUDA emulation mode did not intercept a CUDA
# operation". This block mirrors NVIDIA's reference Nemotron/Parakeet Spaces.
# --------------------------------------------------------------------------- #
_omp = os.environ.get("OMP_NUM_THREADS", "")
if _omp and not _omp.isdigit():
    os.environ["OMP_NUM_THREADS"] = str(
        max(1, int("".join(c for c in _omp if c.isdigit()) or "1"))
    )
# Stop NeMo's RNNT/numba path from initializing CUDA in the main process.
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
# Help torch's allocator play nicely with ZeroGPU's stateless GPU.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# Disable the Xet downloader: on this Space it both fails to write in the GPU
# worker AND crawls at ~1 kB/s for the 16 GB FLUX weights. Forcing the standard
# HTTPS download path makes the main-process pre-fetch actually complete.
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "0")
# Gradio 6 turns SSR on by default whenever Node.js is present in the image. Its
# SSR Node proxy mounts at "/" and serves an empty Gradio shell, shadowing our
# custom Server() routes (landing.html + the React app) -> blank page. This is a
# pure custom-FastAPI app, not a Gradio `demo`, so force SSR off.
os.environ["GRADIO_SSR_MODE"] = "false"

import asyncio
import functools
import hashlib
import gzip
import io
import json
import mimetypes
import re
import secrets
import sqlite3
import threading
import time
import uuid

from fastapi import HTTPException, Request
from fastapi.responses import FileResponse, HTMLResponse, Response, StreamingResponse

from gradio import Server

# --------------------------------------------------------------------------- #
# ZeroGPU support.
#
# On Hugging Face *ZeroGPU* Spaces, a GPU is attached only while a function
# decorated with `@spaces.GPU` is running, and the runtime REFUSES TO START if
# it can't find at least one such function during startup (that is the
# "No @spaces.GPU function detected during startup" error). The `spaces`
# package is preinstalled in the Space image.
#
# Locally (and on any non-ZeroGPU box) the `spaces` package usually isn't
# installed. The decorator is documented to be effect-free off ZeroGPU, so we
# provide a tiny no-op fallback that simply returns the function unchanged.
# This lets the SAME code run both locally and in a Space with no branching:
#   * In a ZeroGPU Space   -> real `spaces.GPU` requests/releases the GPU.
#   * Anywhere else        -> `spaces.GPU` is a transparent pass-through.
try:
    import spaces  # provided by the HF ZeroGPU Space image
    _ZEROGPU = True
except Exception:  # not on ZeroGPU (e.g. local dev, CPU Space, plain GPU Space)
    _ZEROGPU = False

    class _SpacesShim:
        """Minimal stand-in for the `spaces` module off ZeroGPU.

        Supports both `@spaces.GPU` and `@spaces.GPU(...)` call styles and just
        returns the wrapped function untouched."""

        @staticmethod
        def GPU(fn=None, **_kwargs):
            if fn is None:                      # used as @spaces.GPU(duration=...)
                def _wrap(f):
                    return f
                return _wrap
            return fn                            # used as @spaces.GPU

    spaces = _SpacesShim()  # type: ignore

# Load environment variables from a .env file (if present) so SOZAI_*,
# PHOENIX_*, and the Cesium ion token can be set without exporting them.
# Done early so every os.environ.get(...) below sees the values.
try:
    from dotenv import load_dotenv
    load_dotenv()
except Exception:  # python-dotenv not installed → rely on the real environment
    pass

BASE = os.path.dirname(os.path.abspath(__file__))

# --------------------------------------------------------------------------- #
# Frontend assets: read from a real public/ dir if present, else from the zip.
# --------------------------------------------------------------------------- #

def _first_existing(paths):
    return next((p for p in paths if p and os.path.exists(p)), None)


PUBLIC_DIR = _first_existing([
    os.environ.get("STORYBOOK_PUBLIC"),
    os.path.join(BASE, "public"),
    os.path.join(BASE, "storybook_render", "public"),
    os.path.join(BASE, "src", "public"),
])
ZIP_PATH = _first_existing([
    os.environ.get("STORYBOOK_ZIP"),
    os.path.join(BASE, "storybook_render.zip"),
])
_ZIP_PREFIX = None


@functools.lru_cache(maxsize=512)
def _read_from_zip(rel_path: str) -> bytes:
    import zipfile
    global _ZIP_PREFIX
    with zipfile.ZipFile(ZIP_PATH) as z:
        if _ZIP_PREFIX is None:
            _ZIP_PREFIX = ""
            for n in z.namelist():
                idx = n.find("public/")
                if idx != -1:
                    _ZIP_PREFIX = n[:idx]
                    break
        for name in (f"{_ZIP_PREFIX}public/{rel_path}", f"public/{rel_path}", rel_path):
            try:
                return z.read(name)
            except KeyError:
                continue
    raise KeyError(rel_path)


if PUBLIC_DIR is None and ZIP_PATH is None:
    raise SystemExit(
        "Could not find frontend assets. Put storybook_render.zip next to "
        "app.py, or extract its public/ folder alongside app.py."
    )

# --------------------------------------------------------------------------- #
# Easter-egg pet sprites: the 2023 oneko-style icon library (76 animals, 32
# sprites each). Served from a real folder if present, else from neko-main.zip.
# Each animal lives at <root>/2023-icon-library/<animal>/<sprite>.png .
# --------------------------------------------------------------------------- #

SPRITES_DIR = _first_existing([
    os.environ.get("SOZAI_SPRITES_DIR"),
    os.path.join(BASE, "2023-icon-library"),
    os.path.join(BASE, "public", "2023-icon-library"),
    os.path.join(BASE, "sprites"),
])
SPRITES_ZIP = _first_existing([
    os.environ.get("SOZAI_SPRITES_ZIP"),
    os.path.join(BASE, "neko-main.zip"),
])
_SPRITES_ZIP_PREFIX = None

# The canonical sprite filenames shared by every animal in the library.
SPRITE_FILES = [
    "alert", "still", "sleep1", "sleep2", "wash", "yawn", "itch1", "itch2",
    "nrun1", "nrun2", "nerun1", "nerun2", "erun1", "erun2", "serun1", "serun2",
    "srun1", "srun2", "swrun1", "swrun2", "wrun1", "wrun2", "nwrun1", "nwrun2",
    "nscratch1", "nscratch2", "escratch1", "escratch2",
    "sscratch1", "sscratch2", "wscratch1", "wscratch2",
]


@functools.lru_cache(maxsize=1)
def list_sprite_animals():
    """Return the sorted list of available animal sprite-set names."""
    names = set()
    if SPRITES_DIR and os.path.isdir(SPRITES_DIR):
        for d in os.listdir(SPRITES_DIR):
            if os.path.isdir(os.path.join(SPRITES_DIR, d)):
                names.add(d)
    elif SPRITES_ZIP:
        import zipfile
        try:
            with zipfile.ZipFile(SPRITES_ZIP) as z:
                for n in z.namelist():
                    idx = n.find("2023-icon-library/")
                    if idx == -1:
                        continue
                    rest = n[idx + len("2023-icon-library/"):]
                    animal = rest.split("/", 1)[0]
                    if animal and "." not in animal:
                        names.add(animal)
        except Exception:
            pass
    return sorted(names)


@functools.lru_cache(maxsize=4096)
def _read_sprite(animal: str, fname: str) -> bytes:
    """Read a single sprite PNG for an animal from folder or zip."""
    safe_animal = "".join(ch for ch in animal if ch.isalnum() or ch in "-_")
    safe_file = "".join(ch for ch in fname if ch.isalnum() or ch in "-_")
    rel = f"2023-icon-library/{safe_animal}/{safe_file}.png"
    if SPRITES_DIR and os.path.isdir(SPRITES_DIR):
        fp = os.path.join(SPRITES_DIR, safe_animal, safe_file + ".png")
        if os.path.isfile(fp):
            with open(fp, "rb") as fh:
                return fh.read()
    if SPRITES_ZIP:
        import zipfile
        global _SPRITES_ZIP_PREFIX
        with zipfile.ZipFile(SPRITES_ZIP) as z:
            if _SPRITES_ZIP_PREFIX is None:
                _SPRITES_ZIP_PREFIX = ""
                for n in z.namelist():
                    idx = n.find("2023-icon-library/")
                    if idx != -1:
                        _SPRITES_ZIP_PREFIX = n[:idx]
                        break
            for name in (f"{_SPRITES_ZIP_PREFIX}{rel}", rel):
                try:
                    return z.read(name)
                except KeyError:
                    continue
    raise KeyError(rel)


# --------------------------------------------------------------------------- #
# SQLite — created locally on first run.
# --------------------------------------------------------------------------- #

DB_PATH = os.environ.get("SOZAI_DB", os.path.join(BASE, "sozai_rooms.db"))
DB_LOCK = threading.RLock()
_conn = sqlite3.connect(DB_PATH, check_same_thread=False)
_conn.row_factory = sqlite3.Row


def init_db():
    with DB_LOCK:
        # auto_vacuum must be set BEFORE tables exist to take effect on a fresh
        # DB; on an existing file it only applies after a full VACUUM (below).
        # INCREMENTAL lets the cleanup loop reclaim freed pages to the OS so the
        # file doesn't keep growing as rooms are created and deleted.
        _conn.execute("PRAGMA auto_vacuum=INCREMENTAL;")
        _conn.execute("PRAGMA journal_mode=WAL;")
        _conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS rooms (
                room_id          TEXT PRIMARY KEY,
                share_code       TEXT UNIQUE,
                code_length      INTEGER,
                created_at       REAL,
                last_activity    REAL,
                owner_session_id TEXT,
                is_active        INTEGER DEFAULT 1,
                metadata         TEXT
            );
            CREATE TABLE IF NOT EXISTS participants (
                participant_id TEXT PRIMARY KEY,
                room_id        TEXT,
                session_id     TEXT,
                display_name   TEXT,
                joined_at      REAL,
                last_seen      REAL
            );
            CREATE TABLE IF NOT EXISTS room_state (
                room_id    TEXT PRIMARY KEY,
                state      TEXT,
                version    INTEGER DEFAULT 0,
                updated_at REAL
            );
            CREATE INDEX IF NOT EXISTS idx_rooms_code ON rooms(share_code);
            CREATE INDEX IF NOT EXISTS idx_rooms_activity ON rooms(last_activity);
            CREATE INDEX IF NOT EXISTS idx_part_room ON participants(room_id);
            """
        )
        _conn.commit()
        # If this is an existing DB created without auto_vacuum, a one-time full
        # VACUUM is needed for the INCREMENTAL setting to take hold. Cheap on a
        # small/empty DB; safe to run on every startup.
        try:
            mode = _conn.execute("PRAGMA auto_vacuum;").fetchone()[0]
            if mode != 2:  # 2 == INCREMENTAL
                _conn.execute("VACUUM;")
                _conn.commit()
        except Exception:
            pass


init_db()

# --------------------------------------------------------------------------- #
# Identifiers
# --------------------------------------------------------------------------- #

# Unambiguous uppercase alphabet (no 0/O/1/I/L) -> ~4.9 bits/char.
CODE_ALPHABET = "23456789ABCDEFGHJKMNPQRSTVWXYZ"
MIN_CODE_LEN, MAX_CODE_LEN = 4, 12
IDLE_TIMEOUT = int(os.environ.get("SOZAI_IDLE_TIMEOUT", str(60 * 60)))          # 1 hour -> mark inactive
# Delete a room (and its photos/state) after this much inactivity. Keyed on
# last_activity, so the DB reflects live usage and doesn't bloat. Defaults to
# 1 hour + a 10-minute grace window; override with SOZAI_ROOM_DELETE_AFTER_IDLE.
ROOM_DELETE_AFTER_IDLE = int(os.environ.get("SOZAI_ROOM_DELETE_AFTER_IDLE", str(60 * 60 + 600)))
HARD_EXPIRATION = int(os.environ.get("SOZAI_HARD_EXPIRATION", str(24 * 60 * 60)))  # absolute cap


def clamp_code_len(n) -> int:
    try:
        n = int(n)
    except (TypeError, ValueError):
        n = 6
    return max(MIN_CODE_LEN, min(MAX_CODE_LEN, n))


def gen_share_code(length: int) -> str:
    return "".join(secrets.choice(CODE_ALPHABET) for _ in range(length))


def unique_share_code(length: int) -> str:
    length = clamp_code_len(length)
    with DB_LOCK:
        for _ in range(40):
            code = gen_share_code(length)
            row = _conn.execute(
                "SELECT 1 FROM rooms WHERE share_code=? AND is_active=1", (code,)
            ).fetchone()
            if row is None:
                return code
    # extremely unlikely fallback: widen the code
    return gen_share_code(min(MAX_CODE_LEN, length + 2))


def sanitize_name(name, default="Guest"):
    name = (name or "").strip()
    name = "".join(ch for ch in name if ch.isprintable())
    return name[:24] or default


def touch_room(room_id: str):
    with DB_LOCK:
        _conn.execute(
            "UPDATE rooms SET last_activity=? WHERE room_id=?", (time.time(), room_id)
        )
        _conn.commit()


def room_row(room_id: str):
    with DB_LOCK:
        return _conn.execute(
            "SELECT * FROM rooms WHERE room_id=? AND is_active=1", (room_id,)
        ).fetchone()


# --------------------------------------------------------------------------- #
# Gradio Server
# --------------------------------------------------------------------------- #

app = Server()


# Kept from the previous version: a Gradio API endpoint the frontend can call
# through @gradio/client. Still works in collab mode.
@app.api(name="save_to_album")
def save_to_album(album: str = "", caption: str = "", tags: str = "") -> str:
    """Save a photo + caption + tags to an album, returning a confirmation."""
    album = (album or "").strip()
    if not album:
        return "Select an album to save this photo."
    return f'Added to "{album}".'


# --------------------------------------------------------------------------- #
# Auto-caption: a vision-language model that suggests a caption for a photo.
# The model (MiniCPM-V) is heavy, so it is loaded lazily on first use and
# cached. If transformers / llama.cpp / the model weights are unavailable, we
# fall back to a simple suggestion so the feature degrades gracefully instead
# of erroring.
#
# IMPORTANT: MiniCPM-V is a true VISION-LANGUAGE model. Both backends are wired
# to send the ACTUAL IMAGE to the model (not just alt-text), so the suggested
# title / caption / tags are grounded in what the photo actually shows.
# --------------------------------------------------------------------------- #

AUTOCAPTION_MODEL_ID = os.environ.get("SOZAI_CAPTION_MODEL", "openbmb/MiniCPM-V-4")
# GGUF variant used when running inside a Hugging Face Space (where llama.cpp +
# a quantized GGUF is the practical way to run MiniCPM-V). Outside Spaces we use
# the normal transformers model above. Both are overridable via env.
#
# MiniCPM-V-4 is a vision-language model, so the GGUF build needs TWO files: the
# LM weights (CAPTION_GGUF_FILE) and the vision projector / CLIP encoder
# (CAPTION_GGUF_MMPROJ). Without the mmproj the model can't see images. Use the
# plain mmproj-model-f16.gguf — NOT the -IOS variant, which is iOS-only and
# fails to load in llama.cpp.
CAPTION_GGUF_REPO = os.environ.get("SOZAI_CAPTION_GGUF_REPO", "openbmb/MiniCPM-V-4-gguf")
CAPTION_GGUF_FILE = os.environ.get("SOZAI_CAPTION_GGUF_FILE", "ggml-model-Q4_K_M.gguf")
CAPTION_GGUF_MMPROJ = os.environ.get("SOZAI_CAPTION_GGUF_MMPROJ", "mmproj-model-f16.gguf")
_caption_lock = threading.Lock()
# backend: "" until loaded, then "gguf" or "transformers"
_caption_state = {"loaded": False, "ok": False, "model": None, "tokenizer": None,
                  "processor": None, "backend": "",
                  "gguf_model_path": None, "gguf_mmproj_path": None}


def _in_hf_space() -> bool:
    """True when running inside a Hugging Face Space. HF sets SPACE_ID (and
    related SPACE_* vars) automatically, so we use that to autoswitch to the
    GGUF backend. Force on/off with SOZAI_USE_GGUF=1/0 if needed."""
    forced = os.environ.get("SOZAI_USE_GGUF", "").strip()
    if forced in ("1", "true", "True", "yes"):
        return True
    if forced in ("0", "false", "False", "no"):
        return False
    return bool(os.environ.get("SPACE_ID") or os.environ.get("SPACE_HOST"))


def _load_caption_model_gguf() -> bool:
    """Prepare the MiniCPM-V (vision) GGUF backend. Used inside HF Spaces.

    On ZeroGPU, CUDA must NOT be initialized in the main process — so this does
    NOT build the llama.cpp model here. It only (a) checks the deps import, and
    (b) downloads the LM + mmproj GGUF files (CPU/network only). The actual
    `Llama(...)` construction happens lazily in `_ensure_gguf_model()`, which is
    called from inside the @spaces.GPU worker where a GPU is attached.
    """
    try:
        from huggingface_hub import hf_hub_download
        import llama_cpp  # noqa: F401  (import check only)
        from llama_cpp.llama_chat_format import MiniCPMv26ChatHandler  # noqa: F401

        model_path = hf_hub_download(repo_id=CAPTION_GGUF_REPO, filename=CAPTION_GGUF_FILE)
        mmproj_path = hf_hub_download(repo_id=CAPTION_GGUF_REPO, filename=CAPTION_GGUF_MMPROJ)
        # stash the resolved file paths; the model object is built later, lazily
        _caption_state.update(model=None, tokenizer=None, processor=None,
                              gguf_model_path=model_path, gguf_mmproj_path=mmproj_path,
                              ok=True, backend="gguf")
        print(f"[autocaption] GGUF vision backend ready "
              f"({CAPTION_GGUF_REPO}: {CAPTION_GGUF_FILE} + {CAPTION_GGUF_MMPROJ}); "
              f"model will load on first use inside the GPU worker")
        return True
    except Exception as exc:  # llama_cpp/hub missing, download failed…
        print(f"[autocaption] GGUF backend unavailable, will try transformers: {exc}")
        return False


# Lock guarding the lazy in-worker GGUF build. On ZeroGPU each worker builds its
# own copy once; this prevents a double-build if two calls race.
_gguf_build_lock = threading.Lock()


def _ensure_gguf_model():
    """Build (once) and return the llama.cpp `Llama` model. MUST be called from
    inside a @spaces.GPU context on ZeroGPU so the CUDA context is created while
    a GPU is attached. Idempotent and thread-safe; caches into _caption_state.
    Returns None on failure (caller falls back)."""
    model = _caption_state.get("model")
    if model is not None:
        return model
    with _gguf_build_lock:
        model = _caption_state.get("model")
        if model is not None:
            return model
        try:
            from llama_cpp import Llama
            from llama_cpp.llama_chat_format import MiniCPMv26ChatHandler

            model_path = _caption_state.get("gguf_model_path")
            mmproj_path = _caption_state.get("gguf_mmproj_path")
            if not model_path or not mmproj_path:
                # files weren't resolved in the main-process prep step
                from huggingface_hub import hf_hub_download
                model_path = model_path or hf_hub_download(
                    repo_id=CAPTION_GGUF_REPO, filename=CAPTION_GGUF_FILE)
                mmproj_path = mmproj_path or hf_hub_download(
                    repo_id=CAPTION_GGUF_REPO, filename=CAPTION_GGUF_MMPROJ)
            chat_handler = MiniCPMv26ChatHandler(clip_model_path=mmproj_path, verbose=False)
            llm = Llama(
                model_path=model_path,
                chat_handler=chat_handler,
                n_ctx=4096,            # vision tokens need a larger context window
                n_gpu_layers=-1,       # offload all layers (GPU is attached here)
                verbose=False,
            )
            _caption_state["model"] = llm
            print(f"[autocaption] GGUF model built in worker ({CAPTION_GGUF_FILE})")
            return llm
        except Exception as exc:  # noqa: BLE001
            print(f"[autocaption] GGUF in-worker build failed: {exc}")
            return None


def _load_caption_model_transformers() -> bool:
    """Load MiniCPM-V via transformers. Used outside HF Spaces.

    MiniCPM-V exposes a `.chat(image=..., msgs=...)` helper (trust_remote_code),
    which we use so this path is ALSO image-grounded, not text-only.
    """
    try:
        import torch
        from transformers import AutoModel, AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(AUTOCAPTION_MODEL_ID, trust_remote_code=True)
        model = AutoModel.from_pretrained(
            AUTOCAPTION_MODEL_ID,
            trust_remote_code=True,
            attn_implementation="sdpa",
            torch_dtype=torch.bfloat16,
        )
        model = model.eval()
        if torch.cuda.is_available():
            model = model.cuda()
        _caption_state.update(model=model, tokenizer=tokenizer, processor=None,
                              ok=True, backend="transformers")
        print(f"[autocaption] transformers vision backend ready ({AUTOCAPTION_MODEL_ID})")
        return True
    except Exception as exc:  # transformers/torch/model missing, no GPU, etc.
        print(f"[autocaption] transformers backend unavailable: {exc}")
        return False


def _load_caption_model():
    """Load and cache the captioning model once. Returns True on success.

    Autoswitches backend: inside a Hugging Face Space we use the GGUF
    (llama.cpp) build; everywhere else we use the normal transformers model.
    If the preferred backend fails to load we fall back to the other one before
    giving up, so the feature still works wherever it can.
    """
    if _caption_state["loaded"]:
        return _caption_state["ok"]
    with _caption_lock:
        if _caption_state["loaded"]:
            return _caption_state["ok"]
        try:
            if _in_hf_space():
                ok = _load_caption_model_gguf() or _load_caption_model_transformers()
            else:
                ok = _load_caption_model_transformers() or _load_caption_model_gguf()
            _caption_state["ok"] = ok
        finally:
            _caption_state["loaded"] = True
    return _caption_state["ok"]


def _resolve_image_for_model(image_path: str):
    """Resolve whatever the frontend sends into a real local file path the
    vision model can open. Handles three cases:
      1. a base64 data URL  (data:image/...;base64,...) — what room uploads send
      2. a bundled asset URL (/assets/...) or zip entry — sample photos
      3. an already-local file path
    Returns a filesystem path, or None if nothing usable was provided.

    This is the key to image-grounded captions: room photos arrive as data URLs,
    NOT /assets/ paths, so without case (1) the model would never see the pixels
    and would fall back to filename/alt-based guesses.
    """
    if not image_path:
        return None
    s = image_path.strip()
    # 1) data URL → write bytes to a temp file
    if s.startswith("data:"):
        try:
            import base64
            import tempfile
            header, b64 = (s.split(",", 1) if "," in s else ("", s))
            mime = ""
            if header.startswith("data:") and ";" in header:
                mime = header[len("data:"):header.index(";")]
            ext = mimetypes.guess_extension(mime) or ".png"
            raw = base64.b64decode(b64)
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
            tmp.write(raw)
            tmp.close()
            return tmp.name
        except Exception as exc:  # noqa: BLE001
            print(f"[autocaption] could not decode data URL: {exc}")
            return None
    # 2) asset / zip path
    resolved = _resolve_asset_path(s)
    if resolved:
        return resolved
    # 3) already a local file
    if os.path.isfile(s):
        return s
    return None


def _resolve_asset_path(image_path: str):
    """Map a frontend asset URL (/assets/...) to a real file or zip bytes."""
    if not image_path:
        return None
    rel = image_path.split("?")[0]
    if rel.startswith("/assets/"):
        rel = rel[len("/assets/"):]
    rel = _safe_rel(rel)
    if PUBLIC_DIR is not None:
        fp = os.path.join(PUBLIC_DIR, *rel.split("/"))
        if os.path.isfile(fp):
            return fp
    if ZIP_PATH is not None:
        try:
            import tempfile

            data = _read_from_zip(rel)
            suffix = os.path.splitext(rel)[1] or ".png"
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
            tmp.write(data)
            tmp.close()
            return tmp.name
        except KeyError:
            return None
    return None


def _load_pil_image(local_img: str):
    """Open a resolved asset path as an RGB PIL image, or None on failure."""
    if not local_img or not os.path.isfile(local_img):
        return None
    try:
        from PIL import Image
        return Image.open(local_img).convert("RGB")
    except Exception as exc:  # noqa: BLE001
        print(f"[autocaption] could not open image {local_img}: {exc}")
        return None


def _image_data_uri(local_img: str):
    """Encode a local image file as a data: URI for the GGUF chat handler."""
    if not local_img or not os.path.isfile(local_img):
        return None
    try:
        import base64
        mime = mimetypes.guess_type(local_img)[0] or "image/png"
        with open(local_img, "rb") as fh:
            b64 = base64.b64encode(fh.read()).decode()
        return f"data:{mime};base64,{b64}"
    except Exception as exc:  # noqa: BLE001
        print(f"[autocaption] could not encode image {local_img}: {exc}")
        return None


def _fallback_caption(alt: str) -> str:
    base = (alt or "A moment worth remembering").strip().rstrip(".")
    return base[:1].upper() + base[1:] + "."


def _fallback_tags(alt: str) -> str:
    """Comma-separated tag fallback derived from the alt text."""
    words = re.findall(r"[A-Za-z0-9]+", (alt or "").lower())
    stop = {"the", "a", "an", "of", "and", "in", "on", "at", "to", "with", "photo", "image"}
    tags = [w for w in words if w not in stop and len(w) > 2][:6]
    if not tags:
        tags = ["memory", "moment"]
    return ", ".join(dict.fromkeys(tags))   # de-dupe, keep order


def _normalize_tags(text: str) -> str:
    """Coerce model output into clean comma-separated tags. Splits on commas,
    newlines, semicolons or bullet markers; trims; de-dupes; lowercases. Returns
    "" if nothing usable is found (caller treats that as an error)."""
    if not text:
        return ""
    # split on common separators the model might emit
    parts = re.split(r"[,\n;•·\-\u2022]+", text)
    out = []
    seen = set()
    for p in parts:
        tag = re.sub(r"^[\s\d\.\)\(]+", "", p).strip().strip('"').strip("'").lower()
        tag = re.sub(r"\s+", " ", tag)
        if tag and tag not in seen and len(tag) <= 40:
            seen.add(tag)
            out.append(tag)
    return ", ".join(out)


def _is_comma_separated(text: str) -> bool:
    """A valid tag string has at least one tag and, if multiple, comma sep."""
    t = (text or "").strip()
    if not t:
        return False
    # single tag (no spaces-as-separator confusion) or multiple comma-separated
    if "," in t:
        return all(seg.strip() for seg in t.split(","))
    return bool(t)


@app.api(name="caption_model")
def caption_model() -> str:
    """Return the id of the vision-language model used for auto-captioning,
    so the UI can show the user which model is currently in use."""
    return AUTOCAPTION_MODEL_ID


# --------------------------------------------------------------------------- #
# Observability with Arize Phoenix + OpenTelemetry (no API keys).
#
# Phoenix runs IN-PROCESS — it launches automatically when app.py starts, so
# there is no separate `phoenix serve` to run. Every model call (caption / NSFW
# / ASR) is emitted as an OpenTelemetry span with OpenInference semantic
# attributes and shows up in the Phoenix UI, which is embedded directly in the
# app's "LLM Trace" section (via iframe). We also keep a small in-memory ring
# buffer so the in-app list works even if Phoenix can't start.
#
# Port can be overridden with PHOENIX_PORT (default 6006).
# --------------------------------------------------------------------------- #
_trace_lock = threading.Lock()
_trace_buffer: list = []          # most-recent-last ring buffer of spans
_TRACE_MAX = 200
_otel_tracer = None
_otel_loaded = False
_otel_provider = None
_phoenix_session = None
_phoenix_launched = False
_phoenix_lock = threading.Lock()
PHOENIX_PORT = int(os.environ.get("PHOENIX_PORT", "6006"))
PHOENIX_HOST = os.environ.get("PHOENIX_HOST", "localhost")
PHOENIX_PROJECT = os.environ.get("PHOENIX_PROJECT_NAME", "sozai")
# the collector endpoint OTel sends spans to (same in-process server)
PHOENIX_ENDPOINT = os.environ.get(
    "PHOENIX_COLLECTOR_ENDPOINT", f"http://{PHOENIX_HOST}:{PHOENIX_PORT}"
)


def _launch_phoenix():
    """Start the Phoenix app in-process (once). Returns the UI base URL, or
    None if Phoenix isn't installed / failed to start. Idempotent + thread-safe.
    This means the trace UI is available the moment app.py is running — no
    separate `phoenix serve` needed.

    Host/port are configured via the PHOENIX_HOST / PHOENIX_PORT environment
    variables (the launch_app host/port kwargs are deprecated)."""
    global _phoenix_session, _phoenix_launched
    if _phoenix_launched:
        return _phoenix_session_url()
    with _phoenix_lock:
        if _phoenix_launched:
            return _phoenix_session_url()
        _phoenix_launched = True
        # Phoenix reads these at import/launch time — set them before importing.
        os.environ.setdefault("PHOENIX_HOST", PHOENIX_HOST)
        os.environ.setdefault("PHOENIX_PORT", str(PHOENIX_PORT))
        try:
            import phoenix as px
            # launch_app() picks up PHOENIX_HOST / PHOENIX_PORT from the env.
            # If something is already serving on the port, it returns a session
            # pointed at it.
            _phoenix_session = px.launch_app()
            print(f"[phoenix] in-process UI at {_phoenix_session_url()}")
        except Exception as exc:  # noqa: BLE001
            print(f"[phoenix] could not launch in-process app: {exc}")
            _phoenix_session = None
    return _phoenix_session_url()


def _phoenix_session_url():
    if _phoenix_session is not None:
        try:
            return _phoenix_session.url
        except Exception:  # noqa: BLE001
            pass
    return f"http://{PHOENIX_HOST}:{PHOENIX_PORT}"


def _get_tracer():
    """Lazily launch in-process Phoenix, register the OTel tracer provider, and
    return a tracer. Falls back to a raw OTLP/HTTP exporter, then to the local
    buffer only. Returns None if OpenTelemetry isn't installed."""
    global _otel_tracer, _otel_loaded
    if _otel_loaded:
        return _otel_tracer
    _otel_loaded = True
    # make sure the embedded Phoenix collector is up first
    _launch_phoenix()
    # Preferred: phoenix.otel.register — Phoenix-aware OTel defaults, no keys.
    try:
        from phoenix.otel import register
        tracer_provider = register(
            project_name=PHOENIX_PROJECT,
            endpoint=PHOENIX_ENDPOINT.rstrip("/") + "/v1/traces",
            batch=True,                     # use a BatchSpanProcessor (recommended)
            set_global_tracer_provider=True,
        )
        # batched spans may still be queued at exit — flush them on shutdown.
        global _otel_provider
        _otel_provider = tracer_provider
        try:
            import atexit
            atexit.register(lambda: _otel_provider and _otel_provider.shutdown())
        except Exception:  # noqa: BLE001
            pass
        _otel_tracer = tracer_provider.get_tracer(__name__)
        print(f"[otel] Phoenix tracing enabled → {PHOENIX_ENDPOINT} (project '{PHOENIX_PROJECT}')")
        return _otel_tracer
    except Exception as exc:  # noqa: BLE001
        print(f"[otel] phoenix.otel unavailable ({exc}); trying raw OTLP exporter")
    # Fallback: plain OpenTelemetry with an OTLP/HTTP exporter to Phoenix.
    try:
        from opentelemetry import trace as _t
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import BatchSpanProcessor
        from opentelemetry.sdk.resources import Resource
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
        provider = TracerProvider(resource=Resource.create({"service.name": PHOENIX_PROJECT}))
        provider.add_span_processor(BatchSpanProcessor(
            OTLPSpanExporter(endpoint=PHOENIX_ENDPOINT.rstrip("/") + "/v1/traces")))
        _t.set_tracer_provider(provider)
        _otel_tracer = _t.get_tracer(__name__)
        print(f"[otel] OTLP tracing enabled → {PHOENIX_ENDPOINT}")
        return _otel_tracer
    except Exception as exc:  # noqa: BLE001
        print(f"[otel] OpenTelemetry unavailable, using local trace buffer only: {exc}")
        _otel_tracer = None
        return None


def _record_trace(event: dict):
    """Append a trace event to the local ring buffer (newest last)."""
    event = dict(event)
    event.setdefault("ts", time.time())
    with _trace_lock:
        _trace_buffer.append(event)
        if len(_trace_buffer) > _TRACE_MAX:
            del _trace_buffer[: len(_trace_buffer) - _TRACE_MAX]


# OpenInference semantic attribute keys (string fallbacks if the package that
# defines them isn't importable — keeps tracing working without it). These mirror
# the attributes Phoenix reads for its span detail + metrics.
_OI = {
    "kind": "openinference.span.kind",
    "input": "input.value",
    "input_mime": "input.mime_type",
    "output": "output.value",
    "output_mime": "output.mime_type",
    "model": "llm.model_name",
    "provider": "llm.provider",
    "system": "llm.system",
    "invocation": "llm.invocation_parameters",
    "finish": "llm.finish_reason",
    "tok_prompt": "llm.token_count.prompt",
    "tok_completion": "llm.token_count.completion",
    "tok_total": "llm.token_count.total",
    "metadata": "metadata",
    "session": "session.id",
}
try:
    from openinference.semconv.trace import SpanAttributes as _OISpan
    _OI.update({
        "kind": _OISpan.OPENINFERENCE_SPAN_KIND,
        "input": _OISpan.INPUT_VALUE,
        "output": _OISpan.OUTPUT_VALUE,
        "model": _OISpan.LLM_MODEL_NAME,
    })
    for k, attr in (("input_mime", "INPUT_MIME_TYPE"), ("output_mime", "OUTPUT_MIME_TYPE"),
                    ("provider", "LLM_PROVIDER"), ("system", "LLM_SYSTEM"),
                    ("invocation", "LLM_INVOCATION_PARAMETERS"), ("finish", "LLM_FINISH_REASON"),
                    ("tok_prompt", "LLM_TOKEN_COUNT_PROMPT"), ("tok_completion", "LLM_TOKEN_COUNT_COMPLETION"),
                    ("tok_total", "LLM_TOKEN_COUNT_TOTAL"), ("metadata", "METADATA"),
                    ("session", "SESSION_ID")):
        if hasattr(_OISpan, attr):
            _OI[k] = getattr(_OISpan, attr)
except Exception:  # noqa: BLE001
    pass
# back-compat aliases used elsewhere
_OI_KIND = _OI["kind"]; _OI_INPUT = _OI["input"]; _OI_OUTPUT = _OI["output"]; _OI_MODEL = _OI["model"]


def _provider_for(model_id: str) -> str:
    """Best-effort provider label from a model id (for the trace 'provider')."""
    m = (model_id or "").lower()
    if m.startswith("nvidia/") or "nemotron" in m or "nemo" in m:
        return "nvidia"
    if "minicpm" in m or m.startswith("openbmb/"):
        return "openbmb"
    if "falconsai" in m or "nsfw" in m:
        return "falconsai"
    if "/" in m:
        return m.split("/", 1)[0]
    return "local"


def _count_tokens(text: str) -> int:
    """Real token count using the caption model's tokenizer when available;
    otherwise a clearly-approximate word-based estimate."""
    text = text or ""
    if not text:
        return 0
    tok = _caption_state.get("tokenizer") if isinstance(_caption_state, dict) else None
    if tok is not None:
        try:
            return len(tok.encode(text))
        except Exception:  # noqa: BLE001
            try:
                return len(tok(text)["input_ids"])
            except Exception:  # noqa: BLE001
                pass
    return int(round(len(text.split()) / 0.75))



class traced:
    """Context manager that times a model call, records it to the local buffer,
    and emits an OpenTelemetry span (OpenInference 'LLM' kind) to Phoenix.

    Usage:
        with traced("autocaption", model=MODEL_ID, file="x.jpg",
                     meta={"prompt": p}) as t:
            ...
            t.set_output(text)
    """

    def __init__(self, name, model="", file="", meta=None, invocation=None, system=""):
        self.name = name
        self.model = model
        self.file = file
        self.meta = meta or {}
        self.invocation = invocation or {}   # llm.invocation_parameters
        self.system = system                 # llm.system (system prompt, if any)
        self.output = None
        self.finish_reason = "stop"
        self.error = None
        self._t0 = None
        self._span_cm = None
        self._span = None
        self.trace_id = ""
        self.span_id = ""

    def __enter__(self):
        self._t0 = time.time()
        tracer = _get_tracer()
        if tracer is not None:
            try:
                self._span_cm = tracer.start_as_current_span(f"sozai.{self.name}")
                self._span = self._span_cm.__enter__()
                self._span.set_attribute(_OI["kind"], "LLM")
                if self.model:
                    self._span.set_attribute(_OI["model"], self.model)
                    self._span.set_attribute(_OI["provider"], _provider_for(self.model))
                if self.system:
                    self._span.set_attribute(_OI["system"], str(self.system)[:1000])
                if self.invocation:
                    try:
                        self._span.set_attribute(_OI["invocation"], json.dumps(self.invocation)[:1000])
                    except Exception:  # noqa: BLE001
                        pass
                if self.file:
                    self._span.set_attribute("sozai.file", self.file)
                # input value + mime for Phoenix display
                try:
                    self._span.set_attribute(_OI["input"], json.dumps({"file": self.file, **self.meta})[:2000])
                    self._span.set_attribute(_OI["input_mime"], "application/json")
                except Exception:  # noqa: BLE001
                    pass
                for k, v in self.meta.items():
                    try:
                        self._span.set_attribute(f"sozai.{k}", str(v)[:500])
                    except Exception:  # noqa: BLE001
                        pass
                # capture the OTel trace/span ids so the in-app view can show them
                try:
                    ctx = self._span.get_span_context()
                    self.trace_id = format(ctx.trace_id, "032x")
                    self.span_id = format(ctx.span_id, "016x")
                except Exception:  # noqa: BLE001
                    pass
            except Exception as exc:  # noqa: BLE001
                print(f"[otel] span start failed: {exc}")
                self._span = None
        if not self.span_id:
            self.span_id = uuid.uuid4().hex[:16]
        if not self.trace_id:
            self.trace_id = uuid.uuid4().hex
        return self

    def set_output(self, output, finish_reason="stop"):
        self.output = output
        self.finish_reason = finish_reason

    def __exit__(self, exc_type, exc, tb):
        t_end = time.time()
        dur_ms = round((t_end - self._t0) * 1000.0, 1) if self._t0 else None
        if exc is not None:
            self.error = f"{exc_type.__name__}: {exc}" if exc_type else str(exc)
            self.finish_reason = "error"
        try:
            input_payload = json.dumps({"file": self.file, **self.meta})
        except Exception:  # noqa: BLE001
            input_payload = str({"file": self.file, **self.meta})
        out_str = str(self.output) if self.output is not None else ""
        # real token counts via the model tokenizer when available, else estimate
        tok_in = _count_tokens(input_payload)
        tok_out = _count_tokens(out_str)
        tok_total = tok_in + tok_out
        tokens_per_sec = round(tok_out / (dur_ms / 1000.0), 1) if (dur_ms and tok_out) else 0
        provider = _provider_for(self.model)
        _record_trace({
            "id": self.span_id,
            "span_id": self.span_id,
            "trace_id": self.trace_id,
            "name": self.name,
            "kind": "LLM",
            "model": self.model,
            "provider": provider,
            "system": (str(self.system)[:400] if self.system else None),
            "invocation": (self.invocation or None),
            "finish_reason": self.finish_reason,
            "file": self.file,
            "meta": self.meta,
            "input": input_payload[:2000],
            "input_mime": "application/json",
            "output": (out_str[:2000] if out_str else None),
            "output_mime": "text/plain",
            "error": self.error,
            "duration_ms": dur_ms,
            "started": self._t0,
            "ended": t_end,
            "tokens_in": tok_in,
            "tokens_out": tok_out,
            "tokens_total": tok_total,
            "tokens_per_sec": tokens_per_sec,
            "cost_usd": 0.0,                 # local open-source models — no API cost
            "status": "error" if self.error else "ok",
        })
        if self._span is not None:
            try:
                if self.output is not None:
                    self._span.set_attribute(_OI["output"], str(self.output)[:2000])
                    self._span.set_attribute(_OI["output_mime"], "text/plain")
                self._span.set_attribute(_OI["finish"], self.finish_reason)
                if dur_ms is not None:
                    self._span.set_attribute("sozai.duration_ms", dur_ms)
                self._span.set_attribute(_OI["tok_prompt"], tok_in)
                self._span.set_attribute(_OI["tok_completion"], tok_out)
                self._span.set_attribute(_OI["tok_total"], tok_total)
                if self.error:
                    try:
                        from opentelemetry.trace import Status, StatusCode
                        self._span.set_status(Status(StatusCode.ERROR, self.error))
                    except Exception:  # noqa: BLE001
                        pass
                    self._span.set_attribute("sozai.error", self.error)
            except Exception:  # noqa: BLE001
                pass
            try:
                self._span_cm.__exit__(exc_type, exc, tb)
            except Exception:  # noqa: BLE001
                pass
        return False  # never suppress exceptions


@app.get("/api/trace")
def get_trace(limit: int = 50):
    """Return recent model-call traces plus summary stats for the in-app trace
    view. Distributed traces also go to the in-process Phoenix UI."""
    tracer_ok = _get_tracer() is not None
    with _trace_lock:
        all_items = list(_trace_buffer)
        items = _trace_buffer[-max(1, min(limit, _TRACE_MAX)):]

    # aggregate stats across everything in the buffer
    total = len(all_items)
    errors = sum(1 for e in all_items if e.get("status") == "error")
    durations = [e["duration_ms"] for e in all_items if isinstance(e.get("duration_ms"), (int, float)) and e.get("duration_ms")]
    avg_ms = round(sum(durations) / len(durations), 1) if durations else 0

    def _pct(sorted_vals, q):
        if not sorted_vals:
            return 0
        return round(sorted_vals[min(len(sorted_vals) - 1, int(len(sorted_vals) * q))], 1)
    s = sorted(durations)
    p50_ms, p95_ms, p99_ms = _pct(s, 0.50), _pct(s, 0.95), _pct(s, 0.99)

    tokens_prompt = sum(int(e.get("tokens_in") or 0) for e in all_items)
    tokens_completion = sum(int(e.get("tokens_out") or 0) for e in all_items)
    tokens_total = tokens_prompt + tokens_completion
    tps_vals = [e["tokens_per_sec"] for e in all_items if isinstance(e.get("tokens_per_sec"), (int, float)) and e.get("tokens_per_sec")]
    avg_tps = round(sum(tps_vals) / len(tps_vals), 1) if tps_vals else 0
    cost_total = round(sum(float(e.get("cost_usd") or 0) for e in all_items), 4)

    by_type: dict = {}
    by_kind: dict = {}
    for e in all_items:
        nm = e.get("name", "?")
        b = by_type.setdefault(nm, {"count": 0, "errors": 0, "dur": [], "tok": 0})
        b["count"] += 1
        if e.get("status") == "error":
            b["errors"] += 1
        if isinstance(e.get("duration_ms"), (int, float)) and e.get("duration_ms"):
            b["dur"].append(e["duration_ms"])
        b["tok"] += int(e.get("tokens_total") or 0)
        k = e.get("kind", "LLM")
        by_kind[k] = by_kind.get(k, 0) + 1
    by_type_out = {
        nm: {
            "count": b["count"],
            "errors": b["errors"],
            "avg_ms": round(sum(b["dur"]) / len(b["dur"]), 1) if b["dur"] else 0,
            "tokens": b["tok"],
        }
        for nm, b in by_type.items()
    }

    return {
        "phoenix": tracer_ok,
        "phoenix_url": "/phoenix/" if _phoenix_session is not None else None,
        "phoenix_direct": _phoenix_session_url() if _phoenix_session is not None else None,
        "endpoint": PHOENIX_ENDPOINT if tracer_ok else None,
        "project": PHOENIX_PROJECT,
        "stats": {
            "total": total,
            "errors": errors,
            "error_rate": round(errors / total, 3) if total else 0,
            "avg_ms": avg_ms,
            "p50_ms": p50_ms,
            "p95_ms": p95_ms,
            "p99_ms": p99_ms,
            "tokens_total": tokens_total,
            "tokens_prompt": tokens_prompt,
            "tokens_completion": tokens_completion,
            "avg_tokens_per_sec": avg_tps,
            "cost_usd": cost_total,
            "by_type": by_type_out,
            "by_kind": by_kind,
        },
        "events": list(reversed(items)),  # newest first for display
    }


# --- Same-origin reverse proxy for the embedded Phoenix UI ------------------ #
# Phoenix runs in-process on its own port (6006). To embed it in an iframe
# without cross-origin / X-Frame-Options problems, we proxy it under /phoenix/
# on the app's own origin. Frame-blocking headers are stripped on the way back.
import urllib.request as _urlreq
import urllib.error as _urlerr

_HOP_BY_HOP = {
    "content-encoding", "content-length", "transfer-encoding", "connection",
    "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailers",
    "upgrade", "x-frame-options", "content-security-policy",
}


_PHX_PREFIX = "/phoenix"


def _rewrite_phoenix_body(data: bytes, ctype: str) -> bytes:
    """Phoenix's SPA references assets/APIs with root-absolute paths
    (/assets/..., /graphql, /v1/...). Served under /phoenix/, those would
    resolve against the app origin and 404 — leaving the iframe blank. Rewrite
    HTML and JS so everything stays under the proxy prefix."""
    is_html = "text/html" in ctype
    is_js = ("javascript" in ctype) or ("text/css" in ctype)
    if not (is_html or is_js):
        return data
    try:
        text = data.decode("utf-8")
    except Exception:  # noqa: BLE001
        return data  # binary; leave untouched

    if is_html:
        # 1) a <base> so relative URLs resolve under /phoenix/
        if "<base " not in text:
            text = re.sub(r"(<head[^>]*>)", r'\1<base href="' + _PHX_PREFIX + '/">',
                          text, count=1, flags=re.IGNORECASE)
        # 2) prefix root-absolute attribute URLs:  href="/x"  src="/x"
        #    (skip protocol-relative //, data:, and already-prefixed /phoenix)
        text = re.sub(r'(\b(?:href|src)=")/(?!/|phoenix/)',
                      r"\1" + _PHX_PREFIX + "/", text)

    # 3) prefix root-absolute paths used in JS/CSS (fetch, websocket, url()).
    #    Phoenix hits "/graphql", "/v1/...", "/arize_phoenix_version", etc.
    for p in ("/graphql", "/v1/", "/arize_phoenix_version", "/exports",
              "/assets/", "/healthz", "/readyz"):
        # in quotes: "/graphql"  '/v1/...'
        text = text.replace('"' + p, '"' + _PHX_PREFIX + p)
        text = text.replace("'" + p, "'" + _PHX_PREFIX + p)
        # css url(/assets/...)
        text = text.replace("(" + p, "(" + _PHX_PREFIX + p)
    return text.encode("utf-8")


async def _phoenix_proxy(request: Request, path: str):
    base = _phoenix_session_url().rstrip("/")
    target = base + "/" + path
    if request.url.query:
        target += "?" + request.url.query
    body = await request.body()
    method = request.method
    headers = {k: v for k, v in request.headers.items()
               if k.lower() not in ("host", "content-length", "accept-encoding")}
    headers["accept-encoding"] = "identity"
    req = _urlreq.Request(target, data=(body or None), method=method, headers=headers)
    try:
        with _urlreq.urlopen(req, timeout=30) as resp:
            data = resp.read()
            status = resp.status
            ctype = resp.headers.get("Content-Type", "application/octet-stream")
            out_headers = {}
            for k, v in resp.headers.items():
                if k.lower() in _HOP_BY_HOP:
                    continue
                out_headers[k] = v
        data = _rewrite_phoenix_body(data, ctype)
        return Response(content=data, status_code=status, media_type=ctype, headers=out_headers)
    except _urlerr.HTTPError as e:
        try:
            err_body = e.read()
        except Exception:  # noqa: BLE001
            err_body = b""
        return Response(content=err_body, status_code=e.code,
                        media_type=e.headers.get("Content-Type", "text/plain"))
    except Exception as exc:  # noqa: BLE001
        return Response(content=f"Phoenix UI not reachable: {exc}".encode(),
                        status_code=502, media_type="text/plain")


@app.get("/phoenix")
@app.get("/phoenix/")
async def phoenix_root(request: Request):
    return await _phoenix_proxy(request, "")


@app.get("/phoenix/{path:path}")
async def phoenix_proxy_get(request: Request, path: str):
    return await _phoenix_proxy(request, path)


@app.post("/phoenix/{path:path}")
async def phoenix_proxy_post(request: Request, path: str):
    return await _phoenix_proxy(request, path)


# Cesium ion access token. NEVER hardcode this in the HTML — it is read from
# the environment so it can be rotated without touching the frontend. Set it
# with:  export SOZAI_CESIUM_ION_TOKEN="your-token"  (or put it in a .env file).
from dotenv import load_dotenv
load_dotenv()
CESIUM_ION_TOKEN = os.environ.get("SOZAI_CESIUM_ION_TOKEN", "")
# Free Stadia Maps API key (https://client.stadiamaps.com) for the Stamen
# Watercolor basemap on the Cesium globe. Optional — without it the globe keeps
# its default imagery.
STADIA_MAPS_KEY = os.environ.get("SOZAI_STADIA_KEY", "")


@app.get("/api/config")
def client_config():
    """Public, non-secret-ish runtime config for the frontend (e.g. whether a
    Cesium ion token is configured, and the token itself if so). The token is
    only as protected as the page that needs it; keeping it server-side env-only
    means it can be rotated in one place."""
    import platform
    return {
        "cesiumIonToken": CESIUM_ION_TOKEN,
        "cesiumEnabled": bool(CESIUM_ION_TOKEN),
        "stadiaMapsKey": STADIA_MAPS_KEY,
        "nsfwEnabled": True,            # the gate always runs (open if model absent)
        "asrEnabled": True,             # the endpoint always exists (may be unavailable)
        "asrModel": ASR_MODEL_ID,
        "asrOsSupported": platform.system().lower() == "linux",
        "asrLanguages": ASR_LANGUAGE_CHOICES,
        "asrProfiles": list(ASR_CHUNK_PROFILES.keys()),
        "asrDefaultProfile": ASR_DEFAULT_PROFILE,
        "captionModel": AUTOCAPTION_MODEL_ID,
        # "Proceed to Development" watercolour: the endpoint always exists; this
        # only reports whether torch/diffusers are present (a CUDA-free check —
        # the LoRA is optional, stage 1 runs base img2img). The ~14 GB pipeline
        # still builds lazily on first use inside the GPU worker.
        "img2imgEnabled": _img2img_prep(),
        "img2imgModel": IMG2IMG_MODEL_ID,
        "phoenix": _get_tracer() is not None,
    }


# --------------------------------------------------------------------------- #
# NSFW image gate (item 1). Falconsai/nsfw_image_detection via transformers.
# If transformers/torch/model are missing, the gate is OPEN (score 0.0) so
# uploads are never blocked by a missing optional dependency.
# --------------------------------------------------------------------------- #
NSFW_THRESHOLD = float(os.environ.get("SOZAI_NSFW_THRESHOLD", "0.80"))
NSFW_MODEL_ID = "Falconsai/nsfw_image_detection"
_nsfw_lock = threading.Lock()
_nsfw_state = {"loaded": False, "pipe": None, "available": False}


def _get_nsfw_pipe():
    """Check that the NSFW classifier CAN be built, without initializing CUDA in
    the main process (forbidden on ZeroGPU). This only verifies the import +
    that transformers is present; the actual pipeline is constructed lazily
    inside the @spaces.GPU worker by _ensure_nsfw_pipe(). Returns True when the
    gate is available, None/False when it can't be (gate then fails OPEN)."""
    if _nsfw_state["loaded"]:
        return _nsfw_state["available"]
    with _nsfw_lock:
        if _nsfw_state["loaded"]:
            return _nsfw_state["available"]
        try:
            import transformers  # noqa: F401  (import check only — no CUDA here)
            _nsfw_state["available"] = True
        except Exception as exc:  # noqa: BLE001
            print(f"[nsfw] transformers unavailable, gate open: {exc}")
            _nsfw_state["available"] = False
        finally:
            _nsfw_state["loaded"] = True
    return _nsfw_state["available"]


_nsfw_build_lock = threading.Lock()


def _ensure_nsfw_pipe():
    """Build (once) and return the transformers image-classification pipeline.
    MUST run inside a @spaces.GPU context on ZeroGPU so CUDA initializes with a
    GPU attached. Idempotent + thread-safe; caches into _nsfw_state. Returns the
    pipeline, or None on failure (caller then treats the image as safe)."""
    pipe = _nsfw_state.get("pipe")
    if pipe is not None:
        return pipe
    with _nsfw_build_lock:
        pipe = _nsfw_state.get("pipe")
        if pipe is not None:
            return pipe
        try:
            from transformers import pipeline
            import torch
            device = 0 if torch.cuda.is_available() else -1
            pipe = pipeline("image-classification", model=NSFW_MODEL_ID, device=device)
            _nsfw_state["pipe"] = pipe
            print(f"[nsfw] classifier built in worker (device={'cuda' if device == 0 else 'cpu'})")
            return pipe
        except Exception as exc:  # noqa: BLE001
            print(f"[nsfw] classifier build failed, gate open: {exc}")
            return None


@spaces.GPU(duration=30)
def _nsfw_infer(img):
    """Score one PIL image for NSFW. Builds/uses the pipeline INSIDE this
    @spaces.GPU context so CUDA is only touched with a GPU attached. The image
    (picklable) crosses the boundary, never the model."""
    pipe = _ensure_nsfw_pipe()
    if pipe is None:
        return []
    return pipe(img.convert("RGB"))


def _nsfw_score_for(path_or_pil) -> float:
    if not _get_nsfw_pipe():
        return 0.0
    try:
        from PIL import Image
        img = path_or_pil if hasattr(path_or_pil, "convert") else Image.open(path_or_pil)
        for r in _nsfw_infer(img):
            if str(r["label"]).lower() == "nsfw":
                return float(r["score"])
    except Exception as exc:  # noqa: BLE001
        print(f"[nsfw] scoring failed (treating as safe): {exc}")
    return 0.0


@app.api(name="nsfw_check")
def nsfw_check(image_path: str = "") -> dict:
    """Score an uploaded image for NSFW content. Returns {score, blocked,
    threshold, available}. Fails OPEN (never blocks) if the model is absent."""
    fp = _resolve_image_for_model(image_path) or image_path
    available = bool(_get_nsfw_pipe())
    with traced("nsfw_check", model=NSFW_MODEL_ID, file=os.path.basename(image_path or "")) as t:
        score = _nsfw_score_for(fp) if available else 0.0
        t.set_output({"score": round(score, 4)})
    return {
        "score": round(score, 4),
        "blocked": bool(score >= NSFW_THRESHOLD),
        "threshold": NSFW_THRESHOLD,
        "available": available,
    }


@app.post("/api/nsfw")
async def nsfw_rest(request: Request):
    """NSFW gate for client uploads sent as a base64 data URL.
    Body: {data_url: "data:image/...;base64,...", name: "file.jpg"}.
    Fails OPEN on any error so a missing model never blocks uploads."""
    import base64
    try:
        body = await request.json()
    except Exception:  # noqa: BLE001
        return {"score": 0.0, "blocked": False, "available": bool(_get_nsfw_pipe()), "error": "bad request"}
    data_url = (body or {}).get("data_url", "")
    name = (body or {}).get("name", "upload")
    available = bool(_get_nsfw_pipe())
    if not available or not data_url:
        return {"score": 0.0, "blocked": False, "threshold": NSFW_THRESHOLD, "available": available}
    try:
        b64 = data_url.split(",", 1)[1] if "," in data_url else data_url
        raw = base64.b64decode(b64)
        from PIL import Image
        img = Image.open(io.BytesIO(raw))
        with traced("nsfw_check", model=NSFW_MODEL_ID, file=str(name)[:60]) as t:
            score = _nsfw_score_for(img)
            t.set_output({"score": round(score, 4)})
        return {"score": round(score, 4), "blocked": bool(score >= NSFW_THRESHOLD),
                "threshold": NSFW_THRESHOLD, "available": True}
    except Exception as exc:  # noqa: BLE001
        print(f"[nsfw] rest scoring failed (treating as safe): {exc}")
        return {"score": 0.0, "blocked": False, "threshold": NSFW_THRESHOLD, "available": available, "error": str(exc)}


# --------------------------------------------------------------------------- #
# Develop → watercolour via Flux.2-klein + a scene LoRA (in-process img2img).
# This is what "Proceed to Development" triggers: each photo is sent here and
# painted into a watercolour by Flux2KleinPipeline (the source is VAE-encoded
# and used as reference conditioning while the LoRA steers the watercolour
# style).
#
# ZeroGPU port notes (mirrors the caption/NSFW/ASR helpers in this file):
#   * The pipeline is ~14 GB and lives ONLY on the GPU, so it must be built
#     INSIDE a @spaces.GPU context — never in the main process, where CUDA must
#     stay un-initialized. `_img2img_prep()` does the CUDA-free part (import +
#     LoRA file check) so /config and the endpoint can report availability
#     cheaply; `_ensure_img2img_pipe()` does the actual GPU build, lazily, the
#     first time `_img2img_infer` runs in a worker.
#   * @spaces.GPU pickles arguments to/from the worker. The Flux pipe is NOT
#     picklable, so it is fetched from module state inside the worker and never
#     crosses the boundary. Only the PIL image in and the PIL image out cross
#     it (both picklable), exactly like `_nsfw_infer`.
#   * This file is self-contained on the Space (no backend/ package), so the
#     squaring + load + stylize logic from backend/hf/3_infer_img2img.py is
#     inlined here rather than imported.
#
# STAGED ROLLOUT: the scene LoRA is OPTIONAL. Stage 1 (now) is "pure img2img" —
# the base Flux2Klein model develops the photo with a generic watercolour prompt,
# no .safetensors required. Stage 2 is dropping the LoRA file in (or setting
# SOZAI_IMG2IMG_LORA_REPO): it's detected automatically and the trained
# watercolour style + trigger prompt switch on. So a missing LoRA degrades to
# base img2img rather than disabling the feature.
#
# Fails CLOSED (available=false) only when torch/diffusers are absent, so a CPU
# Space or local Mac keeps the original photo instead of erroring/OOMing.
# --------------------------------------------------------------------------- #
IMG2IMG_MODEL_ID = os.environ.get("SOZAI_IMG2IMG_MODEL", "black-forest-labs/FLUX.2-klein-4B")
# Path to the watercolour scene LoRA .safetensors. OPTIONAL: ship the file in the
# repo at this path, or point SOZAI_IMG2IMG_LORA_REPO/_FILE at a Hub repo to have
# it downloaded (CPU/network only) at prep time. Absent = base img2img (stage 1).
IMG2IMG_LORA_PATH = os.environ.get(
    "SOZAI_IMG2IMG_LORA",
    os.path.join(BASE, "data/places/checkpoints/hf-spaces/scene_lora.safetensors"),
)
IMG2IMG_LORA_REPO = os.environ.get("SOZAI_IMG2IMG_LORA_REPO", "")
IMG2IMG_LORA_FILE = os.environ.get("SOZAI_IMG2IMG_LORA_FILE", "scene_lora.safetensors")
IMG2IMG_LORA_SCALE = float(os.environ.get("SOZAI_IMG2IMG_LORA_SCALE", "1.0"))
# Generation defaults (from backend/hf/3_infer_img2img.py).
IMG2IMG_TRIGGER = "wtrclr8"
# Prompt used WITH the LoRA (its trigger token), vs. the base-model fallback used
# in stage 1 when no LoRA is loaded — the trigger token means nothing without it.
IMG2IMG_PROMPT = os.environ.get("SOZAI_IMG2IMG_PROMPT", "WTRCLR8 watercolour painting")
IMG2IMG_BASE_PROMPT = os.environ.get(
    "SOZAI_IMG2IMG_BASE_PROMPT",
    "a soft watercolour painting of this scene, painterly, textured paper, gentle washes",
)
IMG2IMG_SIZE = int(os.environ.get("SOZAI_IMG2IMG_SIZE", "1024"))
IMG2IMG_STEPS = int(os.environ.get("SOZAI_IMG2IMG_STEPS", "8"))
IMG2IMG_GUIDANCE = float(os.environ.get("SOZAI_IMG2IMG_GUIDANCE", "3.5"))
IMG2IMG_SEED = int(os.environ.get("SOZAI_IMG2IMG_SEED", "42"))
_img2img_lock = threading.Lock()
_img2img_state = {"loaded": False, "available": False, "pipe": None,
                  "lora_path": None, "lora_loaded": False,
                  "snapshot_path": None, "snapshot_ready": False, "dl_started": False}


def _img2img_start_bg_download():
    """Download the FLUX weights in the MAIN process (a daemon thread), where
    cache writes succeed. The ZeroGPU worker can request a GPU but CANNOT write
    the HF cache during a download (it fails with 'Permission denied'), so the
    model must be fetched here first; the worker then only READS the cache
    (from_pretrained with local_files_only). Runs once, non-blocking."""
    if _img2img_state.get("dl_started"):
        return
    _img2img_state["dl_started"] = True

    def _run():
        import time
        from huggingface_hub import snapshot_download
        # FLUX.2-klein is ~24 GB across ~25 shards; HF's CDN intermittently drops
        # the connection mid-shard ("peer closed connection"). snapshot_download
        # RESUMES partial files from the cache, so we retry with backoff until the
        # whole repo is present rather than failing the first time a shard drops.
        attempt = 0
        while not _img2img_state.get("snapshot_ready"):
            attempt += 1
            try:
                print(f"[img2img] background download attempt {attempt} of {IMG2IMG_MODEL_ID}…")
                path = snapshot_download(repo_id=IMG2IMG_MODEL_ID, max_workers=2)
                _img2img_state["snapshot_path"] = path
                _img2img_state["snapshot_ready"] = True
                _img2img_state["build_error"] = None
                print(f"[img2img] model cached and ready at {path}")
                return
            except Exception as exc:  # noqa: BLE001 — usually a dropped connection
                _img2img_state["build_error"] = f"download attempt {attempt} failed: {exc}"
                wait = min(60, 5 * attempt)
                print(f"[img2img] download attempt {attempt} failed ({exc}); resuming in {wait}s")
                time.sleep(wait)

    threading.Thread(target=_run, daemon=True).start()


def _img2img_prep() -> bool:
    """CUDA-free availability check, run in the main process. Verifies diffusers
    (with Flux2KleinPipeline) imports; the LoRA is OPTIONAL — if a file is present
    (or downloadable via SOZAI_IMG2IMG_LORA_REPO) its path is cached for stage 2,
    otherwise we fall back to base img2img. Does NOT touch CUDA or build the
    pipeline. Returns True whenever the base feature can run in a GPU worker."""
    if _img2img_state["loaded"]:
        return _img2img_state["available"]
    with _img2img_lock:
        if _img2img_state["loaded"]:
            return _img2img_state["available"]
        try:
            import torch  # noqa: F401  (import check only — no CUDA here)
            from diffusers import Flux2KleinPipeline  # noqa: F401
            # Resolve the LoRA if we can, but its absence is NOT fatal.
            lora_path = IMG2IMG_LORA_PATH if (IMG2IMG_LORA_PATH and
                                              os.path.isfile(IMG2IMG_LORA_PATH)) else None
            if lora_path is None and IMG2IMG_LORA_REPO:
                try:
                    from huggingface_hub import hf_hub_download
                    lora_path = hf_hub_download(repo_id=IMG2IMG_LORA_REPO,
                                                filename=IMG2IMG_LORA_FILE)
                except Exception as lexc:  # noqa: BLE001
                    print(f"[img2img] LoRA download skipped ({lexc}); using base model")
                    lora_path = None
            _img2img_state["lora_path"] = lora_path
            _img2img_state["available"] = True
            print(f"[img2img] available ({IMG2IMG_MODEL_ID}; "
                  f"LoRA={lora_path or 'none — base img2img (stage 1)'})")
            # FLUX is mounted as a read-only volume (HF Space "model volume") at
            # this path → the weights are already present, no download. The GPU
            # worker loads straight from the mount. Falls back to downloading
            # into the cache when the mount isn't there (local dev / no volume).
            mount = os.environ.get("SOZAI_IMG2IMG_MOUNT", "/models/flux")
            if os.path.isdir(mount) and os.listdir(mount):
                _img2img_state["snapshot_path"] = mount
                _img2img_state["snapshot_ready"] = True
                print(f"[img2img] using mounted model at {mount} (no download)")
            else:
                print("[img2img] no mount — pre-fetching weights to cache")
                _img2img_start_bg_download()
        except Exception as exc:  # noqa: BLE001
            print(f"[img2img] unavailable, endpoint will report disabled: {exc}")
            _img2img_state["available"] = False
        finally:
            _img2img_state["loaded"] = True
    return _img2img_state["available"]


_img2img_build_lock = threading.Lock()


def _ensure_img2img_pipe():
    """Build (once) and return the Flux2KleinPipeline on the GPU, applying the
    watercolour LoRA only if one was resolved (stage 2); otherwise it's plain
    base img2img (stage 1). MUST be called from inside a @spaces.GPU context on
    ZeroGPU so the CUDA context is created with a GPU attached. Idempotent +
    thread-safe; caches into _img2img_state. Returns the pipe, or None on failure
    (caller falls back)."""
    pipe = _img2img_state.get("pipe")
    if pipe is not None:
        return pipe
    with _img2img_build_lock:
        pipe = _img2img_state.get("pipe")
        if pipe is not None:
            return pipe
        try:
            import torch
            from diffusers import Flux2KleinPipeline
            try:
                from pillow_heif import register_heif_opener
                register_heif_opener()  # so HEIC/HEIF phone photos open
            except Exception:  # noqa: BLE001
                pass
            # The worker must NOT download (it can't write the cache). Require the
            # main-process background download to have finished, then load READ-ONLY
            # from the local cache snapshot.
            if not _img2img_state.get("snapshot_ready"):
                _img2img_start_bg_download()  # ensure it's running
                raise RuntimeError("model still downloading in the background; try again shortly")
            model_src = _img2img_state.get("snapshot_path") or IMG2IMG_MODEL_ID
            lora_path = _img2img_state.get("lora_path")
            device = "cuda" if torch.cuda.is_available() else "cpu"
            print(f"[img2img] building from cache {model_src} on {device} "
                  f"({'with LoRA' if lora_path else 'base, no LoRA'}) …")
            # low_cpu_mem_usage avoids staging the full ~14 GB in CPU RAM (which
            # can OOM-kill the ZeroGPU worker mid-load). local_files_only forbids
            # any network write from the worker.
            pipe = Flux2KleinPipeline.from_pretrained(
                model_src, torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True, local_files_only=True
            ).to(device)
            if lora_path and os.path.isfile(lora_path):
                # LoRA load is isolated: if it fails (e.g. peft missing, or an
                # incompatible adapter) we keep the BASE pipeline working rather
                # than failing the whole develop.
                try:
                    from pathlib import Path as _Path
                    lf = _Path(lora_path)
                    pipe.load_lora_weights(str(lf.parent), weight_name=lf.name,
                                           adapter_name=IMG2IMG_TRIGGER)
                    pipe.set_adapters([IMG2IMG_TRIGGER], adapter_weights=[IMG2IMG_LORA_SCALE])
                    _img2img_state["lora_loaded"] = True
                    print("[img2img] pipeline ready (watercolour LoRA applied)")
                except Exception as lexc:  # noqa: BLE001
                    _img2img_state["lora_loaded"] = False
                    print(f"[img2img] LoRA load failed ({lexc}); using base model")
            else:
                _img2img_state["lora_loaded"] = False
                print("[img2img] pipeline ready (base model — stage 1, no LoRA)")
            _img2img_state["pipe"] = pipe
            return pipe
        except Exception as exc:  # noqa: BLE001
            import traceback
            _img2img_state["build_error"] = f"{type(exc).__name__}: {exc}"
            print(f"[img2img] in-worker build failed: {exc}")
            traceback.print_exc()
            return None


def _img2img_square(img, size: int = IMG2IMG_SIZE):
    """EXIF-normalise, flatten to RGB, and centre-crop a PIL image to size×size —
    the shape Klein expects (inlined from backend/hf/3_infer_img2img.py)."""
    from PIL import Image as PILImage, ImageOps
    src = ImageOps.exif_transpose(img).convert("RGB")
    w, h = src.size
    s = min(w, h)
    src = src.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))
    return src.resize((size, size), PILImage.LANCZOS)


@spaces.GPU(duration=180)
def _img2img_infer(src):
    """Develop one PIL photo into a watercolour PIL image, inside a @spaces.GPU
    context. Builds/reuses the pipeline here (never in the main process, never
    pickled across the boundary). `src` and the returned image are PIL (both
    picklable), the only things that cross the GPU boundary. Generous duration
    because the FIRST call also pays the ~14 GB cold load; later calls reuse the
    cached pipe and are fast."""
    import torch
    pipe = _ensure_img2img_pipe()
    if pipe is None:
        raise RuntimeError("img2img pipeline could not be initialized: "
                           + str(_img2img_state.get("build_error") or "unknown"))
    sq = _img2img_square(src, size=IMG2IMG_SIZE)
    # Trigger-token prompt only makes sense once the LoRA is loaded; otherwise
    # use the generic base-model watercolour prompt (stage 1).
    prompt = IMG2IMG_PROMPT if _img2img_state.get("lora_loaded") else IMG2IMG_BASE_PROMPT
    device = getattr(pipe, "device", None)
    generator = torch.Generator(
        str(device) if device is not None else "cuda"
    ).manual_seed(IMG2IMG_SEED)
    result = pipe(
        prompt=prompt,
        image=sq,
        height=IMG2IMG_SIZE,
        width=IMG2IMG_SIZE,
        num_inference_steps=IMG2IMG_STEPS,
        guidance_scale=IMG2IMG_GUIDANCE,
        generator=generator,
    )
    return result.images[0]


def _pil_to_data_url(img, fmt: str = "PNG") -> str:
    import base64
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/{fmt.lower()};base64,{b64}"


@app.post("/api/img2img")
async def img2img_rest(request: Request):
    """Develop one uploaded photo into a watercolour painting.
    Body: {data_url: "data:image/...;base64,...", name?: str}.
    Returns {available, image: <png data url>, model}. When the pipeline can't
    run on this host it returns {available: false, image: null} so the caller can
    keep the original photo rather than treating it as an error."""
    import base64
    if not _img2img_prep():
        return {"available": False, "image": None, "model": IMG2IMG_MODEL_ID}
    try:
        body = await request.json()
    except Exception:  # noqa: BLE001
        raise HTTPException(status_code=400, detail="invalid JSON body")
    data_url = (body or {}).get("data_url", "")
    name = str((body or {}).get("name", "upload"))[:60]
    if not data_url:
        raise HTTPException(status_code=400, detail="missing data_url")
    try:
        b64 = data_url.split(",", 1)[1] if "," in data_url else data_url
        raw = base64.b64decode(b64)
        from PIL import Image
        src = Image.open(io.BytesIO(raw))
        with traced("img2img", model=IMG2IMG_MODEL_ID, file=name) as t:
            out = _img2img_infer(src)
            t.set_output({"size": list(out.size)})
        return {"available": True, "image": _pil_to_data_url(out), "model": IMG2IMG_MODEL_ID}
    except HTTPException:
        raise
    except Exception as exc:  # noqa: BLE001
        print(f"[img2img] generation failed: {exc}")
        raise HTTPException(status_code=500, detail=str(exc))


# --------------------------------------------------------------------------- #
# Voice → text via NVIDIA NeMo ASR (item 5). Streaming Nemotron model; needs a
# GPU to be practical. Loaded lazily; if NeMo/torch/GPU are missing the endpoint
# returns {available: false} so the UI can show a graceful message.
# --------------------------------------------------------------------------- #
ASR_MODEL_ID = os.environ.get("SOZAI_ASR_MODEL", "nvidia/nemotron-3.5-asr-streaming-0.6b")
ASR_SAMPLE_RATE = 16_000
_asr_lock = threading.Lock()
_asr_state = {"loaded": False, "model": None, "unsupported_os": False}

# --- ported from the reference Nemotron Spaces app.py (working logic) ------- #
# Cache-aware streaming "lookahead" contexts. Larger = more accurate, higher
# latency. Keys are shown in the UI; values are [left, right] att_context_size.
ASR_CHUNK_PROFILES = {
    "Lowest latency - 80 ms": [56, 0],
    "Fast - 160 ms": [56, 1],
    "Balanced - 320 ms": [56, 3],
    "Accurate - 560 ms": [56, 6],
    "Most accurate - 1.12 s": [56, 13],
}
ASR_DEFAULT_PROFILE = "Most accurate - 1.12 s"   # default per spec

# Ordered (label, value) choices for the UI language picker.
ASR_LANGUAGE_CHOICES = [
    ["Auto-detect", "auto"], ["Arabic", "ar-AR"], ["Bulgarian", "bg-BG"],
    ["Croatian", "hr-HR"], ["Czech", "cs-CZ"], ["Danish", "da-DK"],
    ["Dutch", "nl-NL"], ["English (UK)", "en-GB"], ["English (US)", "en-US"],
    ["Estonian", "et-EE"], ["Finnish", "fi-FI"], ["French (Canada)", "fr-CA"],
    ["French (France)", "fr-FR"], ["German", "de-DE"], ["Greek", "el-GR"],
    ["Hebrew", "he-IL"], ["Hindi", "hi-IN"], ["Hungarian", "hu-HU"],
    ["Italian", "it-IT"], ["Japanese", "ja-JP"], ["Korean", "ko-KR"],
    ["Latvian", "lv-LV"], ["Lithuanian", "lt-LT"], ["Maltese", "mt-MT"],
    ["Mandarin", "zh-CN"], ["Norwegian Bokmal", "nb-NO"], ["Norwegian Nynorsk", "nn-NO"],
    ["Polish", "pl-PL"], ["Portuguese (Brazil)", "pt-BR"], ["Portuguese (Portugal)", "pt-PT"],
    ["Romanian", "ro-RO"], ["Russian", "ru-RU"], ["Slovak", "sk-SK"],
    ["Slovenian", "sl-SI"], ["Spanish (Spain)", "es-ES"], ["Spanish (US)", "es-US"],
    ["Swedish", "sv-SE"], ["Thai", "th-TH"], ["Turkish", "tr-TR"],
    ["Ukrainian", "uk-UA"], ["Vietnamese", "vi-VN"],
]

# 40 supported language-locales (value sent to set_inference_prompt). "auto"
# lets the model auto-detect.
ASR_LANGUAGE_CODES = {
    "auto", "ar-AR", "bg-BG", "hr-HR", "cs-CZ", "da-DK", "nl-NL", "en-GB",
    "en-US", "et-EE", "fi-FI", "fr-CA", "fr-FR", "de-DE", "el-GR", "he-IL",
    "hi-IN", "hu-HU", "it-IT", "ja-JP", "ko-KR", "lv-LV", "lt-LT", "mt-MT",
    "zh-CN", "nb-NO", "nn-NO", "pl-PL", "pt-BR", "pt-PT", "ro-RO", "ru-RU",
    "sk-SK", "sl-SI", "es-ES", "es-US", "sv-SE", "th-TH", "tr-TR", "uk-UA",
    "vi-VN",
}


def _asr_os_supported() -> bool:
    """NeMo / Nemotron streaming ASR is only supported on Linux. On macOS and
    Windows the dependency stack (and ZeroGPU pathways) don't work, so we report
    'not supported' rather than attempting a load that will error out."""
    import platform
    return platform.system().lower() == "linux"


def _get_asr_model():
    if _asr_state["loaded"]:
        return _asr_state["model"]
    with _asr_lock:
        if _asr_state["loaded"]:
            return _asr_state["model"]
        try:
            if not _asr_os_supported():
                import platform
                print(f"[asr] Nemotron is not supported on {platform.system()} "
                      f"(Linux only); ASR disabled.")
                _asr_state["model"] = None
                _asr_state["unsupported_os"] = True
            else:
                import nemo.collections.asr as nemo_asr
                import torch
                torch.set_grad_enabled(False)
                torch.set_float32_matmul_precision("high")
                # Load on CPU at process start. On ZeroGPU the GPU isn't attached
                # yet, so the model must be created on CPU here and only moved to
                # cuda INSIDE the @spaces.GPU call (see _asr_run). map_location +
                # explicit .to("cpu") keep all weight init on CPU.
                model = nemo_asr.models.ASRModel.from_pretrained(
                    model_name=ASR_MODEL_ID,
                    map_location=torch.device("cpu"),
                )
                # Configure greedy RNNT/CTC decoding for cache-aware streaming,
                # mirroring the reference Spaces app so conformer_stream_step works.
                try:
                    from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
                    from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
                    if hasattr(model, "change_decoding_strategy") and hasattr(model, "decoding"):
                        if hasattr(model, "joint"):
                            decoding_cfg = RNNTDecodingConfig(fused_batch_size=-1)
                            decoding_cfg.greedy.use_cuda_graph_decoder = False
                            model.change_decoding_strategy(decoding_cfg)
                        else:
                            model.change_decoding_strategy(CTCDecodingConfig())
                except Exception as dexc:  # noqa: BLE001
                    print(f"[asr] decoding strategy setup skipped: {dexc}")
                model = model.to(device="cpu", dtype=torch.float32)
                model.eval()
                _asr_state["model"] = model
        except Exception as exc:  # noqa: BLE001
            print(f"[asr] NeMo model unavailable: {exc}")
            _asr_state["model"] = None
        finally:
            _asr_state["loaded"] = True
    return _asr_state["model"]


def _prep_wav_16k_mono(src_path: str) -> str:
    """Decode any uploaded audio to 16kHz mono WAV (what the ASR model expects)."""
    import tempfile
    out = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    out.close()
    # Prefer soundfile+librosa; fall back to ffmpeg if present.
    try:
        import soundfile as sf
        import numpy as np  # noqa: F401
        try:
            import librosa
            audio, _sr = librosa.load(src_path, sr=16000, mono=True)
        except Exception:
            data, sr = sf.read(src_path)
            if getattr(data, "ndim", 1) > 1:
                data = data.mean(axis=1)
            audio = data
            if sr != 16000:
                import numpy as np
                # naive resample as a last resort
                idx = (np.arange(int(len(audio) * 16000 / sr)) * sr / 16000).astype(int)
                idx = idx[idx < len(audio)]
                audio = audio[idx]
        sf.write(out.name, audio, 16000, subtype="PCM_16")
        return out.name
    except Exception as exc:  # noqa: BLE001
        # ffmpeg fallback
        try:
            import subprocess
            subprocess.run(["ffmpeg", "-y", "-i", src_path, "-ar", "16000", "-ac", "1", out.name],
                           check=True, capture_output=True)
            return out.name
        except Exception as exc2:  # noqa: BLE001
            print(f"[asr] audio prep failed: {exc} / {exc2}")
            return src_path


def _detect_language(text: str) -> str:
    """Best-effort language id of a transcript (graceful if langdetect absent)."""
    text = (text or "").strip()
    if not text:
        return ""
    try:
        from langdetect import detect
        return detect(text)
    except Exception:  # noqa: BLE001
        return ""


def _asr_configure(model, language: str, chunk_profile: str) -> None:
    """Apply the chosen lookahead profile + language prompt to the model.
    Ported from the reference Nemotron Spaces app."""
    att = ASR_CHUNK_PROFILES.get(chunk_profile, ASR_CHUNK_PROFILES[ASR_DEFAULT_PROFILE])
    if hasattr(model.encoder, "set_default_att_context_size"):
        model.encoder.set_default_att_context_size(att_context_size=att)
    lang = language if language in ASR_LANGUAGE_CODES else "auto"
    if hasattr(model, "set_inference_prompt"):
        try:
            model.set_inference_prompt(lang)
            if hasattr(model, "decoding") and hasattr(model.decoding, "set_strip_lang_tags"):
                # auto mode emits a <xx-XX> tag; strip it for a clean transcript
                model.decoding.set_strip_lang_tags(True, lang_tag_pattern=None)
        except Exception as exc:  # noqa: BLE001
            print(f"[asr] inference-prompt setup skipped: {exc}")


def _asr_transcribe_stream(model, wav_path: str, language: str, chunk_profile: str) -> str:
    """Cache-aware streaming transcription, ported from NVIDIA NeMo's
    Apache-licensed streaming example (via the reference Spaces app.py). Returns
    the final transcript text. Falls back to model.transcribe() on any issue."""
    import torch
    from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer
    from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis

    def _extract(hyps):
        hyps = list(hyps)
        if not hyps:
            return [""]
        if isinstance(hyps[0], Hypothesis):
            return [h.text for h in hyps]
        return [str(h) for h in hyps]

    def _drop_extra(step_num, pad_and_drop):
        if step_num == 0 and not pad_and_drop:
            return 0
        return model.encoder.streaming_cfg.drop_extra_pre_encoded

    _asr_configure(model, language, chunk_profile)
    model_device = next(model.parameters()).device

    buf = CacheAwareStreamingAudioBuffer(model=model, online_normalization=False,
                                         pad_and_drop_preencoded=False)
    buf.append_audio_file(wav_path, stream_id=-1)

    cache_last_channel, cache_last_time, cache_last_channel_len = \
        model.encoder.get_initial_cache_state(batch_size=1)
    previous_hypotheses = None
    previous_pred_out = None
    text = ""
    for step_num, (chunk_audio, chunk_lengths) in enumerate(buf, start=1):
        with torch.inference_mode():
            chunk_audio = chunk_audio.to(device=model_device, dtype=torch.float32)
            chunk_lengths = chunk_lengths.to(device=model_device)
            (previous_pred_out, transcribed_texts, cache_last_channel, cache_last_time,
             cache_last_channel_len, previous_hypotheses) = model.conformer_stream_step(
                processed_signal=chunk_audio,
                processed_signal_length=chunk_lengths,
                cache_last_channel=cache_last_channel,
                cache_last_time=cache_last_time,
                cache_last_channel_len=cache_last_channel_len,
                keep_all_outputs=buf.is_buffer_empty(),
                previous_hypotheses=previous_hypotheses,
                previous_pred_out=previous_pred_out,
                drop_extra_pre_encoded=_drop_extra(step_num - 1, False),
                return_transcription=True,
            )
        text = _extract(transcribed_texts)[0]
    return (text or "").strip()


@spaces.GPU(duration=120)
def _asr_run(wav_path: str, language: str = "auto",
             chunk_profile: str = ASR_DEFAULT_PROFILE) -> str:
    """Transcribe one prepared 16k-mono wav. Tries the cache-aware streaming
    path first (honors language + latency/accuracy profile); falls back to the
    plain transcribe() API if streaming isn't available for this checkpoint.

    ZeroGPU specifics, mirroring NVIDIA's reference Nemotron Space:
      * The model is fetched from module state HERE, not passed as an argument —
        @spaces.GPU pickles arguments to the GPU worker, and we don't want to
        ship the NeMo model across that boundary.
      * The model lives on CPU between calls (loaded on CPU at startup). Inside
        this @spaces.GPU context the GPU is attached, so we move it to cuda for
        the call and back to cpu in `finally`, exactly like the reference app.
        This is what keeps CUDA out of the main process.
    """
    import torch
    model = _asr_state.get("model")
    if model is None:
        return ""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
    try:
        model.to(device=device, dtype=torch.float32)
        try:
            return _asr_transcribe_stream(model, wav_path, language, chunk_profile)
        except Exception as exc:  # noqa: BLE001
            print(f"[asr] streaming path failed ({exc}); falling back to transcribe()")
            try:
                outputs = model.transcribe([wav_path])
                first = outputs[0] if outputs else None
                return (getattr(first, "text", None) or (first if isinstance(first, str) else "") or "").strip()
            except Exception as exc2:  # noqa: BLE001
                print(f"[asr] transcribe() fallback also failed: {exc2}")
                return ""
    finally:
        # release the GPU: move weights back to CPU and clear the cache so the
        # next ZeroGPU request starts clean.
        try:
            if device != "cpu":
                model.to(device="cpu", dtype=torch.float32)
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
        except Exception:  # noqa: BLE001
            pass


@app.api(name="transcribe")
def transcribe(audio_path: str = "", field: str = "") -> dict:
    """Transcribe a short audio clip to text with NeMo ASR and detect its
    language. Returns {text, language, available, field}. If the model can't be
    loaded (no NeMo/GPU), returns available=false with empty text."""
    model = _get_asr_model()
    if model is None:
        _record_trace({"name": "transcribe", "model": ASR_MODEL_ID,
                       "file": os.path.basename(audio_path or ""), "status": "unavailable",
                       "error": "NeMo ASR model not loaded (needs nemo_toolkit + GPU)",
                       "meta": {"field": field}, "output": None, "duration_ms": 0})
        return {"text": "", "language": "", "available": False, "field": field}

    wav = _prep_wav_16k_mono(audio_path)
    text = ""
    with traced("transcribe", model=ASR_MODEL_ID,
                file=os.path.basename(audio_path or ""), meta={"field": field}) as t:
        try:
            # Go through _asr_run so the actual inference runs inside a
            # @spaces.GPU context on ZeroGPU (and a no-op context elsewhere).
            text = _asr_run(wav).strip()
        except Exception as exc:  # noqa: BLE001
            t.error = f"transcribe failed: {exc}"
        t.set_output(text)
    lang = _detect_language(text)
    return {"text": text, "language": lang, "available": True, "field": field}


@app.post("/api/transcribe")
async def transcribe_rest(request: Request):
    """Transcribe recorder audio sent as a base64 data URL.
    Body: {data_url: "data:audio/webm;base64,...", field: "caption"}.
    Returns {text, language, available, field}."""
    import base64
    import tempfile
    try:
        body = await request.json()
    except Exception:  # noqa: BLE001
        return {"text": "", "language": "", "available": False, "field": ""}
    data_url = (body or {}).get("data_url", "")
    field = (body or {}).get("field", "")
    language = (body or {}).get("language", "auto") or "auto"
    chunk_profile = (body or {}).get("chunk_profile", ASR_DEFAULT_PROFILE) or ASR_DEFAULT_PROFILE
    model = _get_asr_model()
    unsupported = bool(_asr_state.get("unsupported_os"))
    if model is None or not data_url:
        if not data_url:
            return {"text": "", "language": "", "available": model is not None,
                    "unsupported_os": unsupported, "field": field}
        _record_trace({"name": "transcribe", "model": ASR_MODEL_ID, "file": "recording",
                       "status": "unavailable",
                       "error": ("Nemotron not supported on this OS (Linux only)" if unsupported
                                 else "NeMo ASR model not loaded (needs nemo_toolkit + GPU)"),
                       "meta": {"field": field}, "output": None, "duration_ms": 0})
        return {"text": "", "language": "", "available": False,
                "unsupported_os": unsupported, "field": field}
    try:
        header, b64 = (data_url.split(",", 1) if "," in data_url else ("", data_url))
        ext = ".webm"
        if "ogg" in header:
            ext = ".ogg"
        elif "wav" in header:
            ext = ".wav"
        elif "mp4" in header or "m4a" in header:
            ext = ".mp4"
        raw = base64.b64decode(b64)
        src = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
        src.write(raw)
        src.close()
        wav = _prep_wav_16k_mono(src.name)
        text = ""
        with traced("transcribe", model=ASR_MODEL_ID, file="recording",
                    meta={"field": field, "language": language, "chunk_profile": chunk_profile}) as t:
            try:
                text = _asr_run(wav, language=language, chunk_profile=chunk_profile)
            except Exception as exc:  # noqa: BLE001
                t.error = f"transcribe failed: {exc}"
            t.set_output(text)
        return {"text": text, "language": _detect_language(text), "available": True,
                "unsupported_os": False, "field": field}
    except Exception as exc:  # noqa: BLE001
        print(f"[asr] rest transcribe failed: {exc}")
        return {"text": "", "language": "", "available": True, "field": field, "error": str(exc)}


# --------------------------------------------------------------------------- #
# Auto-caption inference. MiniCPM-V is a vision model, so the heavy generation
# is isolated in a @spaces.GPU-decorated helper: on ZeroGPU this requests a GPU
# for the duration of the call and releases it after; off ZeroGPU it's a no-op.
# Both backends (gguf via llama.cpp, transformers via .chat) are image-grounded.
# --------------------------------------------------------------------------- #

@spaces.GPU(duration=60)
def _caption_infer(local_img, base_instruction, full_prompt, reply_suffix, max_new):
    """Run one vision generation and return raw model text.

    `full_prompt` already bundles the per-kind instruction + reply directive, so
    the same helper produces a title, a caption, or tags depending on what the
    caller passed in. The IMAGE is sent to the model in both backends.

    IMPORTANT (ZeroGPU): two rules shape this function.
      1. @spaces.GPU ships ARGUMENTS to the GPU worker via pickle, and the
         llama.cpp `Llama` object is NOT picklable ("cannot pickle 'module'
         object"). So the model is never an argument — only plain str/int cross
         the boundary; the model is obtained here.
      2. CUDA must NOT be initialized in the main process on ZeroGPU (GPUs only
         exist inside this decorated call). So for the GGUF backend the model is
         BUILT HERE on first use, while the GPU is attached, and cached in the
         worker — never constructed in the main process.
    """
    backend = _caption_state.get("backend") or "transformers"
    if backend == "gguf":
        # Build (once) and reuse the llama.cpp model inside the GPU worker so the
        # CUDA context is created with a GPU attached. _ensure_gguf_model()
        # caches it in module state; on ZeroGPU that state lives in the worker.
        model = _ensure_gguf_model()
        if model is None:
            raise RuntimeError("GGUF caption model could not be initialized")
        # llama.cpp vision chat completion: send the real image as a data: URI
        # plus the instruction text; the mmproj-backed handler makes visual
        # tokens.
        user_content = [{"type": "text", "text": full_prompt}]
        data_uri = _image_data_uri(local_img)
        if data_uri:
            user_content.insert(0, {"type": "image_url", "image_url": {"url": data_uri}})
        messages = [
            {"role": "system", "content": base_instruction},
            {"role": "user", "content": user_content},
        ]
        out = model.create_chat_completion(
            messages=messages, max_tokens=max_new, temperature=0.3,
        )
        return (out["choices"][0]["message"]["content"] or "").strip()

    # transformers backend: MiniCPM-V's .chat() helper takes the PIL image and a
    # msgs list. This keeps the off-Spaces path image-grounded too.
    model = _caption_state["model"]
    tokenizer = _caption_state["tokenizer"]
    pil = _load_pil_image(local_img)
    msgs = [{"role": "user", "content": [pil, full_prompt] if pil is not None else [full_prompt]}]
    answer = model.chat(
        msgs=msgs,
        image=pil,
        tokenizer=tokenizer,
        sampling=True,
        temperature=0.3,
        max_new_tokens=max_new,
    )
    if isinstance(answer, (list, tuple)):
        answer = answer[0] if answer else ""
    return (str(answer) or "").strip()


@app.api(name="autocaption")
def autocaption(image_path: str = "", alt: str = "", prompt: str = "", kind: str = "caption") -> str:
    """Generate a caption, title, or tag suggestion for a photo by LOOKING AT
    the photo with MiniCPM-V.

    `kind` is "caption", "title", or "tags" — it selects the instruction, the
    reply directive, the generation length, and the trace span name, so the one
    function serves all three call sites (caption field, title field, tags
    field) with image-grounded output.

    `prompt` is an optional user-supplied instruction (from Settings) that lets
    people customize the style. When empty, a sensible default is used.

    On any failure (model missing, generation error) it returns a lightweight
    fallback derived from the photo's alt text so the UI button always works.
    """
    kind_l = str(kind).lower()
    is_title = kind_l == "title"
    is_tags = kind_l == "tags"
    span_name = "autotags" if is_tags else ("autotitle" if is_title else "autocaption")
    default_instruction = (
        "Look at this personal photo and list 3 to 6 short, lowercase tags "
        "describing what is in it, separated by commas."
        if is_tags else
        "Look at this personal photo and write a very short, evocative title "
        "(2–5 words) for it."
        if is_title else
        "Look at this personal photo and write a short, warm one-sentence caption "
        "for a personal photo album."
    )
    reply_suffix = (
        " Reply with only the comma-separated tags." if is_tags else
        " Reply with only the title." if is_title else
        " Reply with only the caption."
    )
    max_new = 32 if is_tags else 24 if is_title else 128

    if not _load_caption_model():
        if is_tags:
            cap = _fallback_tags(alt)
        else:
            cap = _fallback_caption(alt)
        _record_trace({"name": span_name, "kind": "LLM", "model": AUTOCAPTION_MODEL_ID,
                       "file": os.path.basename(image_path or ""), "status": "fallback",
                       "error": "caption model unavailable", "output": cap[:300],
                       "meta": {"prompt": (prompt or "")[:120], "for": kind}, "duration_ms": 0})
        return cap

    with traced(span_name, model=AUTOCAPTION_MODEL_ID,
                file=os.path.basename(image_path or ""),
                meta={"prompt": (prompt or "")[:120], "alt": (alt or "")[:120], "for": kind},
                invocation={"max_new_tokens": max_new, "temperature": 0.3},
                system=default_instruction) as t:
        try:
            local_img = _resolve_image_for_model(image_path)
            # Some frontends pass the photo as a data URL in `alt`; if image_path
            # didn't resolve to a real file, try alt as an image source too.
            if local_img is None and alt and alt.strip().startswith("data:"):
                local_img = _resolve_image_for_model(alt)
            # If we DID get an image, don't pollute the prompt with the filename;
            # the pixels are the context. Only fall back to an alt text hint when
            # no image could be resolved.
            hint = "" if local_img else (f' (context hint: {alt})' if alt else "")
            base_instruction = (prompt or "").strip() or default_instruction
            full_prompt = base_instruction + hint + reply_suffix

            text = _caption_infer(local_img, base_instruction,
                                  full_prompt, reply_suffix, max_new)

            if is_tags:
                # tags must come out comma-separated; normalize whatever the
                # model produced, then validate. If it can't be made into a
                # comma list, fall back (which is always comma-separated).
                norm = _normalize_tags(text)
                if not _is_comma_separated(norm):
                    norm = _fallback_tags(alt)
                result = norm or _fallback_tags(alt)
            else:
                # tidy: first line only, strip surrounding quotes
                text = (text.splitlines()[0].strip().strip('"').strip()) if text else ""
                result = text or _fallback_caption(alt)
            t.set_output(result)
            return result
        except Exception as exc:
            print(f"[{span_name}] generation failed, using fallback: {exc}")
            t.error = f"generation failed: {exc}"
            fb = _fallback_tags(alt) if is_tags else _fallback_caption(alt)
            t.set_output(fb)
            return fb


# ----------------------------- Room REST API ------------------------------- #

_JOIN_HITS: dict = {}


def rate_ok(ip: str, limit=20, window=60) -> bool:
    now = time.time()
    hits = _JOIN_HITS.setdefault(ip, [])
    hits[:] = [t for t in hits if now - t < window]
    if len(hits) >= limit:
        return False
    hits.append(now)
    return True


@app.post("/api/rooms")
async def create_room(request: Request):
    """Create a room. Body: { code_length?, display_name?, session_id? }."""
    try:
        data = await request.json()
    except Exception:
        data = {}
    code_length = clamp_code_len(data.get("code_length", 6))
    name = sanitize_name(data.get("display_name"))
    session_id = (data.get("session_id") or secrets.token_hex(8))[:64]

    room_id = str(uuid.uuid4())          # secure, never sequential
    share_code = unique_share_code(code_length)
    now = time.time()
    with DB_LOCK:
        _conn.execute(
            "INSERT INTO rooms (room_id, share_code, code_length, created_at, "
            "last_activity, owner_session_id, is_active, metadata) "
            "VALUES (?,?,?,?,?,?,1,?)",
            (room_id, share_code, code_length, now, now, session_id, json.dumps({})),
        )
        _conn.commit()
    return {
        "room_id": room_id,
        "share_code": share_code,
        "code_length": code_length,
        "room_path": f"/room/{room_id}",
    }


@app.get("/api/rooms/{room_id}")
def get_room(room_id: str):
    row = room_row(room_id)
    if row is None:
        raise HTTPException(status_code=404, detail="Room not found or expired.")
    return {
        "room_id": row["room_id"],
        "share_code": row["share_code"],
        "code_length": row["code_length"],
        "is_active": bool(row["is_active"]),
        "participants": manager.presence(room_id),
    }


@app.post("/api/join-by-code")
async def join_by_code(request: Request):
    """Resolve a human room code to a room_id. Body: { code }."""
    ip = request.client.host if request.client else "?"
    if not rate_ok(ip):
        raise HTTPException(status_code=429, detail="Too many attempts. Slow down.")
    try:
        data = await request.json()
    except Exception:
        data = {}
    code = (data.get("code") or "").strip().upper()
    code = "".join(ch for ch in code if ch in CODE_ALPHABET)
    if not (MIN_CODE_LEN <= len(code) <= MAX_CODE_LEN):
        raise HTTPException(status_code=400, detail="Invalid code format.")
    with DB_LOCK:
        row = _conn.execute(
            "SELECT room_id FROM rooms WHERE share_code=? AND is_active=1", (code,)
        ).fetchone()
    if row is None:
        raise HTTPException(status_code=404, detail="No active room with that code.")
    return {"room_id": row["room_id"], "room_path": f"/room/{row['room_id']}"}


# --------------------------- Room content persistence ---------------------- #
# A room's scrapbook (photos, captions, pins, transforms) is stored as one JSON
# blob with a monotonic version. The frontend PUTs on change (debounced) and
# GETs on join, so a returning user sees the room as it was left. Live edits
# still flow over the WebSocket; this is the durable backstop + late-join seed.

# cap the stored blob so a room can't be used as unbounded storage (data URLs
# for images can be large; 8 MB is generous for a handful of photos)
MAX_ROOM_STATE_BYTES = int(os.environ.get("SOZAI_MAX_ROOM_STATE_BYTES", str(8 * 1024 * 1024)))


@app.get("/api/rooms/{room_id}/state")
def get_room_state(room_id: str):
    """Return the persisted scrapbook state for a room: {state, version}.
    state is null if nothing has been saved yet."""
    if room_row(room_id) is None:
        raise HTTPException(status_code=404, detail="Room not found or expired.")
    with DB_LOCK:
        row = _conn.execute(
            "SELECT state, version, updated_at FROM room_state WHERE room_id=?", (room_id,)
        ).fetchone()
    if row is None:
        return {"state": None, "version": 0, "updated_at": None}
    try:
        state = json.loads(row["state"]) if row["state"] else None
    except Exception:  # noqa: BLE001
        state = None
    return {"state": state, "version": row["version"], "updated_at": row["updated_at"]}


@app.post("/api/rooms/{room_id}/state")
async def put_room_state(room_id: str, request: Request):
    """Persist the room's scrapbook state. Body: {state, base_version?}.

    Uses optimistic concurrency: if base_version is supplied and no longer
    matches the stored version, the write is rejected with the current state so
    the client can reconcile (last-writer is NOT blindly allowed to clobber a
    newer save). Returns {version, conflict?, state?}.
    """
    if room_row(room_id) is None:
        raise HTTPException(status_code=404, detail="Room not found or expired.")
    try:
        body = await request.json()
    except Exception:  # noqa: BLE001
        raise HTTPException(status_code=400, detail="Invalid JSON.")
    state = (body or {}).get("state")
    base_version = (body or {}).get("base_version")

    blob = json.dumps(state, separators=(",", ":"))
    if len(blob.encode("utf-8")) > MAX_ROOM_STATE_BYTES:
        raise HTTPException(status_code=413, detail="Scrapbook state too large to save.")

    now = time.time()
    with DB_LOCK:
        cur = _conn.execute(
            "SELECT state, version FROM room_state WHERE room_id=?", (room_id,)
        ).fetchone()
        current_version = cur["version"] if cur else 0
        # optimistic concurrency check
        if base_version is not None and cur is not None and int(base_version) != int(current_version):
            try:
                current_state = json.loads(cur["state"]) if cur["state"] else None
            except Exception:  # noqa: BLE001
                current_state = None
            return {"version": current_version, "conflict": True, "state": current_state}
        new_version = current_version + 1
        _conn.execute(
            "INSERT INTO room_state (room_id, state, version, updated_at) VALUES (?,?,?,?) "
            "ON CONFLICT(room_id) DO UPDATE SET state=excluded.state, "
            "version=excluded.version, updated_at=excluded.updated_at",
            (room_id, blob, new_version, now),
        )
        _conn.commit()
    touch_room(room_id)
    return {"version": new_version, "conflict": False}


# --------------------------- Realtime sync (SSE) --------------------------- #

_PALETTE = [
    "#7c5cff", "#e9698f", "#2bb673", "#f2994a", "#3aa0ff",
    "#9b51e0", "#eb5757", "#27ae60", "#f2c94c", "#56ccf2",
]


def color_for(session_id: str) -> str:
    h = int(hashlib.sha1(session_id.encode()).hexdigest(), 16)
    return _PALETTE[h % len(_PALETTE)]


# Events the server relays verbatim to the other participants.
#
# Every collaborative interaction the frontend emits must be listed here, or the
# server silently drops it and collaborators never see the action. Rather than a
# brittle allow-list that has to be hand-updated whenever the frontend grows a
# new event, we relay ANY event that is NOT a server-owned/internal type
# (see NON_RELAY_EVENTS below). This keeps cursors, selections, field typing,
# polaroid drags, card edits, screen-follow, etc. all in sync.
RELAY_EVENTS = {
    # object lifecycle + text field edits (title / caption / date / time)
    "object_created", "object_updated", "object_deleted",
    # chat + presence-style ephemera
    "chat_message", "cursor_position", "selection", "interaction",
    # the photo the editor is currently viewing (drives optional follow)
    "active_image",
    # darkroom / scrapbook card edits (time / location / note)
    "dv_edit",
    # polaroid pins on the map: drop, move, style (rotate/scale), reset
    "photo_pinned", "photo_moved", "photo_props", "photos_reset",
    # navigation that collaborators can optionally follow
    "screen_changed", "variant_changed",
    # LIVE shared image library: the full photo list, relayed on every
    # add / delete / edit so both people always see the same images instantly.
    "photos_state",
    # ambient "what is my collaborator doing / where are they" presence beacon
    # (navigation, AI activity, etc.) — shown as a floating label by the cursor.
    "activity",
    # SYNCED "Play your journey" playback. Only the lock-holding driver emits
    # these; the peer applies them so both watch the SAME flythrough. The lock
    # itself (journey_acquire/release/state) is server-owned (NON_RELAY) below.
    "journey_frame", "journey_transport", "journey_start", "journey_stop",
    # darkroom develop: which prints have been revealed (shake) so the peer's
    # darkroom shows the same reveal progress and developed watercolours.
    "darkroom_reveal", "darkroom_phase", "darkroom_developed",
    # map fullscreen toggle (optional follow, same as screen/variant)
    "map_fullscreen",
}

# Types the server originates or manages itself — these must NEVER be relayed
# back out from a client (a client cannot, e.g., forge presence for others).
NON_RELAY_EVENTS = {
    "init", "me", "presence", "presence_update", "status",
    "user_joined", "user_left", "rename", "error",
    "join_request", "join_pending", "join_declined", "admit_join", "decline_join", "cancel_join",
    # the AI mutual-exclusion lock is server-authoritative: clients REQUEST it,
    # the server grants/denies and broadcasts the resulting ai_state. Clients
    # may not forge ai_state directly, so it's non-relay.
    "ai_acquire", "ai_release", "ai_state",
    # the journey mutual-exclusion lock is likewise server-authoritative.
    "journey_acquire", "journey_release", "journey_refresh", "journey_state",
    # room-full notice is server-originated
    "room_full",
}

# A room holds at most this many simultaneous participants. The owner plus one
# guest = 2. A third connection is rejected with a room_full notice.
ROOM_CAPACITY = int(os.environ.get("SOZAI_ROOM_CAPACITY", "2"))


# NOTE: This transport was a raw WebSocket (/ws/{room_id}). Hugging Face
# Spaces' router silently drops custom (non-Gradio) WebSocket UPGRADE requests
# at the proxy — the handshake never reaches this process, so every connection
# loops in "reconnecting". Plain HTTP routes on this same app work fine, so we
# mirror what Gradio itself does on Spaces: a one-way Server-Sent Events stream
# for server→client messages, plus ordinary POST for client→server messages.
#
# Each client is identified by its session_id and holds one open SSE GET. The
# Manager keys connections by an asyncio.Queue per client instead of a socket;
# "sending" to a client just puts a message on its queue, which the SSE
# generator drains to the browser. Presence, join-approval, the relay
# allow-list, and the DB writes are all unchanged from the WebSocket version.
#
# SINGLE-REPLICA ASSUMPTION: a POST and the matching SSE stream must land on the
# same process. Your Space runs one replica today, so this holds. If you ever
# scale to multiple replicas, move the queues behind Redis pub/sub so a POST on
# replica A reaches a stream on replica B.


# The four fields that are safe to broadcast to other participants. `info`
# never contains anything else now, but we filter through this everywhere it
# crosses the wire so a stray non-JSON value (e.g. an asyncio.Event) can never
# crash the SSE stream's json.dumps.
_PUBLIC_FIELDS = ("session_id", "name", "participant_id", "color")


def _public(info):
    """Project a participant's info down to only the JSON-safe public fields."""
    return {k: info[k] for k in _PUBLIC_FIELDS}


class _Client:
    """One connected participant: their info + the queue feeding their SSE."""
    def __init__(self, info):
        self.info = info
        self.queue: asyncio.Queue = asyncio.Queue()
        self.alive = True

    async def put(self, message):
        if self.alive:
            await self.queue.put(message)


class _PendingJoin:
    """A joiner awaiting host approval. Holds the decision Event SEPARATELY from
    `info` so `info` stays pure-JSON and can be broadcast safely."""
    def __init__(self, info):
        self.info = info
        self.event = asyncio.Event()
        self.admitted = False


class Manager:
    def __init__(self):
        self.rooms = {}     # room_id -> {session_id: _Client}  (admitted)
        self.pending = {}   # room_id -> {session_id: info}     (awaiting host)
        # AI mutual-exclusion lock per room. room_id -> {"session_id","name","action","ts"}
        # or None when free. Only one participant may run AI at a time.
        self.ai_lock = {}
        # Journey mutual-exclusion lock per room. The "Play your journey"
        # flythrough (and the watercolour "develop" pass) are GPU/memory heavy;
        # letting BOTH people trigger them at once risks an out-of-memory crash.
        # So only one participant may drive a journey at a time. The driver's
        # control frames are relayed to the peer so they watch the SAME playback
        # (see the journey_* relay events) — the peer's button stays disabled
        # until the driver finishes or stops. room_id -> {session_id,name,ts}.
        self.journey_lock = {}

    # ---- capacity --------------------------------------------------------- #
    def participant_count(self, room_id):
        """Number of DISTINCT admitted sessions (a person who reconnects with the
        same session_id counts once, so a refresh doesn't burn a seat)."""
        return len({sid for sid in self.rooms.get(room_id, {})})

    def has_seat(self, room_id, session_id):
        """True if this session can occupy/continue a seat: either it's already
        in the room (reconnect) or the room is below capacity."""
        peers = self.rooms.get(room_id, {})
        if session_id in peers:
            return True
        return len(peers) < ROOM_CAPACITY

    # ---- AI lock ---------------------------------------------------------- #
    AI_LOCK_TIMEOUT = 90  # seconds; a stale lock (crashed holder) auto-expires

    def ai_state(self, room_id):
        """Current AI lock for a room as a JSON-safe dict (or {'holder': None})."""
        lock = self.ai_lock.get(room_id)
        if lock and (time.time() - lock["ts"]) > self.AI_LOCK_TIMEOUT:
            # stale — the holder went away without releasing
            self.ai_lock.pop(room_id, None)
            lock = None
        if not lock:
            return {"holder": None}
        return {"holder": {k: lock[k] for k in ("session_id", "name", "action")}}

    def ai_acquire(self, room_id, session_id, name, action):
        """Try to take the AI lock. Returns True if acquired (or already held by
        this same session — re-acquire/refresh), False if held by someone else."""
        lock = self.ai_lock.get(room_id)
        if lock and (time.time() - lock["ts"]) > self.AI_LOCK_TIMEOUT:
            lock = None
        if lock and lock["session_id"] != session_id:
            return False
        self.ai_lock[room_id] = {
            "session_id": session_id, "name": name,
            "action": action, "ts": time.time(),
        }
        return True

    def ai_release(self, room_id, session_id):
        """Release the lock if this session holds it."""
        lock = self.ai_lock.get(room_id)
        if lock and lock["session_id"] == session_id:
            self.ai_lock.pop(room_id, None)
            return True
        return False

    def ai_release_all(self, room_id, session_id):
        """Drop the lock unconditionally if held by this session (used on
        disconnect so a vanished holder never blocks the room)."""
        return self.ai_release(room_id, session_id)

    # ---- Journey lock ----------------------------------------------------- #
    # The flythrough can run for a while, so its lock lives longer than the AI
    # one. A crashed driver still auto-expires so the room never gets stuck.
    JOURNEY_LOCK_TIMEOUT = 600  # seconds

    def journey_state(self, room_id):
        """Current journey lock for a room as a JSON-safe dict (or {'driver': None})."""
        lock = self.journey_lock.get(room_id)
        if lock and (time.time() - lock["ts"]) > self.JOURNEY_LOCK_TIMEOUT:
            self.journey_lock.pop(room_id, None)
            lock = None
        if not lock:
            return {"driver": None}
        return {"driver": {k: lock[k] for k in ("session_id", "name")}}

    def journey_acquire(self, room_id, session_id, name):
        """Try to take the journey lock. True if acquired (or re-acquired by the
        same session), False if another participant is already driving."""
        lock = self.journey_lock.get(room_id)
        if lock and (time.time() - lock["ts"]) > self.JOURNEY_LOCK_TIMEOUT:
            lock = None
        if lock and lock["session_id"] != session_id:
            return False
        self.journey_lock[room_id] = {"session_id": session_id, "name": name, "ts": time.time()}
        return True

    def journey_refresh(self, room_id, session_id):
        """Keep a long-running journey lock alive while the driver streams frames."""
        lock = self.journey_lock.get(room_id)
        if lock and lock["session_id"] == session_id:
            lock["ts"] = time.time()
            return True
        return False

    def journey_release(self, room_id, session_id):
        """Release the journey lock if this session holds it."""
        lock = self.journey_lock.get(room_id)
        if lock and lock["session_id"] == session_id:
            self.journey_lock.pop(room_id, None)
            return True
        return False

    # ---- admitted participants -------------------------------------------- #
    def add_client(self, room_id, client):
        self.rooms.setdefault(room_id, {})[client.info["session_id"]] = client

    def get_client(self, room_id, session_id):
        return self.rooms.get(room_id, {}).get(session_id)

    def disconnect(self, room_id, session_id):
        peers = self.rooms.get(room_id)
        if peers and session_id in peers:
            peers[session_id].alive = False
            del peers[session_id]
            if not peers:
                self.rooms.pop(room_id, None)
        self.remove_pending(room_id, session_id)

    # ---- pending (awaiting host approval) --------------------------------- #
    def admit(self, room_id, info):
        """Promote a pending session to a full participant; returns its _Client."""
        self.pending.get(room_id, {}).pop(info["session_id"], None)
        client = _Client(info)
        self.add_client(room_id, client)
        return client

    def add_pending(self, room_id, pending):
        self.pending.setdefault(room_id, {})[pending.info["session_id"]] = pending

    def remove_pending(self, room_id, session_id):
        p = self.pending.get(room_id)
        if p and session_id in p:
            del p[session_id]
            if not p:
                self.pending.pop(room_id, None)

    def find_pending(self, room_id, session_id):
        """Return the _PendingJoin for a session, or None."""
        return self.pending.get(room_id, {}).get(session_id)

    # ---- host resolution (who approves joins) ----------------------------- #
    def host_client(self, room_id, owner_session_id):
        """The client that should approve joins: the owner if connected, else
        the earliest-connected participant (so approval never deadlocks)."""
        peers = self.rooms.get(room_id, {})
        if not peers:
            return None
        if owner_session_id and owner_session_id in peers:
            return peers[owner_session_id]
        return next(iter(peers.values()), None)

    # ---- presence --------------------------------------------------------- #
    def presence(self, room_id):
        peers = self.rooms.get(room_id, {})
        seen, out = set(), []
        for client in peers.values():
            info = client.info
            if info["session_id"] in seen:
                continue
            seen.add(info["session_id"])
            out.append(_public(info))
        return out

    # ---- delivery --------------------------------------------------------- #
    async def broadcast(self, room_id, message, exclude=None):
        for sid, client in list(self.rooms.get(room_id, {}).items()):
            if sid == exclude:
                continue
            try:
                await client.put(message)
            except Exception:
                pass

    async def send_to(self, room_id, session_id, message):
        client = self.get_client(room_id, session_id)
        if client is None:
            return False
        try:
            await client.put(message)
            return True
        except Exception:
            return False


manager = Manager()

# Pending joiners don't have a _Client/queue yet (they're not admitted), but
# they still need to receive join_pending / join_declined over their own SSE
# stream. We give every SSE connection a queue up front, keyed by session_id,
# so the stream exists from the moment the browser connects — even while the
# session is still waiting for host approval.
_sse_queues = {}  # (room_id, session_id) -> asyncio.Queue


def _sse_queue(room_id, session_id):
    """Return a FRESH queue for a new SSE connection and mark it as the current
    one for this (room, session). If the same session reconnects (e.g. a page
    refresh opens a new stream before the old one's cleanup runs), the newest
    connection wins: the old generator detects it is no longer current and bows
    out WITHOUT evicting the live participant. This prevents a reconnect from
    dropping your own seat or AI lock."""
    key = (room_id, session_id)
    q = asyncio.Queue()
    _sse_queues[key] = q
    return q


def _is_current_sse(room_id, session_id, q):
    return _sse_queues.get((room_id, session_id)) is q


async def _relay_incoming(room_id, session_id, data):
    """Apply one client→server event (formerly a WS receive_json). Shared by the
    POST /send endpoint. `data` already has a 'type'."""
    client = manager.get_client(room_id, session_id)
    info = client.info if client else None

    t = data.get("type")

    if t == "rename" and info is not None:
        info["name"] = sanitize_name(data.get("name"))
        with DB_LOCK:
            _conn.execute(
                "UPDATE participants SET display_name=?, last_seen=? WHERE participant_id=?",
                (info["name"], time.time(), info["participant_id"]),
            )
            _conn.commit()
        await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})
        return

    # ---- host approves / declines a pending joiner ----
    if t in ("admit_join", "decline_join"):
        target_sid = data.get("session_id")
        pend = manager.find_pending(room_id, target_sid)
        if pend is not None:
            pend.admitted = (t == "admit_join")
            pend.event.set()
        return

    # ---- AI mutual-exclusion lock ----
    # A client asks to start an AI action; the server grants it only if no one
    # else holds the lock, then broadcasts the authoritative ai_state to BOTH
    # people so the other user's AI controls disable. The requester is told
    # whether it won via a direct ai_state (with a `granted` hint).
    if t == "ai_acquire" and info is not None:
        action = str(data.get("action") or "AI")[:40]
        ok = manager.ai_acquire(room_id, session_id, info["name"], action)
        # tell the requester the outcome (granted true/false)
        await manager.send_to(room_id, session_id, {
            "type": "ai_state", "granted": ok, **manager.ai_state(room_id),
        })
        # tell everyone the new state (so the peer disables their buttons)
        await manager.broadcast(room_id, {"type": "ai_state", **manager.ai_state(room_id)}, exclude=session_id)
        touch_room(room_id)
        return

    if t == "ai_release" and info is not None:
        manager.ai_release(room_id, session_id)
        await manager.broadcast(room_id, {"type": "ai_state", **manager.ai_state(room_id)})
        touch_room(room_id)
        return

    # ---- Journey mutual-exclusion lock ----
    # Like the AI lock, but for the "Play your journey" flythrough / develop.
    # The requester is told whether it won; everyone is told the new driver so
    # the peer's Play button disables. The driver then streams journey_frame
    # events (relayed normally) so the peer watches the same playback.
    if t == "journey_acquire" and info is not None:
        ok = manager.journey_acquire(room_id, session_id, info["name"])
        await manager.send_to(room_id, session_id, {
            "type": "journey_state", "granted": ok, **manager.journey_state(room_id),
        })
        await manager.broadcast(room_id, {"type": "journey_state", **manager.journey_state(room_id)}, exclude=session_id)
        touch_room(room_id)
        return

    if t == "journey_refresh" and info is not None:
        manager.journey_refresh(room_id, session_id)
        return

    if t == "journey_release" and info is not None:
        manager.journey_release(room_id, session_id)
        await manager.broadcast(room_id, {"type": "journey_state", **manager.journey_state(room_id)})
        touch_room(room_id)
        return

    # explicitly known (RELAY_EVENTS) or simply not a server-owned type
    # (NON_RELAY_EVENTS) — this way the frontend can introduce new interaction
    # events without the server silently dropping them.
    if t and t not in NON_RELAY_EVENTS and info is not None:
        data["from"] = _public(info)
        await manager.broadcast(room_id, data, exclude=session_id)
        # cursors + transient interaction pings are too chatty to keep the room
        # alive on; everything else counts as real activity.
        if t not in ("cursor_position", "interaction", "activity", "photos_state"):
            touch_room(room_id)


@app.get("/api/rooms/{room_id}/stream")
async def room_stream(room_id: str, request: Request):
    """Server-Sent Events stream: server→client messages for one participant.

    Replaces the inbound half of the old WebSocket. The browser opens this with
    EventSource(...). Query params: session_id, name."""
    if room_row(room_id) is None:
        # mimic the old 4404: emit a single error event then end the stream.
        async def _gone():
            yield 'data: {"type":"error","message":"Room not found or expired."}\n\n'
        return StreamingResponse(_gone(), media_type="text/event-stream")

    session_id = (request.query_params.get("session_id") or secrets.token_hex(8))[:64]
    info = {
        "session_id": session_id,
        "name": sanitize_name(request.query_params.get("name")),
        "participant_id": secrets.token_hex(6),
        "color": color_for(session_id),
    }

    q = _sse_queue(room_id, session_id)

    # ---- join approval (same policy as the WS version) ----
    row = room_row(room_id)
    owner_session = row["owner_session_id"] if row else None
    is_owner = bool(owner_session) and session_id == owner_session
    host = manager.host_client(room_id, owner_session)
    needs_approval = (not is_owner) and (host is not None)

    # ---- capacity gate: at most ROOM_CAPACITY (2) people per room ----
    # The owner is never locked out of their own room. Anyone else is rejected
    # when the room is already full AND they don't already hold a seat (a plain
    # reconnect with the same session_id keeps its seat). We emit a room_full
    # notice, then close the stream so the client can show "Room is full".
    if not is_owner and not manager.has_seat(room_id, session_id):
        async def _full():
            yield f'data: {json.dumps({"type": "room_full", "capacity": ROOM_CAPACITY}, separators=(",", ":"))}\n\n'
        _sse_queues.pop((room_id, session_id), None)
        return StreamingResponse(_full(), media_type="text/event-stream", headers={"Cache-Control": "no-cache"})

    if needs_approval:
        pend = _PendingJoin(info)
        manager.add_pending(room_id, pend)
        await q.put({"type": "join_pending", "you": _public(info)})
        await host.put({"type": "join_request", "participant": _public(info)})
        try:
            await asyncio.wait_for(pend.event.wait(), timeout=120)
        except Exception:
            pass
        if not pend.admitted:
            manager.remove_pending(room_id, session_id)
            await q.put({"type": "join_declined"})
            await q.put({"__close__": True})
            # fall through to the generator, which will flush join_declined then end
            client = None
        else:
            client = manager.admit(room_id, info)
            client.queue = q  # reuse the queue the SSE stream is already draining
    else:
        # New connection OR a reconnect by the same session. If the session is
        # already seated (reconnect), keep its existing _Client but repoint its
        # queue at this new SSE stream so broadcasts reach the live connection.
        existing = manager.get_client(room_id, session_id)
        if existing is not None:
            existing.queue = q
            existing.alive = True
            client = existing
        else:
            client = _Client(info)
            client.queue = q
            manager.add_client(room_id, client)

    if client is not None:
        now = time.time()
        with DB_LOCK:
            _conn.execute(
                "INSERT OR REPLACE INTO participants (participant_id, room_id, "
                "session_id, display_name, joined_at, last_seen) VALUES (?,?,?,?,?,?)",
                (info["participant_id"], room_id, session_id, info["name"], now, now),
            )
            _conn.commit()
        touch_room(room_id)
        await q.put({"type": "init", "you": _public(info), "participants": manager.presence(room_id)})
        # seed the current AI-lock state so a late joiner immediately knows if
        # the other person is mid-AI-action (and disables their own AI buttons).
        await q.put({"type": "ai_state", **manager.ai_state(room_id)})
        # likewise seed the journey-lock state, so a late joiner's Play button
        # is correctly disabled if the other person is already driving a journey.
        await q.put({"type": "journey_state", **manager.journey_state(room_id)})
        await manager.broadcast(room_id, {"type": "user_joined", "participant": _public(info)}, exclude=session_id)
        await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})

    async def event_gen():
        try:
            while True:
                # heartbeat keeps the connection alive through idle proxies
                try:
                    msg = await asyncio.wait_for(q.get(), timeout=20)
                except asyncio.TimeoutError:
                    yield ": keep-alive\n\n"
                    continue
                if isinstance(msg, dict) and msg.get("__close__"):
                    break
                yield f"data: {json.dumps(msg, separators=(',', ':'))}\n\n"
        except asyncio.CancelledError:
            pass
        finally:
            # Only tear down if THIS connection is still the current one for the
            # session. If the same session reconnected (refresh), a newer stream
            # already took over — this superseded generator must NOT evict the
            # live participant, drop their seat, or free their AI lock.
            superseded = not _is_current_sse(room_id, session_id, q)
            if not superseded:
                _sse_queues.pop((room_id, session_id), None)
            if not superseded and manager.get_client(room_id, session_id) is not None:
                # if this person was holding the AI lock, free it so the room
                # isn't stuck "someone is generating…" forever after they leave.
                released = manager.ai_release_all(room_id, session_id)
                # same for the journey lock: a driver who leaves must not block
                # the peer's Play button forever.
                journey_released = manager.journey_release(room_id, session_id)
                manager.disconnect(room_id, session_id)
                await manager.broadcast(room_id, {"type": "user_left", "participant": _public(info)})
                await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})
                if released:
                    await manager.broadcast(room_id, {"type": "ai_state", **manager.ai_state(room_id)})
                if journey_released:
                    await manager.broadcast(room_id, {"type": "journey_state", **manager.journey_state(room_id)})

    return StreamingResponse(event_gen(), media_type="text/event-stream", headers={
        "Cache-Control": "no-cache",
        "X-Accel-Buffering": "no",   # disable proxy buffering so events flush live
        "Connection": "keep-alive",
    })


@app.post("/api/rooms/{room_id}/send")
async def room_send(room_id: str, request: Request):
    """Client→server messages: one event the client used to send over the
    socket (object_updated, chat_message, cursor_position, admit_join, ...)."""
    if room_row(room_id) is None:
        raise HTTPException(status_code=404, detail="Room not found or expired.")
    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON.")
    session_id = (body or {}).get("session_id")
    if not session_id:
        raise HTTPException(status_code=400, detail="Missing session_id.")
    data = (body or {}).get("event") or {}
    if not isinstance(data, dict) or not data.get("type"):
        raise HTTPException(status_code=400, detail="Missing event.type.")
    await _relay_incoming(room_id, session_id, data)
    return {"ok": True}


# ------------------------------ Assets ------------------------------------- #

def _safe_rel(path: str) -> str:
    rel = os.path.normpath(path).lstrip("/\\")
    if rel.startswith("..") or os.path.isabs(rel):
        raise HTTPException(status_code=400, detail="bad path")
    return rel.replace("\\", "/")


@app.get("/assets/{path:path}")
async def assets(request: Request, path: str):
    rel = _safe_rel(path)
    if PUBLIC_DIR is not None:
        fp = os.path.join(PUBLIC_DIR, *rel.split("/"))
        if os.path.isfile(fp):
            return FileResponse(fp)
    if ZIP_PATH is not None:
        try:
            data = _read_from_zip(rel)
            media = mimetypes.guess_type(rel)[0] or "application/octet-stream"
            return Response(content=data, media_type=media)
        except KeyError:
            pass
    # Not a Sozai asset → it's almost certainly a Phoenix bundle asset (its JS
    # requests /assets/vendor-*.js etc. with an absolute path). Proxy it to the
    # in-process Phoenix server so the embedded trace UI loads correctly.
    if _phoenix_session is not None:
        return await _phoenix_proxy(request, "assets/" + rel)
    raise HTTPException(status_code=404, detail=f"asset not found: {rel}")


@app.get("/favicon.ico")
def favicon():
    rel = "icon.svg"
    if PUBLIC_DIR is not None:
        fp = os.path.join(PUBLIC_DIR, rel)
        if os.path.isfile(fp):
            return FileResponse(fp)
    if ZIP_PATH is not None:
        try:
            return Response(content=_read_from_zip(rel), media_type="image/svg+xml")
        except KeyError:
            pass
    raise HTTPException(status_code=404, detail="favicon not found")


# ------------------------- Easter-egg sprite API --------------------------- #

@app.get("/api/sprites")
def sprite_animals():
    """List the available pet sprite sets (animals) for the easter-egg picker."""
    return {"animals": list_sprite_animals(), "frames": SPRITE_FILES}


@app.get("/sprites/{animal}/{fname}")
def sprite_image(animal: str, fname: str):
    """Serve one sprite PNG, e.g. /sprites/tabby/still or /sprites/dog/nrun1."""
    name = fname[:-4] if fname.lower().endswith(".png") else fname
    try:
        data = _read_sprite(animal, name)
    except KeyError:
        raise HTTPException(status_code=404, detail="sprite not found")
    return Response(content=data, media_type="image/png")


# ------------------------- Vendored libraries ----------------------------- #
# Serve MapLibre GL JS (and CSS) from our OWN origin. On first request we fetch
# it once from a CDN and cache it to disk; thereafter it is served locally.
# This makes the map work even when the *browser* can't reach public CDNs
# (locked-down networks, offline kiosks, etc.) as long as the SERVER has
# outbound internet. Falls back to multiple CDNs in the page if this is empty.

_VENDOR_DIR = os.path.join(BASE, ".vendor_cache")
_MAPLIBRE_VER = os.environ.get("SOZAI_MAPLIBRE_VER", "4.7.1")
_VENDOR_FILES = {
    "maplibre-gl.js": [
        f"https://unpkg.com/maplibre-gl@{_MAPLIBRE_VER}/dist/maplibre-gl.js",
        f"https://cdn.jsdelivr.net/npm/maplibre-gl@{_MAPLIBRE_VER}/dist/maplibre-gl.js",
        f"https://cdnjs.cloudflare.com/ajax/libs/maplibre-gl/{_MAPLIBRE_VER}/maplibre-gl.js",
    ],
    "maplibre-gl.css": [
        f"https://unpkg.com/maplibre-gl@{_MAPLIBRE_VER}/dist/maplibre-gl.css",
        f"https://cdn.jsdelivr.net/npm/maplibre-gl@{_MAPLIBRE_VER}/dist/maplibre-gl.css",
        f"https://cdnjs.cloudflare.com/ajax/libs/maplibre-gl/{_MAPLIBRE_VER}/maplibre-gl.css",
    ],
}
_vendor_lock = threading.Lock()


def _vendor_fetch(fname: str):
    """Return cached bytes for a vendored file, downloading once if needed."""
    if fname not in _VENDOR_FILES:
        return None
    cache_path = os.path.join(_VENDOR_DIR, fname)
    if os.path.isfile(cache_path) and os.path.getsize(cache_path) > 0:
        with open(cache_path, "rb") as fh:
            return fh.read()
    with _vendor_lock:
        if os.path.isfile(cache_path) and os.path.getsize(cache_path) > 0:
            with open(cache_path, "rb") as fh:
                return fh.read()
        os.makedirs(_VENDOR_DIR, exist_ok=True)
        import urllib.request
        for url in _VENDOR_FILES[fname]:
            try:
                req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 sozai"})
                with urllib.request.urlopen(req, timeout=20) as resp:
                    data = resp.read()
                if data:
                    with open(cache_path, "wb") as fh:
                        fh.write(data)
                    print(f"[vendor] cached {fname} from {url} ({len(data)} bytes)")
                    return data
            except Exception as exc:
                print(f"[vendor] failed {url}: {exc}")
        return None


_VENDOR_GZIP: dict[str, bytes] = {}


@app.get("/vendor/{fname}")
def vendor_file(fname: str, request: Request):
    """Serve a vendored library (e.g. /vendor/maplibre-gl.js) from local cache.

    MapLibre (~800 KB) is lazy-loaded from here on demand rather than inlined
    into every page, so gzip it on the wire when the client supports it
    (~800 KB -> ~220 KB) and memoize the compressed bytes (vendor files are
    immutable, so compress each one only once)."""
    data = _vendor_fetch(fname)
    if data is None:
        raise HTTPException(status_code=404, detail="vendor file unavailable")
    media = "text/javascript" if fname.endswith(".js") else ("text/css" if fname.endswith(".css") else "application/octet-stream")
    if "gzip" in request.headers.get("accept-encoding", "").lower():
        gz = _VENDOR_GZIP.get(fname)
        if gz is None:
            gz = _VENDOR_GZIP[fname] = gzip.compress(data, 6)
        return Response(content=gz, media_type=media,
                        headers={"Content-Encoding": "gzip", "Vary": "Accept-Encoding"})
    return Response(content=data, media_type=media)


# ------------------------------- Pages ------------------------------------- #

def _read_html(name: str) -> str:
    with open(os.path.join(BASE, name), "r", encoding="utf-8") as fh:
        return fh.read()


_APP_HTML_CACHE: str | None = None


def _app_html() -> str:
    """App HTML, read once and cached in memory (the file doesn't change at
    runtime, so re-reading ~500 KB per request is waste). MapLibre is no longer
    inlined here — the frontend lazy-loads it from /vendor on demand.

    Prefers `index.build.html` — the precompiled output of build_frontend.py,
    which ships plain JS instead of making the browser download @babel/standalone
    and transpile the whole app on first paint. Falls back to `index.html` (the
    editable source, which still works via the in-browser Babel fallback) when
    no build exists, e.g. during local development."""
    global _APP_HTML_CACHE
    if _APP_HTML_CACHE is None:
        built = os.path.join(BASE, "index.build.html")
        if os.path.exists(built):
            if os.path.getmtime(os.path.join(BASE, "index.html")) > os.path.getmtime(built):
                print("[frontend] WARNING: index.html is newer than index.build.html — "
                      "rerun `python build_frontend.py` before deploying.")
            name = "index.build.html"
        else:
            print("[frontend] no index.build.html — serving un-precompiled index.html "
                  "(slow first boot). Run `python build_frontend.py`.")
            name = "index.html"
        _APP_HTML_CACHE = _read_html(name)
    return _APP_HTML_CACHE


# The app document is ~500 KB but gzips to ~130 KB. Nothing in this process adds
# Content-Encoding (no GZipMiddleware — it would buffer/break the SSE collab
# stream), so without this the whole document ships uncompressed on every load,
# which dominates first paint. Pre-gzip the cached document once and serve it
# compressed when the client advertises gzip; this touches ONLY these static
# HTML routes and leaves /api/rooms/{id}/stream (SSE) untouched.
_APP_HTML_GZIP: bytes | None = None


def _serve_html(html: str, request: Request, gz: bytes | None = None) -> Response:
    if "gzip" in request.headers.get("accept-encoding", "").lower():
        body = gz if gz is not None else gzip.compress(html.encode("utf-8"), 6)
        return Response(
            content=body,
            media_type="text/html; charset=utf-8",
            headers={"Content-Encoding": "gzip", "Vary": "Accept-Encoding"},
        )
    return HTMLResponse(html)


def _app_html_gzip() -> bytes:
    global _APP_HTML_GZIP
    if _APP_HTML_GZIP is None:
        _APP_HTML_GZIP = gzip.compress(_app_html().encode("utf-8"), 6)
    return _APP_HTML_GZIP


@app.get("/", response_class=HTMLResponse)
async def landing(request: Request):
    return _serve_html(_read_html("landing.html"), request)


@app.get("/app", response_class=HTMLResponse)
async def solo_app(request: Request):
    return _serve_html(_app_html(), request, _app_html_gzip())


@app.get("/room/{room_id}", response_class=HTMLResponse)
async def room_app(room_id: str, request: Request):
    return _serve_html(_app_html(), request, _app_html_gzip())


# --------------------------- Cleanup daemon -------------------------------- #

def _cleanup_loop():
    runs = 0
    while True:
        time.sleep(300)  # every 5 minutes
        try:
            now = time.time()
            runs += 1
            with DB_LOCK:
                # 1) flag rooms idle for > IDLE_TIMEOUT as inactive (so they
                #    can't be joined) — kept briefly for a grace window.
                _conn.execute(
                    "UPDATE rooms SET is_active=0 WHERE is_active=1 AND ?-last_activity>?",
                    (now, IDLE_TIMEOUT),
                )
                # 2) DELETE rooms (and their content) once they've been idle
                #    past the deletion window. This is keyed on last_activity,
                #    not creation time, so the DB tracks live usage and the
                #    heavy room_state blobs don't pile up.
                _conn.execute(
                    "DELETE FROM rooms WHERE ?-last_activity>?", (now, ROOM_DELETE_AFTER_IDLE)
                )
                # absolute safety cap: nothing survives past HARD_EXPIRATION
                _conn.execute("DELETE FROM rooms WHERE ?-created_at>?", (now, HARD_EXPIRATION))
                # 3) cascade: drop orphaned participants + saved scrapbook state
                _conn.execute(
                    "DELETE FROM participants WHERE room_id NOT IN (SELECT room_id FROM rooms)"
                )
                _conn.execute(
                    "DELETE FROM room_state WHERE room_id NOT IN (SELECT room_id FROM rooms)"
                )
                _conn.commit()
                # 4) reclaim disk space. SQLite doesn't shrink the file on its
                #    own after deletes; incremental_vacuum returns freed pages to
                #    the OS. Run occasionally (hourly) to keep the file compact.
                if runs % 12 == 0:
                    try:
                        _conn.execute("PRAGMA incremental_vacuum;")
                        _conn.commit()
                    except Exception:
                        pass
        except Exception:
            pass


threading.Thread(target=_cleanup_loop, daemon=True).start()


if __name__ == "__main__":
    # Bring up the in-process Phoenix UI + OTel tracing immediately, so the
    # embedded "LLM Trace" section is live as soon as the app starts — no
    # separate `phoenix serve` required.
    try:
        _get_tracer()
        print(f"[phoenix] embedded trace UI ready at {_phoenix_session_url()}")
    except Exception as exc:  # noqa: BLE001
        print(f"[phoenix] startup tracing init skipped: {exc}")

    # On Windows, asyncio's Proactor loop logs a noisy but harmless
    # "ConnectionResetError: [WinError 10054]" whenever a client (browser tab,
    # WebSocket, hot-reload) drops a connection abruptly. It's cosmetic — the
    # connection simply closed — so we wrap the internal callback to ignore that
    # one error. This affects nothing else and is a no-op off Windows.
    import platform as _platform
    if _platform.system() == "Windows":
        try:
            from asyncio.proactor_events import _ProactorBasePipeTransport

            _orig_ccl = _ProactorBasePipeTransport._call_connection_lost

            def _quiet_ccl(self, exc):
                if isinstance(exc, ConnectionResetError):
                    return  # ignore the benign 10054 reset
                return _orig_ccl(self, exc)

            _ProactorBasePipeTransport._call_connection_lost = _quiet_ccl
        except Exception:
            pass  # if internals change, just leave the default behavior

    app.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)