"""booper-chat, served.

A 34M-parameter model answers in well under a second on CPU, so the whole
design here is about not squandering that: weights are baked into the image
(no cold download), the page is one static file with no framework to boot,
and tokens go out over SSE as they are produced rather than after the reply
is finished.
"""

from __future__ import annotations

import json
import os
import threading
from pathlib import Path

from fastapi import FastAPI
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from booper.engine import Booper, SamplingParams, Turn

ROOT = Path(__file__).parent
MODEL_DIR = Path(os.environ.get("BOOPER_MODEL_DIR", ROOT / "weights"))
MODEL_ID = "ProCreations/booper-chat"

# One generation at a time. The Space runs on 2 vCPUs and torch already uses
# both for a single decode, so overlapping two requests makes each of them
# slower rather than making the pair faster; queueing keeps every reply fast
# and the tok/s number honest.
_generation_lock = threading.Lock()
MAX_QUEUE = int(os.environ.get("BOOPER_MAX_QUEUE", "24"))
_waiting = threading.Semaphore(MAX_QUEUE)

# Input caps. The 1024-token window truncates anything longer anyway, but the
# tokenizer would still have to chew through whatever arrives first.
MAX_TURNS = 400
MAX_CHARS = 60_000

booper = Booper(MODEL_DIR / "model.pt", MODEL_DIR / "tokenizer.json")
print(
    f"[booper] {booper.n_params:,} params · block {booper.block_size} · "
    f"vocab {booper.tok.vocab_size} · backend {booper.backend}",
    flush=True,
)

app = FastAPI(title="booper-chat", docs_url=None, redoc_url=None)


class Message(BaseModel):
    role: str = "user"
    content: str = ""


class ChatRequest(BaseModel):
    messages: list[Message] = Field(default_factory=list)
    max_new_tokens: int = 96
    temperature: float = 0.7
    top_k: int = 40
    top_p: float = 0.9
    repetition_penalty: float = 1.15
    frequency_penalty: float = 0.12
    presence_penalty: float = 0.0
    no_repeat_ngram_size: int = 3
    seed: int | None = None


def _sse(event: str, payload: dict) -> str:
    return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n"


@app.get("/api/meta")
def meta() -> JSONResponse:
    return JSONResponse(
        {
            "model": MODEL_ID,
            "params": booper.n_params,
            "context": booper.block_size,
            "vocab": booper.tok.vocab_size,
            "layers": booper.config.n_layer,
            "width": booper.config.n_embd,
            "heads": booper.config.n_head,
            "quantized": booper.quantized,
            "defaults": SamplingParams().__dict__,
        }
    )


@app.post("/api/chat")
def chat(request: ChatRequest) -> StreamingResponse:
    turns = [
        Turn("assistant" if m.role == "assistant" else "user", m.content)
        for m in request.messages[-MAX_TURNS:]
        if m.content and m.content.strip()
    ]
    total = sum(len(t.content) for t in turns)
    while total > MAX_CHARS and len(turns) > 1:
        total -= len(turns.pop(0).content)
    if turns:
        turns[-1] = Turn(turns[-1].role, turns[-1].content[:MAX_CHARS])

    params = SamplingParams(
        max_new_tokens=request.max_new_tokens,
        temperature=request.temperature,
        top_k=request.top_k,
        top_p=request.top_p,
        repetition_penalty=request.repetition_penalty,
        frequency_penalty=request.frequency_penalty,
        presence_penalty=request.presence_penalty,
        no_repeat_ngram_size=request.no_repeat_ngram_size,
        seed=request.seed,
    )

    def body():
        if not turns:
            yield _sse("error", {"message": "say something first"})
            return
        if not _waiting.acquire(blocking=False):
            yield _sse("error", {"message": "too many people talking to booper at once — try again"})
            return
        try:
            queued = not _generation_lock.acquire(blocking=False)
            if queued:
                yield _sse("queued", {})
                _generation_lock.acquire()
            try:
                for kind, payload in booper.stream(turns, params):
                    if kind == "token":
                        yield _sse("token", {"t": payload})
                    else:
                        yield _sse("done", payload)
            finally:
                _generation_lock.release()
        except Exception as exc:  # never leave the client hanging on a dead stream
            yield _sse("error", {"message": f"{type(exc).__name__}: {exc}"})
        finally:
            _waiting.release()

    return StreamingResponse(
        body(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache, no-transform",
            "Connection": "keep-alive",
            # Spaces sit behind a proxy that will happily buffer a whole SSE
            # stream and deliver it as one lump, which would silently undo
            # every bit of the streaming above.
            "X-Accel-Buffering": "no",
        },
    )


@app.get("/")
def index() -> FileResponse:
    return FileResponse(ROOT / "static" / "index.html", headers={"Cache-Control": "no-cache"})


app.mount("/", StaticFiles(directory=ROOT / "static"), name="static")