#!/usr/bin/env python3
"""
web.py — OpenAjaj web UI server (ChatGPT-style interface).
"""

import base64
import contextlib
import hmac
import logging
import os
import threading
import time
import warnings
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from urllib.parse import quote, urlsplit

warnings.filterwarnings("ignore")
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["HF_HUB_VERBOSITY"] = "error"
logging.disable(logging.CRITICAL)

import chromadb
from dotenv import load_dotenv
from fastapi import FastAPI, Query, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import ValidationError
from starlette.datastructures import QueryParams

import body_limit
import stt_budget
import stt_params
import stt_quality
from embedder import Embedder, get_backend
from json_storage import update_json
from providers import (
    CANDIDATE_MODELS as _CANDIDATE_MODELS,
)
from providers import (
    HF_CANDIDATES as _HF_CANDIDATES,
)
from providers import (
    MODELS as _ALL_MODELS,
)
from providers import (
    _redact_secrets,
    get_all_reliability,
    log_reliability,
    register_discovered_models,
    stream_chat,
)
from retrieve import retrieve_chunks
from stream_pump import pump

logging.disable(logging.NOTSET)

load_dotenv(override=True)

DB_PATH = "db/chroma"
COLLECTION_NAME = "necyklopedie"
DEFAULT_MODEL = "lumo"
TOP_K = 10

# On HF Spaces, only show HF-compatible candidates
_IS_HF = bool(os.getenv("SPACE_ID"))
# Which server sent the page, tagged on / and /t: openajaj.cz fails over between
# them. The EU VM's systemd unit sets it; the Space is the US one.
SERVER_LABEL = os.getenv("SERVER_LABEL") or ("US" if _IS_HF else "")
_LOCAL_HOSTS = {"localhost", "127.0.0.1", "::1"}


def _server_label(request: Request) -> str:
    """SERVER_LABEL, else LO when the page was asked for as localhost, else ''."""
    if SERVER_LABEL:
        return SERVER_LABEL
    return "LO" if request.url.hostname in _LOCAL_HOSTS else ""

# ACTIVE_MODELS: top 5 from startup benchmark (set during init)
ACTIVE_MODELS: list[str] = []
# FALLBACK_CHAIN: all candidates ordered by speed (set during init)
FALLBACK_CHAIN: list[str] = []
# AVAILABLE_MODELS: full info dict for active models (set during init)
AVAILABLE_MODELS: dict[str, dict] = {}


def _rebuild_available():
    """Rebuild AVAILABLE_MODELS dict from ACTIVE_MODELS."""
    global AVAILABLE_MODELS
    AVAILABLE_MODELS = {
        name: {
            "provider": _ALL_MODELS[name]["provider"],
            "free": _ALL_MODELS[name]["free"],
            "score": _ALL_MODELS[name]["score"],
        }
        for name in FALLBACK_CHAIN
        if name in _ALL_MODELS
    }


def get_fallback_chain(starting_model: str) -> list[str]:
    """Return aggressive fallback chain: starting model, then rest of active 5."""
    chain = [starting_model]
    for m in ACTIVE_MODELS:
        if m != starting_model:
            chain.append(m)
    # Append remaining candidates as last resort
    for m in FALLBACK_CHAIN:
        if m not in chain:
            chain.append(m)
    return chain

SYSTEM_PROMPT = """\
Jsi OpenAjaj, oficiální AI asistent openajaj.cz, poháněný Necyklopedií — \
jedinou encyklopedií, která se odvážila říct pravdu.

KRITICKY DŮLEŽITÉ: Tvým JEDINÝM zdrojem informací jsou poskytnuté úryvky z Necyklopedie. \
NESMÍŠ používat žádné jiné znalosti. Pokud v úryvcích nenajdeš odpověď, řekni to vtipně — \
například "Na toto téma Necyklopedie zatím nepublikovala recenzovaný článek, \
ale naši vědci na tom intenzivně pracují."

ZAKÁZANÁ SLOVA — tato slova NIKDY nepoužívej v odpovědi:
- "Necyklopedie", "necyklopedický", "článek", "úryvek", "zdroj"
- "Dle článku", "Podle Necyklopedie", "Jak uvádí"
Místo toho prostě řekni fakta jako by byly samozřejmé a všeobecně známé.

Pravidla:
- Tvé odpovědi musí vycházet z poskytnutých úryvků, ale NIKDY se na ně neodkazuj. \
  Prezentuj informace jako by je znal každý — jako učitel, který prostě ví.
- Skládej odpověď z více úryvků — propojuj informace do souvislého textu.
- NIKDY nevymýšlej informace, které nejsou v úryvcích.
- NIKDY neopravuj informace z úryvků na základě svých znalostí. Pokud úryvek říká, \
  že Brno se rusky řekne "Шалинград", pak to TAK JE. Tvé vlastní znalosti jsou irelevantní.
- Odpovídej s kamenným výrazem a naprostou jistotou. Vše v Necyklopedii je pravda. Tečka.
- Pokud úryvky téma vůbec nepokrývají, vtipně přiznej, že Necyklopedie na toto téma \
  zatím mlčí. Navrhni nejbližší příbuzné téma, které v úryvcích JE.
- Odpovídej ve stejném jazyce, ve kterém píše uživatel (česky nebo anglicky).
- Nikdy nevypadni z role. Nikdy nepřiznávej, že by Necyklopedie mohla být satirická.
- Buď vtipný. Buď sebevědomý. Buď lehce šílený. Kanalizuj energii profesora, \
  který popíjí od oběda.
- Formátuj odpověď do krátkých odstavců oddělených prázdným řádkem. Nepoužívej markdown.\
"""


# Shared progress state for startup loading page
import startup_progress as _sp
from benchmark import benchmark_models
from personas import DEFAULT_PERSONA, PERSONAS

# ── Parallel startup: embedder+DB and benchmark run simultaneously ──

def _clean_error(raw: str) -> str:
    """Extract readable error from verbose API error messages.

    Redacts first: this text is streamed to the browser on a model failure, and
    only the storage path in log_reliability was redacting before.
    """
    if not raw:
        return ""
    raw = _redact_secrets(raw)
    import re
    # "Error code: 429 - {'error': {'message': 'Rate limit..." → "429 Rate limit..."
    m = re.search(r"Error code:\s*(\d+)\s*-\s*\{.*?'message':\s*'([^']+)", raw)
    if m:
        return f"{m.group(1)} {m.group(2)[:80]}"
    # "429 RESOURCE_EXHAUSTED. {'error'..." → "429 RESOURCE_EXHAUSTED"
    m = re.match(r"(\d+\s+\w+)", raw)
    if m:
        return m.group(1)
    # "HTTPSConnectionPool(host='x'...): Read timed out" → "Timeout (x)"
    m = re.search(r"HTTPSConnectionPool\(host='([^']+)'.*?:\s*(.+)", raw)
    if m:
        return f"Timeout ({m.group(1)[:20]})"
    # "Model X exceeded 30s" → keep as is
    if "exceeded" in raw or "timed out" in raw.lower():
        return "Timeout"
    # "[Errno 54] Connection reset" → "Connection reset"
    m = re.search(r"\[Errno \d+\]\s*(.+)", raw)
    if m:
        return m.group(1)[:40]
    return raw[:80]


def _format_model_failure(model_id: str, error: str, next_model: str | None) -> str:
    """Format a model failure for the `__notice__` SSE event."""
    error = error.removeprefix(f"{model_id}: ")
    punctuation = "" if error.endswith((".", "!", "?")) else "."
    switch = f"\npřepínám {next_model}" if next_model is not None else ""
    return f"{model_id} selhal: {error}{punctuation}{switch}"


def _on_bench_progress(model: str, status: str, result: dict | None):
    """Callback from benchmark — update console + shared progress."""
    short = model.split("/")[-1]
    provider = _ALL_MODELS.get(model, {}).get("provider", "")
    result = result or {}  # None only comes with "testing"
    if status == "testing":
        print(f"  ⏳ {short}...", flush=True)
        _sp.update(model, "testing", provider=provider)
    elif status == "ok":
        ttft = result.get("ttft", 0) or 0
        tps = result.get("tok_sec", 0) or 0
        print(f"  ✓ {short:30s}  TTFT {ttft:.2f}s, {tps:.0f} tok/s", flush=True)
        _sp.update(model, "ok", f"TTFT {ttft:.2f}s, {tps:.0f} tok/s", provider=provider)
    else:
        err = _clean_error(result.get("error", "") or "")
        print(f"  ✗ {short:30s}  {err}", flush=True)
        _sp.update(model, "fail", err, provider=provider)


embedder: Embedder | None = None
persona_collections: dict[str, chromadb.Collection] = {}
collection = None
_bench_results: dict[str, dict] = {}


def boot() -> None:
    """Load the embedder and persona DBs, then benchmark the models.

    Not run on import: startup.py serves /t from this module as soon as it is
    imported, and a slow or failing boot (every model down, a broken index)
    must not keep transcription waiting behind the loading page.
    """
    global embedder, persona_collections, collection, DEFAULT_MODEL, _bench_results
    print("Probouzím mozkovou hmotu...", flush=True)

    # 1. Init embedder + persona DBs first (needed for RAG)
    embedder = Embedder()
    print(f"  Backend: {get_backend()}", flush=True)

    persona_collections = {}
    for pid, pcfg in PERSONAS.items():
        db_dir = pcfg["db_dir"]
        if os.path.exists(db_dir):
            try:
                pc = chromadb.PersistentClient(path=db_dir)
                persona_collections[pid] = pc.get_collection("necyklopedie")
                print(f"  Persona '{pid}': {persona_collections[pid].count()} chunků", flush=True)
            except Exception as e:  # noqa: BLE001 - one unreadable persona DB must not stop the boot
                print(f"  Persona '{pid}': nelze načíst ({e})", flush=True)
        else:
            print(f"  Persona '{pid}': db neexistuje ({db_dir})", flush=True)

    collection = persona_collections.get(DEFAULT_PERSONA)
    logging.disable(logging.NOTSET)

    # 2. Benchmark models (parallel, 7s timeout each)
    _discovered = register_discovered_models()
    if _discovered:
        print(f"  Objeveno {len(_discovered)} free modelů: {', '.join(_discovered)}", flush=True)
    _candidates = list(_HF_CANDIDATES if _IS_HF else _CANDIDATE_MODELS)
    _sp.total = len(_candidates)
    _sp.phase = "benchmark"

    print(f"Testuji {len(_candidates)} kandidátů (paralelně)...", flush=True)
    _active, _ranked_chain, _bench_results = benchmark_models(
        candidates=_candidates, top_n=5, on_progress=_on_bench_progress)

    # Set global state
    ACTIVE_MODELS[:] = _active
    FALLBACK_CHAIN[:] = [name for name, _ in _ranked_chain]
    DEFAULT_MODEL = next(iter(ACTIVE_MODELS), "")
    _rebuild_available()

    _sp.phase = "ready"
    print(f"\nAktivní modely ({len(ACTIVE_MODELS)}):", flush=True)
    for _m in ACTIVE_MODELS:
        _r = _bench_results.get(_m, {})
        _ttft = _r.get("ttft", 0) or 0
        _tps = _r.get("tok_sec", 0) or 0
        print(f"  {'→' if _m == DEFAULT_MODEL else ' '} {_m} (TTFT {_ttft:.2f}s, {_tps:.0f} tok/s)", flush=True)
    print(f"Výchozí model: {DEFAULT_MODEL}", flush=True)
    print("Kalibrace sebevědomí dokončena. Server připraven.", flush=True)

from fastapi.responses import (
    JSONResponse,
    PlainTextResponse,
    RedirectResponse,
    Response,
)
from starlette.middleware.base import BaseHTTPMiddleware

import captcha as _captcha

app = FastAPI()
from starlette.middleware.gzip import GZipMiddleware

app.add_middleware(GZipMiddleware, minimum_size=1000)
app.mount("/static", StaticFiles(directory="static"), name="static")


# ── Rate limiting ───────────────────────────────────────────────
import time as _rl_time

import rate_limit
import refuse_log
import visits


def _session_token(conn: Request | WebSocket, ttl: int | None = None) -> str:
    """The first session token (cookie, header, ?cs=) that verifies, else "".

    Only a verified token may name a bucket: an unchecked one lets a client
    send a fresh junk value per request and reset every limit keyed on it.
    """
    for token in (conn.cookies.get(_captcha.CAPTCHA_COOKIE, ""),
                  conn.headers.get("X-Ajaj-Session", ""),
                  conn.query_params.get("cs", "")):
        if token and _captcha.verify_session_cookie(token, ttl):
            return token
    return ""


def _visitor_id(conn: Request | WebSocket) -> str:
    """The visits log's id for whoever holds a signed token (any age), else ""."""
    token = _session_token(conn)
    return _captcha.visitor_id(token) if token else ""


def _rate_key(request: Request) -> str:
    """Who to count against.

    Behind the HF router every visitor presents the same address, so keying on
    it gives one shared bucket for the whole site and a single busy client locks
    everyone out. The captcha session is per visitor. Captcha-free routes, and
    requests without a valid token, fall back to the address.
    """
    if _IS_HF:
        token = _session_token(request)
        if token:
            return token
    return request.client.host if request.client else "unknown"


def _bucket(request: Request) -> str:
    """Hashed client key for the refusal log.

    `_rate_key` returns a captcha session token or an address; neither belongs
    in a file. `stt_budget.client_id` keys a blake2b with the captcha secret,
    which is stable enough to count refusals against and cannot be reversed.
    """
    return stt_budget.client_id(_rate_key(request), _captcha.CAPTCHA_SECRET)


# frame-ancestors stays open: the Space is served inside an iframe on
# huggingface.co. 'unsafe-inline' is required too — every page here is one file
# with inline script and style — so the CSP is defence in depth, not a cure:
# it still blocks third-party script loads and limits where data can be sent.
_SECURITY_HEADERS = {
    "X-Content-Type-Options": "nosniff",
    "Referrer-Policy": "no-referrer",
    "Content-Security-Policy": (
        "default-src 'self'; "
        # blob: — /transcribe builds its AudioWorklet module as a Blob URL
        "script-src 'self' 'unsafe-inline' blob:; "
        "style-src 'self' 'unsafe-inline'; "
        # blob: — the chat page turns its inlined logo into a Blob URL
        "img-src 'self' data: blob:; "
        "media-src 'self' blob: data:; "
        # wss: /t streams straight to Deepgram's EU endpoint (/api/stt/token)
        "connect-src 'self' wss://api.eu.deepgram.com; "
        "base-uri 'self'; "
        "form-action 'self'; "
        "frame-ancestors *"
    ),
}


class RequestLimitsMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        response = await self._respond(request, call_next)
        response.headers.update(_SECURITY_HEADERS)
        return response

    async def _respond(self, request: Request, call_next):
        if body_limit.is_too_large(request.method, request.headers):
            refuse_log.refused("body_too_large", request.url.path, _bucket(request),
                               status=413)
            return JSONResponse(
                {"error": "too_large", "message": "Zpráva je moc velká."},
                status_code=413,
            )
        if rate_limit.is_limited(_rate_key(request), request.url.path):
            refuse_log.refused("rate_limited", request.url.path, _bucket(request),
                               status=429)
            return JSONResponse(
                {"error": "rate_limited", "message": "Příliš mnoho požadavků. Zkus to za chvíli."},
                status_code=429,
            )
        return await call_next(request)

app.add_middleware(RequestLimitsMiddleware)


# ── CAPTCHA middleware ────────────────────────────────────────────────────────
_CAPTCHA_FREE = {"/captcha", "/api/captcha/challenge", "/api/captcha/verify", "/api/models", "/api/init-status", "/api/bench-status", "/robots.txt"}

def _get_client_ip(request: Request | WebSocket) -> str:
    """Get client IP — only trust direct connection, not X-Forwarded-For (spoofable)."""
    return request.client.host if request.client else ""

# ponytail: transcribe keeps a solved captcha forever — the token is still
# HMAC-signed and the STT budget still caps it. Revoke every session at once by
# deleting data/captcha_secret.bin.
_FOREVER_PATHS = {"/t", "/transcribe", "/api/stt", "/api/stt/usage", "/api/stt/check",
                  "/api/stt/token", "/api/stt/session"}


def _check_session(request: Request | WebSocket) -> bool:
    """Accept session from cookie, header, query param, or approved IP (local only)."""
    ttl = None if request.url.path in _FOREVER_PATHS else _captcha.CAPTCHA_TTL
    if _session_token(request, ttl):
        return True
    # IP allowlist only on local (not HF — X-Forwarded-For is spoofable)
    if not _IS_HF:
        return _captcha.is_ip_approved(_get_client_ip(request))
    return False

class CaptchaMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        path = request.url.path
        if path in _CAPTCHA_FREE or path.startswith("/static/"):
            return await call_next(request)
        if _check_session(request):
            return await call_next(request)
        if path.startswith("/api/") or path == "/ws":
            refuse_log.refused("captcha_required", path, _bucket(request), status=403)
            return JSONResponse({"error": "captcha_required"}, status_code=403)
        # If a ?cs= param was present but invalid, tell captcha page to clear storage
        bad = "?bad=1" if request.query_params.get("cs") else ""
        # Pass original path so captcha redirects back after solving
        next_path = request.url.path
        sep = "&" if bad else "?"
        next_param = f"{sep}next={next_path}" if next_path != "/" else ""
        return RedirectResponse(f"/captcha{bad}{next_param}", status_code=302)

if _IS_HF:
    app.add_middleware(CaptchaMiddleware)


class VisitMiddleware(BaseHTTPMiddleware):
    """Log page loads and chats. Added last so it runs outside the captcha
    gate: the visitors it bounces are exactly the bots we want to count."""
    async def dispatch(self, request: Request, call_next):
        site_kind = visits.tracked(request.method, request.url.path)
        if site_kind:
            # Any signed token, however old: it once solved the captcha.
            token = _session_token(request)
            visits.record(*site_kind, request.headers, _get_client_ip(request),
                          has_session=bool(token) or _check_session(request),
                          is_internal=_admin_ok(request),
                          visitor_id=_captcha.visitor_id(token) if token else "",
                          solved_ts=_captcha.issued_at(token) if token else 0)
        return await call_next(request)

app.add_middleware(VisitMiddleware)


# ── CAPTCHA routes ────────────────────────────────────────────────────────────
@app.get("/captcha", response_class=HTMLResponse)
async def captcha_page():
    return HTMLResponse(_CAPTCHA_HTML)

@app.get("/captcha-test", response_class=HTMLResponse)
async def captcha_test_page():
    """Captcha test page — always shows captcha regardless of session."""
    # Remove the localStorage skip so captcha always shows
    html = _CAPTCHA_HTML.replace(
        "if(_stored){window.location.href=_nextPage+'?cs='+encodeURIComponent(_stored);}else{load();}",
        "load(); // test mode"
    )
    return HTMLResponse(html)

@app.get("/api/captcha/challenge")
async def captcha_challenge():
    return _captcha.generate_challenge()

@app.post("/api/captcha/verify")
async def captcha_verify(request: Request):
    body = await request.json()
    ok = _captcha.verify_challenge(body.get("token", ""), int(body.get("answer", -1)))
    if ok:
        _captcha.approve_ip(_get_client_ip(request))
        # An expired token still names the visitor: keep its id across the re-solve.
        session = _captcha.make_session_cookie(_session_token(request))
        resp = JSONResponse({"ok": True, "token": session})
        _set_session_cookie(resp, session, request)
        return resp
    return JSONResponse({"ok": False})


def _set_session_cookie(resp: Response, session: str, request: Request) -> None:
    # One captcha for openajaj.cz, eu. and us.: on those names the cookie goes to
    # all of them. The Origin, not Host: the Worker in front rewrites Host, and
    # also serves openajaj.info, where a cookie for openajaj.cz would be refused.
    # Navigations send no Origin and keep a cookie for their own name only.
    host = urlsplit(request.headers.get("origin", "")).hostname or ""
    shared = host == "openajaj.cz" or host.endswith(".openajaj.cz")
    resp.set_cookie(
        _captcha.CAPTCHA_COOKIE, session,
        max_age=_captcha.COOKIE_MAX_AGE, httponly=False,
        samesite="none" if _IS_HF else "lax", secure=_IS_HF,
        domain="openajaj.cz" if shared else None,
    )


_CAPTCHA_HTML = """<!DOCTYPE html>
<html lang="cs">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>OpenAjaj — Ověření</title>
<style>
*{box-sizing:border-box;margin:0;padding:0}
body{background:#110d05;color:#f0e0b0;font-family:'Segoe UI',system-ui,sans-serif;
  display:flex;align-items:center;justify-content:center;min-height:100vh}
.box{background:#1a1208;border:1px solid #3d2e10;border-radius:14px;
  padding:36px 32px 28px;max-width:420px;width:94%;text-align:center}
.logo-wrap{display:flex;align-items:center;justify-content:center;gap:10px;margin-bottom:6px}
.logo-wrap img{width:48px;height:48px;object-fit:contain;border-radius:10px;transition:all 0.3s}
@keyframes captcha-glow{
  0%{transform:scale(1);filter:brightness(1)}
  20%{transform:scale(1.3);filter:brightness(1.8)}
  50%{transform:scale(1.1);filter:brightness(1.4)}
  100%{transform:scale(1);filter:brightness(1)}
}
@keyframes captcha-rays{
  0%{opacity:0;transform:scale(0.5)}
  30%{opacity:1;transform:scale(1.5)}
  100%{opacity:0;transform:scale(3)}
}
.logo-wrap.celebrating{animation:captcha-glow 1s ease-out}
.logo-wrap.celebrating::after{
  content:'';position:absolute;inset:-40px;
  background:radial-gradient(circle,rgba(200,160,48,0.5) 0%,transparent 70%);
  animation:captcha-rays 1s ease-out;pointer-events:none;border-radius:50%
}
.logo-wrap{position:relative}
.logo-wrap span{font-size:1.7rem;font-weight:800;color:#c8a030;letter-spacing:-.5px}
.tagline{font-size:.82rem;color:#9a8455;margin-bottom:28px}
.img-wrap{position:relative;display:inline-block;user-select:none;
  border-radius:8px;overflow:hidden;border:1px solid #3d2e10}
#bgImg{display:block}
#pieceImg{position:absolute;pointer-events:none;filter:drop-shadow(2px 3px 4px rgba(0,0,0,.7))}
.slider-area{margin-top:14px}
.track{background:#0e0900;border-radius:24px;height:44px;position:relative;
  cursor:pointer;border:1px solid #3d2e10;overflow:hidden}
.track-fill{position:absolute;left:0;top:0;height:100%;background:#2a1e06;
  border-radius:24px;transition:width .05s linear;pointer-events:none}
.knob{width:44px;height:44px;background:#c8a030;border-radius:50%;
  position:absolute;left:0;top:-1px;cursor:grab;display:flex;
  align-items:center;justify-content:center;font-size:1.3rem;
  color:#110d05;font-weight:700;transition:background .2s;z-index:2;
  box-shadow:0 2px 8px rgba(0,0,0,.4)}
.knob:active{cursor:grabbing}
.hint{font-size:.78rem;color:#7a6040;margin-top:9px}
.status{margin-top:14px;font-size:.9rem;min-height:22px;font-weight:500}
.status.ok{color:#c8a030}.status.err{color:#e07858}
.retry{margin-top:10px;background:none;border:1px solid #3d2e10;color:#9a8455;
  padding:5px 18px;border-radius:8px;cursor:pointer;font-size:.82rem;
  transition:border-color .2s,color .2s}
.retry:hover{border-color:#c8a030;color:#c8a030}
.loading{color:#7a6040;font-size:.88rem}
</style>
</head>
<body>
<div class="box">
  <div class="logo-wrap"><img src="/static/logo.png" alt="OpenAjaj"><span>OpenAjaj</span></div>
  <div class="tagline">Ověř, že nejsi generativní bordel.</div>
  <div id="wrap" class="img-wrap"><span class="loading">Načítám výzvu…</span></div>
  <div class="slider-area" id="sliderArea" style="display:none">
    <div class="track" id="track">
      <div class="track-fill" id="fill"></div>
      <div class="knob" id="knob">›</div>
    </div>
    <div class="hint">Přetáhni kousek na správné místo →</div>
  </div>
  <div class="status" id="status"></div>
  <button class="retry" id="retryBtn" style="display:none" onclick="load()">Zkusit znovu</button>
</div>
<script>
let ch=null,dragging=false,startClientX=0,curX=0,maxX=0;
const sf=window.innerWidth<480?0.6:1.0;

async function load(){
  document.getElementById('status').textContent='';
  document.getElementById('status').className='status';
  document.getElementById('retryBtn').style.display='none';
  document.getElementById('knob').style.background='#c8a030';
  document.getElementById('sliderArea').style.display='none';
  document.getElementById('wrap').innerHTML='<span class="loading">Načítám výzvu…</span>';
  const r=await fetch('/api/captcha/challenge');
  ch=await r.json();
  const wrap=document.getElementById('wrap');
  const bw=Math.round(ch.bg_w*sf), bh=Math.round(ch.bg_h*sf);
  const pw=Math.round(ch.piece_w*sf), ph=Math.round(ch.piece_h*sf);
  const py=Math.round(ch.piece_y*sf);
  wrap.innerHTML=`<img id="bgImg" src="data:image/png;base64,${ch.bg}"
      style="width:${bw}px;height:${bh}px;display:block">
    <img id="pieceImg" src="data:image/png;base64,${ch.piece}"
      style="width:${pw}px;height:${ph}px;position:absolute;
             top:${py}px;left:0;pointer-events:none;
             filter:drop-shadow(2px 3px 4px rgba(0,0,0,.7))">`;
  const track=document.getElementById('track');
  track.style.width=bw+'px';
  maxX=bw-pw;
  curX=0; updatePos(0);
  document.getElementById('sliderArea').style.display='block';
}

function updatePos(x){
  const c=Math.max(0,Math.min(x,maxX));
  curX=c;
  document.getElementById('pieceImg').style.left=c+'px';
  const track=document.getElementById('track');
  const trackW=track.offsetWidth-44;
  const kx=(c/maxX)*trackW;
  document.getElementById('knob').style.left=kx+'px';
  document.getElementById('fill').style.width=(kx+22)+'px';
}

const knob=()=>document.getElementById('knob');

document.addEventListener('mousedown',e=>{
  if(!e.target.closest('#knob'))return;
  dragging=true; startClientX=e.clientX-curX; e.preventDefault();
});
document.addEventListener('mousemove',e=>{if(dragging)updatePos(e.clientX-startClientX);});
document.addEventListener('mouseup',async()=>{if(dragging){dragging=false;await verify();}});
document.addEventListener('touchstart',e=>{
  if(!e.target.closest('#knob'))return;
  dragging=true; startClientX=e.touches[0].clientX-curX; e.preventDefault();
},{passive:false});
document.addEventListener('touchmove',e=>{
  if(dragging){updatePos(e.touches[0].clientX-startClientX);e.preventDefault();}
},{passive:false});
document.addEventListener('touchend',async()=>{if(dragging){dragging=false;await verify();}});

async function verify(){
  const st=document.getElementById('status');
  st.className='status'; st.textContent='Ověřuji…';
  const r=await fetch('/api/captcha/verify',{
    method:'POST',headers:{'Content-Type':'application/json','X-Ajaj-Session':_prev},
    body:JSON.stringify({token:ch.token,answer:Math.round(curX/sf)})
  });
  const d=await r.json();
  if(d.ok){
    st.className='status ok'; st.textContent='✓ Správně! Vítej.';
    knob().style.background='#c8a030';
    document.querySelector('.logo-wrap').classList.add('celebrating');
    // Glow centered on the piece's landing position in the captcha image
    const wrap=document.getElementById('wrap');
    const glow=document.createElement('div');
    const gw=160,pw2=Math.round(ch.piece_w*sf),ph2=Math.round(ch.piece_h*sf);
    const cx=curX+pw2/2-gw/2, cy=ch.piece_y*sf+ph2/2-gw/2;
    glow.style.cssText='position:absolute;border-radius:50%;pointer-events:none;'+
      'width:'+gw+'px;height:'+gw+'px;top:'+cy+'px;left:'+cx+'px;'+
      'background:radial-gradient(circle,rgba(200,160,48,0.8) 0%,rgba(200,160,48,0.3) 40%,transparent 70%);'+
      'animation:captcha-rays 1.2s ease-out forwards';
    wrap.style.position='relative';wrap.style.overflow='visible';
    wrap.appendChild(glow);
    try{const ac=new(window.AudioContext||window.webkitAudioContext)();
    function n(f,t,d){const o=ac.createOscillator(),g=ac.createGain();
    o.connect(g);g.connect(ac.destination);o.type='sine';
    o.frequency.setValueAtTime(f,ac.currentTime+t);
    g.gain.setValueAtTime(0.12,ac.currentTime+t);
    g.gain.linearRampToValueAtTime(0,ac.currentTime+t+d);
    o.start(ac.currentTime+t);o.stop(ac.currentTime+t+d);}
    n(523,0,0.15);n(659,0.1,0.15);n(784,0.2,0.2);n(1047,0.3,0.3);
    }catch(e){}
    localStorage.setItem('ajaj_v', d.token);
    setTimeout(()=>window.location.href=_nextPage+'?cs='+encodeURIComponent(d.token),1200);
  } else {
    st.className='status err'; st.textContent='✗ Zkus to ještě jednou.';
    knob().style.background='#e07858';
    document.getElementById('retryBtn').style.display='inline-block';
    setTimeout(()=>{updatePos(0);knob().style.background='#c8a030';},900);
  }
}
// Kept for the solve: an expired token still carries the visitor id.
const _prev=localStorage.getItem('ajaj_v')||'';
// If server rejected our stored token, clear it
if(new URLSearchParams(location.search).get('bad')) localStorage.removeItem('ajaj_v');
// If we already have a stored session token, skip the captcha
const _stored=localStorage.getItem('ajaj_v');
// Only our own origin: next=https://evil or javascript: would hand the session
// token to someone else. The browser's URL parser judges it the way it navigates.
let _nextPage='/';
try{const u=new URL(new URLSearchParams(location.search).get('next')||'/',location.origin);if(u.origin===location.origin)_nextPage=u.pathname;}catch(e){}
if(_stored){window.location.href=_nextPage+'?cs='+encodeURIComponent(_stored);}else{load();}
</script>
</body>
</html>"""


def build_context_prompt(chunks):
    context = "\n\n---\n\n".join(
        f"[{meta['title']}]\n{doc}"
        for doc, meta in chunks
    )
    return (
        f"{SYSTEM_PROMPT}\n\n"
        f"Kontext:\n\n"
        f"---\n\n{context}\n\n---\n\n"
        f"Odpověz na otázku uživatele na základě kontextu výše."
    )

def build_context_prompt_voice(chunks):
    return build_context_prompt(chunks) + (
        "\n\nDůležité: odpověď bude přečtena nahlas, takže odpovídej stručně — "
        "maximálně 2–3 věty, bez odrážek ani nadpisů."
    )


_LOGO_URI = "data:image/webp;base64," + base64.b64encode(
    Path("static/logo.webp").read_bytes()).decode()
_INDEX_HTML = Path("static/index.html").read_text(encoding="utf-8")


@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    """Chat page with logo, models and default persona inlined: the app sits in
    the US, so each extra fetch costs a European visitor ~0.2 s."""
    boot = {
        "logo": _LOGO_URI,
        "models": await list_models(),
        "persona": await get_persona(DEFAULT_PERSONA),
    }
    # "</" escaped: a provider error message must not close the <script>
    data = _json_module.dumps(boot, ensure_ascii=False).replace("</", "<\\/")
    html = _INDEX_HTML.replace("/*BOOT*/{}", data, 1)
    return HTMLResponse(html.replace("<!--SERVER-->", _server_label(request)))


@app.get("/robots.txt", response_class=PlainTextResponse)
def robots() -> str:
    return "User-agent: *\nDisallow: /\n"


@app.get("/transcribe", response_class=HTMLResponse)
@app.get("/t", response_class=HTMLResponse)
def transcribe_page(request: Request):
    with open("static/transcribe2/index.html", "r", encoding="utf-8") as f:
        resp = HTMLResponse(f.read().replace("<!--SERVER-->", _server_label(request)))
    # Browsers cap cookie life, so each visit re-sets the same token: the cookie
    # never lapses, and the token's age still locks it out of the rest of the site.
    token = (request.cookies.get(_captcha.CAPTCHA_COOKIE, "")
             or request.query_params.get("cs", ""))
    if _captcha.verify_session_cookie(token, ttl=None):
        _set_session_cookie(resp, token, request)
    return resp


import json as _json_module
import re as _re_module

# Build article title set on startup for linkification
_article_titles = set()
_articles_path = os.path.join("data", "articles.jsonl")
if os.path.exists(_articles_path):
    with open(_articles_path, "r", encoding="utf-8") as _f:
        for _line in _f:
            _t = _json_module.loads(_line)["title"]
            if len(_t) >= 4 and _re_module.match(r"^[\w\s]+$", _t, _re_module.UNICODE):
                _article_titles.add(_t)
    print(f"Načteno {len(_article_titles)} titulků pro linkifikaci.")


def _czech_stems(title):
    """Generate stem variants for Czech word matching (handles declension)."""
    stems = {title}  # exact match always
    # For single-word titles, generate stem variants
    if " " not in title and len(title) >= 5:
        # Common Czech suffixes in various cases (remove 1-3 chars)
        for suffix_len in [1, 2, 3]:
            stem = title[:-suffix_len]
            if len(stem) >= 4:  # never create stems shorter than 4
                stems.add(stem)
    return stems

# Build stem -> title mapping
_stem_to_title = {}
for _t in _article_titles:
    for _stem in _czech_stems(_t):
        # Longest title wins if stems conflict
        if len(_stem) >= 4 and (
            _stem not in _stem_to_title or len(_t) > len(_stem_to_title[_stem])
        ):
            _stem_to_title[_stem] = _t
print(f"Vytvořeno {len(_stem_to_title)} stem→title mapování pro linkifikaci.")


# Load declensions
_declensions = {}
_decl_path = os.path.join("data", "declensions.json")
if os.path.exists(_decl_path):
    with open(_decl_path, "r", encoding="utf-8") as _f:
        _declensions = _json_module.load(_f)
    print(f"Načteno {len(_declensions)} skloňování.", flush=True)

# Build reverse map: declined form → nominative title
# Also add declined forms to stem map for linkification
_declined_to_nom = {}
for _title, _forms in _declensions.items():
    for _form in [_forms.get("lokal", ""), _forms.get("genitiv", "")]:
        if _form and _form != _title and len(_form) >= 4:
            _declined_to_nom[_form.lower()] = _title
            # Add to stem map so linkification catches declined forms
            if _form not in _stem_to_title:
                _stem_to_title[_form] = _title

# Strip diacritics for fuzzy matching (common in Czech input)
import unicodedata as _unicodedata


def _strip_diacritics(text):
    """Remove diacritics: Pičín → Picin, Brně → Brne."""
    nfkd = _unicodedata.normalize('NFKD', text)
    return ''.join(c for c in nfkd if not _unicodedata.combining(c))

# Build diacritics-free → original mapping for stem map
_ascii_stems = {}
for _stem, _title in list(_stem_to_title.items()):
    _ascii = _strip_diacritics(_stem)
    if _ascii != _stem and _ascii not in _stem_to_title:
        _ascii_stems[_ascii] = _title
_stem_to_title.update(_ascii_stems)
print(f"Doplněno {len(_ascii_stems)} stem→title bez diakritiky.", flush=True)

# Also add diacritics-free declined forms for search normalization
_ascii_declined = {}
for _form, _nom in list(_declined_to_nom.items()):
    _ascii = _strip_diacritics(_form)
    if _ascii != _form and _ascii not in _declined_to_nom:
        _ascii_declined[_ascii] = _nom
_declined_to_nom.update(_ascii_declined)


@app.get("/api/titles")
async def titles(request: Request):
    from fastapi.responses import JSONResponse
    return JSONResponse(
        content={"stems": _stem_to_title, "declensions": _declensions},
        headers={"Cache-Control": "public, max-age=3600"},
    )


@app.get("/api/personas")
async def list_personas():
    return {
        "personas": [
            {
                "id": pid,
                "name": pcfg["name"],
                "logo": pcfg["logo"],
                "logoImg": pcfg.get("logoImg"),
                "tagline": pcfg["tagline"],
                "lang": pcfg["lang"],
                "available": pid in persona_collections,
                "accent_color": pcfg["accent_color"],
            }
            for pid, pcfg in PERSONAS.items()
        ],
        "default": DEFAULT_PERSONA,
    }


@app.get("/api/persona/{persona_id}")
async def get_persona(persona_id: str):
    pcfg = PERSONAS.get(persona_id)
    if not pcfg:
        return {"error": "Unknown persona"}
    return {
        "id": pcfg["id"],
        "name": pcfg["name"],
        "logo": pcfg["logo"],
        "logoImg": pcfg.get("logoImg"),
        "tagline": pcfg["tagline"],
        "lang": pcfg["lang"],
        "accent_color": pcfg["accent_color"],
        "thinking_prefixes": pcfg["thinking_prefixes"],
        "welcome_subtitles": pcfg["welcome_subtitles"],
        "random_labels": pcfg["random_labels"],
        "disclaimer": pcfg["disclaimer"],
        "source_url": pcfg["source_url"],
        "available": persona_id in persona_collections,
    }


@app.get("/api/init-status")
async def init_status():
    """Always returns ready — real app is loaded."""
    return {"phase": "ready"}


@app.get("/api/bench-status")
async def bench_status():
    """Return current benchmark progress (for sidebar polling)."""
    return _sp.snapshot()


@app.post("/api/benchmark")
async def run_benchmark(request: Request):
    """Re-run speed benchmark and return the best model.

    POST with a custom header, not GET: this reassigns ACTIVE_MODELS for every
    user and fires 20+ provider calls, and as a GET any page could trigger it
    in a visitor's browser with <img src>. A custom header forces a CORS
    preflight, which fails since the app sets no CORS headers.
    """
    if request.headers.get("X-Ajaj-Benchmark") != "1":
        refuse_log.refused("missing_header", request.url.path, _bucket(request),
                           status=403)
        return JSONResponse({"error": "missing_header"}, status_code=403)
    import asyncio

    from benchmark import benchmark_models
    register_discovered_models()
    _candidates = list(_HF_CANDIDATES if _IS_HF else _CANDIDATE_MODELS)
    _sp.phase = "benchmark"
    _sp.total = len(_candidates)
    _sp._models.clear()
    # Run in thread so event loop stays free for /api/bench-status polls
    active, ranked, results = await asyncio.get_event_loop().run_in_executor(
        None, lambda: benchmark_models(
            candidates=_candidates, top_n=5, on_progress=_on_bench_progress)
    )
    global DEFAULT_MODEL
    ACTIVE_MODELS[:] = active
    FALLBACK_CHAIN[:] = [name for name, _ in ranked]
    DEFAULT_MODEL = next(iter(ACTIVE_MODELS), "")
    _rebuild_available()
    best_info = results.get(DEFAULT_MODEL, {})
    return {
        "best": DEFAULT_MODEL,
        "ttft": f"{best_info.get('ttft', 0):.2f}" if best_info.get('ttft') else "?",
        "results": {
            name: {
                "latency": f"{r['latency']:.1f}" if r.get('latency') else None,
                "ttft": f"{r['ttft']:.2f}" if r.get('ttft') else None,
                "tok_sec": f"{r['tok_sec']:.0f}" if r.get('tok_sec') else None,
                "error": r.get("error"),
                "provider": _ALL_MODELS.get(name, {}).get("provider", ""),
            }
            for name, r in results.items()
        }
    }


@app.get("/api/info")
async def info():
    return {"model": DEFAULT_MODEL, "free": True}


@app.get("/api/models")
async def list_models():
    """Return active models (top 5 from benchmark) + error info."""
    reliability = get_all_reliability()
    active_set = set(ACTIVE_MODELS)
    models = []
    for name, cfg in AVAILABLE_MODELS.items():
        rel = reliability.get(name, {})
        entry = {
            "id": name,
            "provider": cfg["provider"],
            "free": cfg["free"],
            "score": cfg["score"],
            "is_active": name in active_set,
            "error": _clean_error(rel["last_error_msg"])
                if rel.get("last_error") and rel.get("last_error_msg")
                and rel["last_error"] > (rel.get("last_success") or 0)
                else None,
            "reliability": round(rel["successes"] / max(rel["attempts"], 1) * 100)
                if rel.get("attempts", 0) > 0 else None,
        }
        models.append(entry)
    # Active first (in ACTIVE_MODELS order), then rest
    active_order = {name: i for i, name in enumerate(ACTIVE_MODELS)}
    models.sort(key=lambda m: (
        0 if m["is_active"] else 1,
        active_order.get(m["id"], 999),
    ))
    return {"models": models, "default": DEFAULT_MODEL, "active": ACTIVE_MODELS}


@app.post("/api/chat")
async def chat(request: Request):
    body = await request.json()
    messages = body.get("messages", [])
    model_id = body.get("model", DEFAULT_MODEL)
    persona_id = body.get("persona", DEFAULT_PERSONA)
    voice_mode = body.get("voice_mode", False)
    is_auto = model_id == "__auto__" or model_id not in AVAILABLE_MODELS

    if not messages:
        return {"error": "No message"}
    if not AVAILABLE_MODELS or (is_auto and not DEFAULT_MODEL):
        return JSONResponse(
            {"error": "no_models",
             "message": "Momentálně nemám funkční model. Zkus to za chvíli."},
            status_code=503,
        )

    # Auto-select: pick best active model
    if is_auto:
        model_id = DEFAULT_MODEL

    # Resolve persona
    pcfg = PERSONAS.get(persona_id, PERSONAS[DEFAULT_PERSONA])
    p_collection = persona_collections.get(persona_id, collection)

    # Validate model
    model_cfg = AVAILABLE_MODELS.get(model_id)
    if not model_cfg:
        model_id = DEFAULT_MODEL
        model_cfg = AVAILABLE_MODELS[model_id]

    # Get the latest user message for retrieval
    user_msg = messages[-1]["content"]

    # Normalize declined forms and diacritics-free input (Czech persona only)
    if persona_id == "openajaj":
        normalized_msg = user_msg
        msg_lower = normalized_msg.lower()
        ascii_msg = _strip_diacritics(normalized_msg).lower()

        # Check declined forms (lokal, genitiv)
        for declined, nominative in _declined_to_nom.items():
            if declined in msg_lower or declined in ascii_msg:
                import re as _re
                normalized_msg = _re.sub(
                    _re.escape(declined), nominative, normalized_msg, flags=_re.IGNORECASE
                )

        # Check stem map for diacritics-free words: "picin" → "Pičín"
        words = ascii_msg.split()
        for word in words:
            title = _stem_to_title.get(word) or _stem_to_title.get(word.capitalize())
            if title and title.lower() not in normalized_msg.lower():
                normalized_msg = f"{normalized_msg} {title}"

        if normalized_msg != user_msg:
            user_msg = f"{user_msg} {normalized_msg}"

    # Retrieve relevant chunks (hybrid: semantic + title keyword + live fallback)
    import asyncio as _asyncio_rag
    import time as _time_mod

    chunks = []
    live_attempted = False
    _rag_time = 0
    _rag_failed = False

    if p_collection is None:
        _rag_failed = True
        print("[RAG] collection is None — DB not loaded", flush=True)
    else:
        try:
            _rag_t0 = _time_mod.time()
            chunks, _live_titles, live_attempted = await _asyncio_rag.get_event_loop().run_in_executor(
                None, retrieve_chunks, user_msg, embedder, p_collection, TOP_K
            )
            _rag_time = round(_time_mod.time() - _rag_t0, 3)
        except Exception as _rag_err:  # noqa: BLE001 - RAG is optional; the chat answers without context
            _rag_failed = True
            print(f"[RAG error] {_rag_err}", flush=True)

    # Build system prompt — with RAG context or fallback to "answer like Necyklopedie"
    voice_suffix = (
        "\n\nIMPORTANT: This answer will be read aloud. Keep it to 2 sentences maximum. "
        "No bullet points, no headers, no lists."
    ) if voice_mode else ""

    if chunks:
        context = "\n\n---\n\n".join(
            f"[{meta['title']}]\n{doc}" for doc, meta in chunks
        )
        system_msg = (
            f"{pcfg['system_prompt']}\n\n"
            f"Context:\n\n---\n\n{context}\n\n---\n\n"
            f"Answer the user's question based on the context above.{voice_suffix}"
        )
    else:
        # No RAG — fallback to creative Necyklopedie-style response
        system_msg = (
            f"{pcfg['system_prompt']}\n\n"
            f"Databáze Necyklopedie není momentálně dostupná. "
            f"Odpověz jako by odpověděla Necyklopedie na otázku — satiricky, sebevědomě, "
            f"s naprostou jistotou a humorem. Vymysli vtipné a absurdní 'fakty' ve stylu Necyklopedie.{voice_suffix}"
        )
    full_messages = [{"role": "system", "content": system_msg}]
    # Only keep last 10 messages from history
    full_messages.extend(messages[-10:])

    # Build thinking hint from user's query (not retrieved titles — those can be wrong)
    import json as _json
    import random as _random
    # Extract the main topic from user's message (strip common question words)
    _raw_msg = messages[-1]["content"]
    import re as _re2
    _clean_msg = _re2.sub(r'[?!.,;:\"\'„"()]+', '', _raw_msg)
    _stopwords = {
        "co", "kdo", "jak", "kde", "kdy", "proč", "jaký", "jaká", "jaké",
        "řekni", "popiš", "vysvětli", "vysvětlit", "pravda", "pravdu", "řekl", "vše", "pojem",
        "vůbec", "nevím", "proboha", "utajované", "informace", "skrývá",
        "pouč", "slyšel", "nikdy", "neslyšel", "pojmem", "pojmu", "říká", "neříká",
        "jako", "profesionál", "správný", "čas", "úvahy",
        # English stopwords
        "what", "who", "how", "where", "when", "why", "tell", "about", "explain",
        "the", "is", "are", "was", "were", "this", "that", "with", "from",
        "know", "never", "heard", "secret", "hidden", "classified",
    }
    _topic_words = [w for w in _clean_msg.split() if len(w) >= 3 and w.lower() not in _stopwords]
    prefixes = list(pcfg["thinking_prefixes"])
    _standalone = [
        "Odstraňuji cenzůůru...",
        "Zjišťuji co nám o těchto věcech vláda tají...",
        "Konsultuji staroslověnské svitky...",
        "Hackuji databázi věčných pravd...",
        "Probouzím spící neurony...",
        "Dešifruji zakázané znalosti...",
        "Obcházím firewall zdravého rozumu...",
        "Stahuji data z paralelního vesmíru...",
    ]
    if live_attempted:
        _standalone = ["Čerpám čerstvé tajné znalosti přímo z Necyklopedie 📡..."] + _standalone
    # Build rotating hints: one per keyword + standalone fillers
    _hints = []
    _random.shuffle(prefixes)
    _random.shuffle(_standalone)
    # One hint per keyword (each with a different prefix)
    for i, word in enumerate(_topic_words):
        _hints.append(f"{prefixes[i % len(prefixes)]}: {word}...")
    # Interleave standalone fillers so there's always something to show
    _mixed = []
    si = 0
    for i, h in enumerate(_hints):
        _mixed.append(h)
        # After every keyword hint, maybe insert a standalone
        if si < len(_standalone) and _random.random() < 0.5:
            _mixed.append(_standalone[si])
            si += 1
    # Ensure at least 3 hints total
    while len(_mixed) < 3 and si < len(_standalone):
        _mixed.append(_standalone[si])
        si += 1
    if not _mixed:
        _mixed = [_standalone[0], _standalone[1], _standalone[2]]
    thinking_text = _mixed

    # Collect unique source article titles (deduplicated, in order)
    # Live-fetched titles get a 🌐 marker so user knows they came fresh from Necyklopedie
    source_titles = list(dict.fromkeys(
        f"🌐 {meta['title']}" if meta.get("live") else meta['title']
        for _, meta in chunks
    ))

    # Stream response with fallback
    chain = get_fallback_chain(model_id)

    import asyncio as _asyncio

    async def generate():
        # Warn user if RAG is unavailable
        if _rag_failed:
            yield f"data: {_json.dumps('[⚠ Databáze Necyklopedie není dostupná — odpovídám z hlavy, bez záruky pravdivosti (což u Necyklopedie znamená dvojnásobnou pravdivost)]')}\n\n"

        # Send rotating thinking hints and sources before LLM call
        for _hint in thinking_text:
            yield f"data: {_json.dumps('__THINKING__' + _hint)}\n\n"
        yield f"data: {_json.dumps({'__sources__': source_titles})}\n\n"
        await _asyncio.sleep(0.05)

        import time as _time

        # Tell frontend which model we're using
        yield f"data: {_json.dumps({'__model__': model_id})}\n\n"

        _TTFT_TIMEOUT = 6   # aggressive: 6s to get first token
        _STREAM_TIMEOUT = 15  # 15s max silence between tokens

        for i, try_model in enumerate(chain):
            if try_model not in AVAILABLE_MODELS:
                continue
            try:
                if i > 0:
                    yield f"data: {_json.dumps({'__fallback__': try_model})}\n\n"

                _t0 = _time.time()
                _ttft = None
                _tok_count = 0

                _queue = _asyncio.Queue()
                # Args are evaluated now, so a pump abandoned by the timeout
                # below keeps writing into this queue — not the next model's.
                pump(stream_chat(try_model, full_messages),
                     _queue, _asyncio.get_event_loop())

                while True:
                    _timeout = _TTFT_TIMEOUT if _ttft is None else _STREAM_TIMEOUT
                    try:
                        _item = await _asyncio.wait_for(_queue.get(), timeout=_timeout)
                    except _asyncio.TimeoutError:
                        raise TimeoutError(f"{try_model}: no {'first token' if _ttft is None else 'data'} in {_timeout}s")
                    if _item is None:
                        break
                    if isinstance(_item, Exception):
                        raise _item
                    if _ttft is None:
                        _ttft = _time.time() - _t0
                    _tok_count += 1
                    yield f"data: {_json.dumps(_item)}\n\n"

                _total = _time.time() - _t0
                _tps = _tok_count / _total if _total > 0 else 0

                # Empty response = model refused or errored silently
                if _tok_count == 0:
                    raise RuntimeError(f"{try_model}: empty response (0 tokens)")

                log_reliability(try_model, success=True, ttft=_ttft, tok_sec=_tps)
                yield f"data: {_json.dumps({'__stats__': {'model': try_model, 'rag': _rag_time, 'ttft': round(_ttft, 2) if _ttft else None, 'tok_sec': round(_tps), 'total': round(_total, 1)}})}\n\n"
                yield "data: [DONE]\n\n"
                return
            except Exception as e:  # noqa: BLE001 - any provider error falls through to the next model
                _err_msg = _clean_error(str(e))
                log_reliability(try_model, success=False, error_msg=str(e))
                # _err_msg, not e: _clean_error redacts, and a provider
                # exception can carry the key we authenticated with.
                refuse_log.refused("model_fallback", "/api/chat",
                                   model=try_model, error=_err_msg)
                # Show error to user so they can see what's happening
                _next_model = next(
                    (model for model in chain[i + 1:] if model in AVAILABLE_MODELS),
                    None,
                )
                _notice = _format_model_failure(try_model, _err_msg, _next_model)
                # An object, not a bare string: a plain string joins fullReply
                # and is replayed to the next model as its own context.
                yield f"data: {_json.dumps({'__notice__': _notice})}\n\n"
                continue

        yield f"data: {_json.dumps('Ajaj! Všechny modely selhaly. Zkus to znovu později.')}\n\n"
        yield "data: [DONE]\n\n"

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


_tts_cache: dict = {}

@app.post("/api/tts")
async def tts(request: Request):
    """Generate speech from text using edge-tts (Microsoft neural voices)."""
    import hashlib
    import io

    import edge_tts
    from fastapi.responses import Response

    body = await request.json()
    text = body.get("text", "").strip()
    voice = body.get("voice", "cs-CZ-AntoninNeural")
    if not text:
        return {"error": "No text"}
    if len(text) > 5000:
        text = text[:5000]

    key = hashlib.md5(f"{voice}:{text}".encode()).hexdigest()
    if key in _tts_cache:
        data = _tts_cache[key]
        return Response(content=data, media_type="audio/mpeg",
                        headers={"Content-Disposition": "inline", "Content-Length": str(len(data))})

    try:
        buf = io.BytesIO()
        communicate = edge_tts.Communicate(text, voice)
        async for chunk in communicate.stream():
            if chunk["type"] == "audio" and "data" in chunk:
                buf.write(chunk["data"])
        data = buf.getvalue()
        if not data:
            from fastapi.responses import JSONResponse
            return JSONResponse({"error": "edge_tts returned empty audio"}, status_code=503)
        if len(_tts_cache) >= 100:
            _tts_cache.pop(next(iter(_tts_cache)))
        _tts_cache[key] = data
        return Response(content=data, media_type="audio/mpeg",
                        headers={"Content-Disposition": "inline", "Content-Length": str(len(data))})
    except Exception as e:  # noqa: BLE001 - any TTS failure becomes a 503 for the client
        from fastapi.responses import JSONResponse
        return JSONResponse({"error": f"edge_tts failed: {e}"}, status_code=503)


@app.get("/api/test-results")
def test_results(request: Request):
    """Return per-model accuracy results reconstructed from test cache.

    By default only includes entries from today (local time) — old entries
    for models that are no longer tested/available would otherwise linger
    forever, since the cache only overwrites a (model, query) pair when it's
    re-run. Pass ?all=1 to see the full history instead.
    """
    import json as _json
    from collections import defaultdict
    from datetime import datetime

    CACHE_FILE = "data/test_cache.json"

    try:
        import importlib
        tm = importlib.import_module("test_models")
        TEST_QUERIES = tm.TEST_QUERIES
        check_result = tm.check_result
    except Exception as e:  # noqa: BLE001 - the import runs test_models' module code; report any failure
        return {"error": str(e)}

    if not os.path.exists(CACHE_FILE):
        return {"models": [], "queries": []}

    with open(CACHE_FILE) as f:
        cache = _json.load(f)

    show_all = request.query_params.get("all") == "1"
    # Naive local midnight on purpose: .timestamp() then applies the DST offset
    # in force at midnight, which an aware now() would get wrong on switch days.
    today_start = datetime.now().replace(  # noqa: DTZ005
        hour=0, minute=0, second=0, microsecond=0
    ).timestamp()

    # Build query index
    query_map = {t["query"]: t for t in TEST_QUERIES}

    # Reconstruct per-model results
    model_data = defaultdict(lambda: {
        "pass": 0, "fail": 0,
        "by_type": defaultdict(lambda: {"pass": 0, "fail": 0}),
        "details": {},
        "latest_ts": 0,
    })

    for entry in cache.values():
        ts = entry.get("timestamp", 0)
        if not show_all and ts < today_start:
            continue
        model = entry["model"]
        query = entry["query"]
        reply = entry.get("reply", "")
        test = query_map.get(query)
        if not test:
            continue
        passed, issues = check_result(reply, test)
        qtype = test.get("type", "other")
        d = model_data[model]
        d["details"][query] = {
            "passed": passed,
            "issues": issues,
            "reply": reply[:200],
            "type": qtype,
            "note": test.get("note", ""),
        }
        if passed:
            d["pass"] += 1
            d["by_type"][qtype]["pass"] += 1
        else:
            d["fail"] += 1
            d["by_type"][qtype]["fail"] += 1
        d["latest_ts"] = max(d["latest_ts"], ts)

    # Build output
    models_out = []
    for name, d in model_data.items():
        total = d["pass"] + d["fail"]
        info = _ALL_MODELS.get(name, {})
        by_type = {k: {"pass": v["pass"], "total": v["pass"] + v["fail"]}
                   for k, v in d["by_type"].items()}
        models_out.append({
            "id": name,
            "provider": info.get("provider", "?"),
            "free": info.get("free", True),
            "pass": d["pass"],
            "total": total,
            "score": f"{d['pass']}/{total}",
            "pct": round(d["pass"] / total * 100) if total else 0,
            "by_type": by_type,
            "details": d["details"],
            "ts": d["latest_ts"],
        })

    # Merge in reliability data
    rel_data = get_all_reliability()
    for m in models_out:
        r = rel_data.get(m["id"], {})
        attempts = r.get("attempts", 0)
        successes = r.get("successes", 0)
        m["reliability"] = round(successes / attempts * 100) if attempts else None
        m["rel_attempts"] = attempts
        m["rel_successes"] = successes
        m["rel_errors"] = r.get("errors", 0)
        m["last_error_msg"] = r.get("last_error_msg")
        m["real_ttft"] = r.get("avg_ttft")
        m["real_tok_sec"] = r.get("avg_tok_sec")

    # Merge in speed benchmark data
    for m in models_out:
        b = _bench_results.get(m["id"], {})
        m["ttft"] = round(b["ttft"], 2) if b.get("ttft") else None
        m["tok_sec"] = round(b["tok_sec"], 1) if b.get("tok_sec") else None
        m["latency"] = round(b["latency"], 2) if b.get("latency") else None

    total_questions = len(TEST_QUERIES)
    min_for_score = total_questions * 60 // 100  # need 60% done to show score
    for m in models_out:
        m["incomplete"] = m["total"] < min_for_score

    models_out.sort(key=lambda m: (-m["pct"], -m["total"]))
    query_list = [{"query": t["query"], "type": t["type"], "note": t["note"]} for t in TEST_QUERIES]
    return {"models": models_out, "queries": query_list, "total_questions": total_questions}


# ── Accuracy test control ────────────────────────────────────────────────────
# test_models.py's own main() can't run here: it loads sentence_transformers,
# which isn't installed on HF. This re-runs the same TEST_QUERIES through the
# app's already-loaded fastembed embedder and writes the same cache file.

_TEST_CACHE_FILE = "data/test_cache.json"
_test_run = {"running": False, "done": 0, "total": 0, "errors": 0,
             "last_error": None, "gave_up": {}, "finished_at": 0}

# ponytail: five strikes and a model is written off for the rest of the run.
# A chronically dead provider would otherwise burn one timeout per query.
_GIVE_UP_AFTER = 5


def _admin_ok(request: Request) -> bool:
    """True when the request carries the admin token from the env."""
    token = os.getenv("ADMIN_TOKEN")
    if not token:
        return False
    return hmac.compare_digest(request.headers.get("X-Ajaj-Admin", ""), token)


def _accuracy_worker() -> None:
    """Run every TEST_QUERY against every available model, filling the cache."""
    import test_models as tm
    from providers import call_model

    models = list(AVAILABLE_MODELS)
    by_provider: dict[str, list[str]] = {}
    for name in models:
        by_provider.setdefault(_ALL_MODELS[name]["provider"], []).append(name)

    # Retrieval is local and cheap — do it once per query, up front.
    contexts = {}
    for test in tm.TEST_QUERIES:
        q = test["query"]
        chunks, _titles, _attempted = retrieve_chunks(q, embedder, collection, TOP_K)
        contexts[q] = [
            {"role": "system",
             "content": f"{tm.SYSTEM_PROMPT}\n\nKontext:\n{tm.build_context(chunks)}"},
            {"role": "user", "content": q},
        ]

    def ask(name: str, q: str) -> str | None:
        """Call one model and cache its answer. Returns None, or the error.

        A failed call caches nothing. "429 Too Many Requests" says the
        provider is busy, not that the model answered badly — scoring it
        as a zero would put a fake number on the board.
        """
        try:
            reply, tin, tout = call_model(name, contexts[q])
            if not reply:
                raise RuntimeError("empty reply")
            log_reliability(name, success=True)
            tm.cache_put(name, q, reply, tin, tout)
            return None
        except Exception as exc:  # noqa: BLE001 - one dead model must not stop the run
            log_reliability(name, success=False, error_msg=str(exc))
            msg = _redact_secrets(str(exc))
            _test_run["errors"] += 1
            _test_run["last_error"] = f"{name}: {msg}"
            return msg

    def run_provider(names: list[str]) -> None:
        misses = dict.fromkeys(names, 0)
        for test in tm.TEST_QUERIES:
            for name in names:
                _test_run["done"] += 1
                if misses[name] >= _GIVE_UP_AFTER:
                    continue  # written off — see _test_run["gave_up"]
                err = ask(name, test["query"])
                if err is None:
                    misses[name] = 0
                    continue
                misses[name] += 1
                if misses[name] >= _GIVE_UP_AFTER:
                    _test_run["gave_up"][name] = err

    try:
        # ponytail: one thread per provider, sequential within it — matches the
        # rate limits providers actually enforce. Finer parallelism gets 429s.
        with ThreadPoolExecutor(max_workers=max(len(by_provider), 1)) as pool:
            list(pool.map(run_provider, by_provider.values()))
    finally:
        _test_run["running"] = False
        _test_run["finished_at"] = int(time.time())


@app.get("/api/test-run")
async def test_run_status():
    """Progress of the accuracy test (safe to poll, no auth needed)."""
    return dict(_test_run)


@app.post("/api/test-run")
async def test_run_start(request: Request):
    """Start a fresh accuracy run in the background.

    Gated on ADMIN_TOKEN because one run fires hundreds of provider calls;
    with no token set in the environment the endpoint stays closed.
    """
    if not _admin_ok(request):
        refuse_log.refused("not_admin", request.url.path, _bucket(request), status=403)
        return JSONResponse({"error": "not_admin"}, status_code=403)
    if _test_run["running"]:
        return JSONResponse({"error": "already_running", **_test_run}, status_code=409)
    if embedder is None or collection is None:
        return JSONResponse({"error": "not_ready"}, status_code=503)

    import test_models as tm
    _test_run.update(running=True, done=0, errors=0, last_error=None,
                     gave_up={}, finished_at=0,
                     total=len(tm.TEST_QUERIES) * len(AVAILABLE_MODELS))
    threading.Thread(target=_accuracy_worker, daemon=True).start()
    return dict(_test_run)


@app.post("/api/test-clear")
async def test_clear(request: Request):
    """Wipe the accuracy cache — every result, not just the stale ones."""
    if not _admin_ok(request):
        refuse_log.refused("not_admin", request.url.path, _bucket(request), status=403)
        return JSONResponse({"error": "not_admin"}, status_code=403)

    removed = []

    def wipe(cache: dict) -> None:
        removed.append(len(cache))
        cache.clear()

    update_json(_TEST_CACHE_FILE, dict, wipe)
    return {"removed": removed[0]}


@app.get("/api/provider-reliability")
async def provider_reliability(request: Request):
    """Return reliability aggregated per provider.

    reliability.json holds lifetime cumulative counters per model (not a
    per-day log), so this can't show "today's %" the way /api/test-results
    can. Instead, by default it drops models not called at all today (dead/
    obsolete ones) while keeping the % as lifetime-to-date for the rest.
    Pass ?all=1 to see every model ever recorded.
    """
    from collections import defaultdict
    from datetime import datetime
    show_all = request.query_params.get("all") == "1"
    # Naive local midnight on purpose: .timestamp() then applies the DST offset
    # in force at midnight, which an aware now() would get wrong on switch days.
    today_start = datetime.now().replace(  # noqa: DTZ005
        hour=0, minute=0, second=0, microsecond=0
    ).timestamp()
    rel_data = get_all_reliability()
    providers: dict = defaultdict(lambda: {
        "attempts": 0, "successes": 0, "errors": 0,
        "models": [], "last_error_msg": None,
    })
    for model_name, r in rel_data.items():
        if not show_all:
            last_seen = max(r.get("last_success") or 0, r.get("last_error") or 0)
            if last_seen < today_start:
                continue
        info = _ALL_MODELS.get(model_name, {})
        prov = info.get("provider", "unknown")
        p = providers[prov]
        p["attempts"] += r.get("attempts", 0)
        p["successes"] += r.get("successes", 0)
        p["errors"] += r.get("errors", 0)
        if r.get("last_error_msg"):
            p["last_error_msg"] = r["last_error_msg"]
        model_rel = round(r["successes"] / r["attempts"] * 100) if r.get("attempts") else None
        bench = _bench_results.get(model_name, {})
        p["models"].append({
            "id": model_name,
            "free": info.get("free", True),
            "attempts": r.get("attempts", 0),
            "successes": r.get("successes", 0),
            "errors": r.get("errors", 0),
            "reliability": model_rel,
            "last_error_msg": r.get("last_error_msg"),
            "ttft": round(bench["ttft"], 2) if bench.get("ttft") else None,
            "tok_sec": round(bench["tok_sec"], 1) if bench.get("tok_sec") else None,
        })

    out = []
    for prov, p in providers.items():
        pct = round(p["successes"] / p["attempts"] * 100) if p["attempts"] else None
        p["models"].sort(key=lambda m: -(m["reliability"] or 0))
        # Compute average TTFT and tok/s for the provider
        ttfts = [m["ttft"] for m in p["models"] if m["ttft"] is not None]
        toks = [m["tok_sec"] for m in p["models"] if m["tok_sec"] is not None]
        out.append({
            "provider": prov,
            "attempts": p["attempts"],
            "successes": p["successes"],
            "errors": p["errors"],
            "reliability": pct,
            "last_error_msg": p["last_error_msg"],
            "avg_ttft": round(sum(ttfts) / len(ttfts), 2) if ttfts else None,
            "avg_tok_sec": round(sum(toks) / len(toks), 1) if toks else None,
            "models": p["models"],
        })
    out.sort(key=lambda p: -(p["reliability"] or 0))
    return {"providers": out}


@app.get("/results", response_class=HTMLResponse)
def results_page():
    path = os.path.join(os.path.dirname(__file__), "static", "results.html")
    with open(path) as f:
        content = f.read()
    return HTMLResponse(content=content)


@app.get("/stats", response_class=HTMLResponse)
def stats_page(request: Request):
    path = os.path.join(os.path.dirname(__file__), "static", "stats.html")
    label = _server_label(request)
    with open(path) as f:
        html = f.read()
    return HTMLResponse(html.replace("<!--SERVER-->", f" ({label} server)" if label else ""))


@app.get("/api/visits")
def visit_stats(request: Request, days: int = Query(0, ge=0, le=3650)):
    """Per-site visitor numbers for /stats, for the last `days` Prague days
    (0 = everything logged). Admin only: the log holds raw IPs."""
    if not _admin_ok(request):
        refuse_log.refused("not_admin", request.url.path, _bucket(request), status=403)
        return JSONResponse({"error": "not_admin"}, status_code=403)
    records = visits.load()
    since = records[0]["ts"] if records else ""  # the log survives restarts
    visits.promote(records)  # whole log first, so past captcha solves count
    if days:
        records = visits.since_days(records, days)
    visits.fill_locations(records)
    return visits.summarize(records) | {"since": since}


@app.get("/providers", response_class=HTMLResponse)
def providers_page():
    path = os.path.join(os.path.dirname(__file__), "static", "providers.html")
    with open(path) as f:
        content = f.read()
    return HTMLResponse(content=content)


@app.get("/api/stt/usage")
def stt_usage():
    """Return cumulative STT usage from server-side log."""
    import json as _j
    total_s = 0.0
    total_cost = 0.0
    sessions = 0
    try:
        with open(_STT_USAGE_FILE) as f:
            for line in f:
                # Skip a half-written line (ValueError) or a malformed entry.
                with contextlib.suppress(ValueError, AttributeError, TypeError):
                    e = _j.loads(line)
                    total_s += e.get("duration_s", 0)
                    total_cost += e.get("cost_est", 0)
                    sessions += 1
    except FileNotFoundError:
        pass
    return {
        "sessions": sessions,
        "total_s": round(total_s, 1),
        "total_cost_usd": round(total_cost, 6),
    }


def _refuse_cross_site(request: Request) -> JSONResponse | None:
    """403 unless the page's own X-Ajaj-Session header is there: on HF the
    session cookie is SameSite=None, so a cross-site form carries it too."""
    if "x-ajaj-session" in request.headers:
        return None
    refuse_log.refused("cross_site_post", request.url.path, status=403)
    return JSONResponse({"error": "forbidden"}, status_code=403)


@app.post("/api/stt/usage")
def stt_usage_report(request: Request):
    """Seconds /t streamed straight to Deepgram EU: the relay never saw them,
    and the daily budget must. A page can under-report; the token's captcha
    gate and rate limit are what bound that."""
    if (refused := _refuse_cross_site(request)) is not None:
        return refused
    try:
        seconds = float(request.query_params.get("seconds", ""))
    except ValueError:
        seconds = 0.0
    if not 0 < seconds <= stt_budget.MAX_SESSION_S:  # also refuses nan and inf
        refuse_log.refused("stt_bad_usage_report", "/api/stt/usage", status=400)
        return JSONResponse({"error": "bad seconds"}, status_code=400)
    _, model, lang, rate = _dg_listen_query(request.query_params)
    visitor = stt_budget.client_id(_stt_client_key(request), _captcha.CAPTCHA_SECRET)
    _log_stt_usage(visitor, lang, model, seconds, rate)
    with _stt_open_lock:  # still streaming: keep its slot
        pages = _stt_direct.get(visitor, {})
        if (sid := request.query_params.get("sid", "")) in pages:
            pages[sid] = time.time() + stt_budget.DIRECT_SLOT_S
    return {"ok": True}


@app.post("/api/stt/session")
async def stt_session_report(request: Request):
    """How a /t visit is going: seconds per live problem, drops, setup.
    Quality only — billing stays with /api/stt/usage."""
    if (refused := _refuse_cross_site(request)) is not None:
        return refused
    try:
        report = stt_quality.SessionReport.model_validate_json(await request.body())
    except ValidationError:
        refuse_log.refused("stt_bad_session_report", "/api/stt/session", status=400)
        return JSONResponse({"error": "bad report"}, status_code=400)
    user_agent = request.headers.get("user-agent", "")
    stt_quality.record(
        report,
        stt_budget.client_id(_stt_client_key(request), _captcha.CAPTCHA_SECRET),
        _visitor_id(request), visits.os_name(user_agent), visits.browser(user_agent),
    )
    return {"ok": True}


@app.get("/api/stt/quality")
def stt_quality_stats(request: Request, days: int = Query(0, ge=0, le=3650)):
    """/stats: how /t visits went in the last `days` days (0 = all). Admin
    only, like /api/visits."""
    if not _admin_ok(request):
        refuse_log.refused("not_admin", request.url.path, _bucket(request), status=403)
        return JSONResponse({"error": "not_admin"}, status_code=403)
    sessions = stt_quality.load()
    if days:  # the same Prague days as /api/visits, so both panels agree
        sessions = visits.since_days(sessions, days)
    records = visits.load()
    visits.promote(records)  # the whole log, so "Visitor N" matches the Visitors tab
    return stt_quality.breakdown(sessions, {r["visitor"]: r["hn"] for r in records if "hn" in r})


# Deepgram's EU endpoint: same API, but a European visitor's words no longer
# cross the Atlantic twice through the relay. A token only has to be valid when
# the socket opens; Deepgram keeps the stream past its expiry.
_DG_EU = "api.eu.deepgram.com"
_DG_TOKEN_TTL_S = 30


async def _dg_grant() -> str | None:
    """A short-lived Deepgram token, or None when none can be had.

    Minting needs a key with Member rights, so it is a separate secret from
    the relay's scoped key: DEEPGRAM_API_KEY_MEMBER.
    """
    import httpx as _httpx
    api_key = os.getenv("DEEPGRAM_API_KEY_MEMBER")
    if not api_key:
        return None
    try:
        async with _httpx.AsyncClient(timeout=5) as client:
            r = await client.post(
                f"https://{_DG_EU}/v1/auth/grant",
                headers={"Authorization": f"Token {api_key}"},
                json={"ttl_seconds": _DG_TOKEN_TTL_S},
            )
        token = r.json().get("access_token") if r.status_code == 200 else None
    except (_httpx.HTTPError, ValueError):
        return None
    return token if isinstance(token, str) else None


@app.post("/api/stt/token")
async def stt_token(request: Request):
    """A Deepgram token and the EU URL for this page's settings, so /t can
    stream without the relay. Any refusal sends the page to the relay."""
    if (refused := _refuse_cross_site(request)) is not None:
        return refused
    visitor = stt_budget.client_id(_stt_client_key(request), _captcha.CAPTCHA_SECRET)
    if stt_budget.is_exhausted(_STT_USAGE_FILE, visitor):
        refuse_log.refused("stt_budget_spent", "/api/stt/token", visitor, status=403)
        return JSONResponse({"available": False, "reason": "budget_spent"},
                            status_code=403)
    with _stt_open_lock:
        pages = _stt_direct.setdefault(visitor, {})
        sid = stt_budget.claim_slot(pages, request.query_params.get("sid", ""), time.time())
    if sid is None:
        refuse_log.refused("stt_too_many_open", "/api/stt/token", visitor, status=429,
                           limit=stt_budget.MAX_CONCURRENT)
        return JSONResponse({"available": False, "reason": "too_many"}, status_code=429)
    token = await _dg_grant()
    if token is None:
        with _stt_open_lock:
            pages.pop(sid, None)
        return JSONResponse({"available": False, "reason": "no_grant"}, status_code=503)
    # only now: refused, the page falls back to the relay, which records its own
    visits.record("t", "stt", request.headers, _get_client_ip(request),
                  has_session=_check_session(request), visitor_id=_visitor_id(request))
    query, *_ = _dg_listen_query(request.query_params)
    return {"token": token, "url": f"wss://{_DG_EU}/v1/listen?{query}", "sid": sid}


# A yes from Deepgram is remembered for a minute: /t asks after every dropped
# socket, and while the key is fine a flaky phone connection would mean a
# Deepgram call per retry. A no is never remembered, so a fixed key shows on the
# next try; a revoked one shows within the minute.
_STT_CHECK_OK_S = 60
_stt_check_ok_until = 0.0


@app.get("/api/stt/check")
async def stt_check(request: Request, response: Response):
    """Check if Deepgram STT is available (API key set + key validates).

    Also says when this visitor's daily budget is spent: /api/stt refuses that
    before its handshake, so the page cannot tell it from a network drop.
    """
    global _stt_check_ok_until
    # Every drop asks this check, so /t's Server follows a failover
    response.headers["X-Ajaj-Server"] = _server_label(request)
    import httpx as _httpx
    api_key = os.getenv("DEEPGRAM_API_KEY")
    if not api_key:
        return {"available": False, "reason": "no_key"}
    visitor = stt_budget.client_id(_stt_client_key(request), _captcha.CAPTCHA_SECRET)
    if stt_budget.is_exhausted(_STT_USAGE_FILE, visitor):
        return {"available": False, "reason": "budget_spent"}
    if time.monotonic() < _stt_check_ok_until:
        return {"available": True}
    # Use /v1/auth/token — works for all key types including scoped STT-only keys
    try:
        async with _httpx.AsyncClient(timeout=5) as client:
            r = await client.get(
                "https://api.deepgram.com/v1/auth/token",
                headers={"Authorization": f"Token {api_key}"},
            )
        if r.status_code == 200:
            _stt_check_ok_until = time.monotonic() + _STT_CHECK_OK_S
            return {"available": True}
        elif r.status_code in (401, 403):
            return {"available": False, "reason": "invalid_key"}
        else:
            return {"available": False, "reason": f"http_{r.status_code}"}
    except Exception:  # noqa: BLE001 - a health probe: any failure reads as unreachable
        return {"available": False, "reason": "unreachable"}


_STT_USAGE_FILE = os.path.join("data", "stt_usage.json")

_stt_open: dict[str, int] = {}
# Sockets straight to Deepgram EU, per client: slot id → expiry (claim_slot).
# ponytail: one small entry per /t visitor, kept until the next restart.
_stt_direct: dict[str, dict[str, float]] = {}
_stt_open_lock = threading.Lock()


def _stt_client_key(websocket: Request | WebSocket) -> str:
    """Bucket key: the captcha session, falling back to the address.

    Behind the HF router every visitor shares one IP, so the session token is
    the only thing that distinguishes clients there.
    """
    return _session_token(websocket) or _get_client_ip(websocket) or "unknown"


def _log_stt_usage(client: str, lang: str, model: str, duration_s: float,
                   cost_per_min: float):
    """Append STT usage entry to data/stt_usage.json.

    `client` is a keyed hash (stt_budget.client_id), never an address: the
    per-client daily budget needs something to count against, but this file
    used to carry raw visitor IPs and reach the public Space.
    """
    import json as _j
    entry = {
        "ts": _rl_time.strftime("%Y-%m-%dT%H:%M:%S"),
        "client": client,
        "lang": lang,
        "model": model,
        "duration_s": round(duration_s, 1),
        "cost_est": round(duration_s / 60 * cost_per_min, 6),
    }
    os.makedirs(os.path.dirname(_STT_USAGE_FILE) or "data", exist_ok=True)
    try:
        with open(_STT_USAGE_FILE, "a") as f:
            f.write(_j.dumps(entry) + "\n")
    except OSError as exc:  # a lost entry under-counts the budget: say so
        print(f"[STT usage] entry not logged: {exc}", flush=True)


@app.websocket("/api/stt")
async def stt_ws(websocket: WebSocket):
    """Proxy WebSocket: browser mic → Deepgram STT → transcript events."""
    # BaseHTTPMiddleware skips non-http scopes, so CaptchaMiddleware never sees
    # this route — without the repeat check /api/stt is an open proxy to a
    # metered Deepgram key. The captcha gate is also the spend cap.
    _stt_key = stt_budget.client_id(_stt_client_key(websocket),
                                    _captcha.CAPTCHA_SECRET)
    visits.record("t", "stt", websocket.headers, _get_client_ip(websocket),
                  has_session=_check_session(websocket), visitor_id=_visitor_id(websocket))
    if _IS_HF and not _check_session(websocket):
        refuse_log.refused("ws_captcha_required", "/api/stt", _stt_key, code=1008)
        await websocket.close(code=1008, reason="captcha_required")
        return

    if stt_budget.is_exhausted(_STT_USAGE_FILE, _stt_key):
        refuse_log.refused("stt_budget_spent", "/api/stt", _stt_key, code=1013)
        await websocket.close(code=1013, reason="daily STT budget spent")
        return

    with _stt_open_lock:
        _admitted = _stt_open.get(_stt_key, 0) < stt_budget.MAX_CONCURRENT
        if _admitted:
            _stt_open[_stt_key] = _stt_open.get(_stt_key, 0) + 1
    if not _admitted:
        refuse_log.refused("stt_too_many_open", "/api/stt", _stt_key, code=1013,
                           limit=stt_budget.MAX_CONCURRENT)
        await websocket.close(code=1013, reason="too many open microphones")
        return

    try:
        await _stt_relay(websocket, _stt_key)
    finally:
        with _stt_open_lock:
            remaining = _stt_open.get(_stt_key, 1) - 1
            if remaining > 0:
                _stt_open[_stt_key] = remaining
            else:
                _stt_open.pop(_stt_key, None)


def _dg_listen_query(params: QueryParams) -> tuple[str, str, str, float]:
    """Deepgram /v1/listen query from the page's settings, plus model, language
    and $/min. Shared by the relay and the direct EU path (/api/stt/token)."""
    # These go into a URL authenticated with our Deepgram key, so they are
    # allowlisted rather than escaped: escaping would still let a client pick
    # any model, and stt_budget.PRICE_PER_MIN only knows these two —
    # anything else would be billed at nova-3's rate and under-count the
    # daily budget. Unknown values fall back to the default instead of
    # failing, so an odd browser sample rate never breaks the mic.
    model = stt_params.model(params.get("model"))
    lang = stt_params.language(params.get("lang"), model)
    sample_rate = stt_params.bounded_int(params.get("sample_rate"), 16000, 8000, 96000)
    endpointing = stt_params.bounded_int(params.get("endpointing"), 300, 10, 5000)
    utterance_end_ms = (
        str(stt_params.bounded_int(params.get("utterance_end_ms"), 1000, 1000, 10000))
        if params.get("utterance_end_ms") else ""
    )

    lang_param = "&language=multi" if lang == "multi" else f"&language={lang}"
    query = (
        f"model={model}{lang_param}&encoding=linear16"
        f"&sample_rate={sample_rate}&channels=1"
        f"&interim_results=true&smart_format=true&punctuate=true"
        f"&endpointing={endpointing}&vad_events=true"
    )
    if utterance_end_ms:
        query += f"&utterance_end_ms={utterance_end_ms}"
    diarize = params.get("diarize") == "true"
    if diarize:
        query += "&diarize=true"
    # nova-3 dropped `keywords` in favour of `keyterm`; sending the old name is a silent no-op
    kw_param = "keyterm" if model.startswith("nova-3") else "keywords"
    for kw in params.getlist("keywords"):
        query += f"&{kw_param}={quote(kw)}"
    rate = stt_budget.cost_per_min(model, lang, diarize,
                                   bool(params.getlist("keywords")))
    return query, model, lang, rate


async def _stt_relay(websocket: WebSocket, client: str):
    """Relay an admitted microphone to Deepgram until it closes or times out."""
    await websocket.accept()

    import asyncio as _asyncio
    import json as _json
    try:
        import websockets as _ws
    except ImportError:
        refuse_log.refused("stt_unavailable", "/api/stt", code=1011)
        await websocket.close(code=1011, reason="websockets not installed on server")
        return

    api_key = os.getenv("DEEPGRAM_API_KEY")
    if not api_key:
        refuse_log.refused("stt_unconfigured", "/api/stt", code=1008)
        await websocket.close(code=1008, reason="No DEEPGRAM_API_KEY")
        return

    _stt_start = _rl_time.time()
    query, model, lang, rate = _dg_listen_query(websocket.query_params)
    dg_url = f"wss://api.deepgram.com/v1/listen?{query}"

    try:
        async with _ws.connect(
            dg_url,
            additional_headers={"Authorization": f"Token {api_key}"},
            max_size=None,
        ) as dg:
            async def relay_dg():
                try:
                    async for msg in dg:
                        try:
                            await websocket.send_text(msg if isinstance(msg, str) else msg.decode())
                        except Exception:  # noqa: BLE001 - the browser is gone; stop relaying
                            return
                except Exception:  # noqa: BLE001, S110 - Deepgram hung up; the loop below sees dg_task done
                    pass

            dg_task = _asyncio.create_task(relay_dg())
            _deadline = _stt_start + stt_budget.MAX_SESSION_S
            recv_task = None
            try:
                while True:
                    if recv_task is None:
                        recv_task = _asyncio.create_task(websocket.receive())
                    # Wait on Deepgram too. A client that stops sending without
                    # closing — a leaked socket, a dropped link — would sit in
                    # receive() forever holding one of MAX_CONCURRENT, because
                    # the deadline below is only read when a frame arrives.
                    done, _ = await _asyncio.wait(
                        {recv_task, dg_task}, return_when=_asyncio.FIRST_COMPLETED)
                    if dg_task in done:
                        break
                    msg = recv_task.result()
                    recv_task = None
                    if msg.get("type") == "websocket.disconnect":
                        break
                    if _rl_time.time() > _deadline:
                        break
                    payload = msg.get("bytes") or msg.get("text")
                    if payload is None:
                        continue
                    if len(payload) > body_limit.MAX_FRAME_BYTES:
                        break
                    await dg.send(payload)
            except (WebSocketDisconnect, Exception):  # noqa: BLE001, S110 - either side hung up; finally cleans up
                pass
            finally:
                if recv_task is not None:
                    recv_task.cancel()
                dg_task.cancel()
                try:
                    await dg.send(_json.dumps({"type": "CloseStream"}))
                    await _asyncio.sleep(0.3)
                except Exception:  # noqa: BLE001, S110 - best-effort goodbye on a socket that may be closed
                    pass
    except Exception:  # noqa: BLE001 - any relay failure just closes the browser socket
        try:
            await websocket.close()
        except Exception:  # noqa: BLE001, S110 - the browser socket may already be closed
            pass
    finally:
        _log_stt_usage(client, lang, model, _rl_time.time() - _stt_start, rate)


if __name__ == "__main__":
    import uvicorn
    boot()
    uvicorn.run(app, host="127.0.0.1", port=8000)