"""Run: uv run python server.py --port 8043."""
from __future__ import annotations

import argparse
import threading
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Literal

import uvicorn
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, ConfigDict, Field

from runtime import Runtime, MAX_TOKENS, MODEL_ID, GEMMA_ID, MODELS

ROOT = Path(__file__).parent
ModelChoice = Literal[MODEL_ID, GEMMA_ID]
runtimes = {model_id: Runtime(model_id) for model_id in MODELS}


class Message(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str = Field(max_length=50000)


class Input(BaseModel):
    model_config = ConfigDict(extra="forbid")
    model: ModelChoice = MODEL_ID
    mode: Literal["chat", "raw"] = "chat"
    messages: list[Message] = Field(default_factory=list, max_length=40)
    text: str = Field(default="", max_length=100000)
    thinking: bool = False
    max_new_tokens: int = Field(default=0, ge=0, le=128)


class AppendInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    text: str | None = Field(default=None, max_length=20000)


@asynccontextmanager
async def lifespan(app):
    def load_models():
        for runtime in runtimes.values():
            runtime.load()
    threading.Thread(target=load_models, daemon=True).start()
    yield


app = FastAPI(lifespan=lifespan)


@app.get("/api/status")
def status(model: ModelChoice = MODEL_ID):
    return dict(**runtimes[model].status(), models=[dict(id=mid, **spec) for mid, spec in MODELS.items()])


@app.post("/api/prepare")
def prepare(payload: Input):
    runtime = runtimes[payload.model]
    try:
        text, ids = runtime.prepare(payload)
        return dict(text=text, count=len(ids), tokens=[runtime.token(i) for i in ids],
                    max_tokens=MAX_TOKENS, fits=len(ids) + payload.max_new_tokens <= MAX_TOKENS)
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc


@app.post("/api/analyze")
def analyze(payload: Input):
    try:
        return runtimes[payload.model].analyze(payload)
    except BlockingIOError as exc:
        raise HTTPException(409, str(exc)) from exc
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc
    except Exception as exc:
        import traceback
        traceback.print_exc()
        raise HTTPException(500, f"Computation failed: {exc}") from exc


def get_session(sid):
    for runtime in runtimes.values():
        try:
            return runtime, runtime.session(sid)
        except KeyError:
            continue
    raise HTTPException(404, "This result has been evicted from memory. Run the analysis again.")


@app.get("/api/runs/{sid}/inspect")
def inspect(sid: str, layer: int = Query(ge=0), query: int = Query(ge=0)):
    runtime, s = get_session(sid)
    if layer >= s["response"]["layers"]:
        raise HTTPException(422, "Layer is outside this model.")
    if query >= len(s["response"]["tokens"]):
        raise HTTPException(422, "Token position is outside the context.")
    return runtime.inspect(sid, layer, query)


@app.post("/api/runs/{sid}/append-preview")
def append_preview(sid: str, payload: AppendInput):
    runtime, session = get_session(sid)
    try:
        ids = runtime.append_ids(sid, payload.text)
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc
    total = len(session["response"]["tokens"]) + len(ids)
    return dict(tokens=[runtime.token(tid) for tid in ids], count=len(ids), total=total,
                max_tokens=MAX_TOKENS, fits=total <= MAX_TOKENS)


@app.post("/api/runs/{sid}/continue")
def continue_run(sid: str, payload: AppendInput | None = None):
    runtime, s = get_session(sid)
    previous = s["response"]
    ids = [token["id"] for token in previous["tokens"]]
    try:
        # Preserve the exact token IDs, including a possible partial UTF-8 token.
        text = payload.text if payload is not None else None
        suffix = runtime.append_ids(sid, text)
        sources = previous["token_sources"] + ["predicted" if text is None else "manual"] * len(suffix)
        return runtime.analyze(Input(model=runtime.model_id, mode="raw"),
                               prefix=(previous["context"], ids + suffix),
                               input_count=previous["input_count"], token_sources=sources)
    except BlockingIOError as exc:
        raise HTTPException(409, str(exc)) from exc
    except ValueError as exc:
        raise HTTPException(422, str(exc)) from exc


@app.get("/api/runs/{sid}/attention")
def attention(sid: str, layer: int = Query(ge=0), head: int | Literal["mean"] = Query(...)):
    _, s = get_session(sid)
    if layer >= s["response"]["layers"] or (head != "mean" and not 0 <= head < s["response"]["heads"]):
        raise HTTPException(422, "Layer or head is outside this model.")
    heads = s["attention"][layer]
    # Average stored post-softmax probabilities, accumulating in float32.
    matrix = heads.mean(axis=0, dtype="<f4") if head == "mean" else heads[head].astype("<f4")
    return Response(matrix.tobytes(), media_type="application/octet-stream",
                    headers={"X-Token-Count": str(matrix.shape[0]), "Cache-Control": "private, max-age=3600"})


@app.get("/api/runs/{sid}/attention-layers")
def attention_layers(sid: str, query: int = Query(ge=0), head: int | Literal["mean"] = "mean"):
    runtime, session = get_session(sid)
    if query >= len(session["response"]["tokens"]):
        raise HTTPException(422, "Token position is outside the context.")
    if head != "mean" and not 0 <= head < session["response"]["heads"]:
        raise HTTPException(422, "Head is outside this model.")
    return runtime.attention_layers(sid, query, head)


@app.get("/")
def index():
    return FileResponse(ROOT / "static" / "index.html")


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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8043)
    parser.add_argument("--host", default="127.0.0.1")
    args = parser.parse_args()
    uvicorn.run(app, host=args.host, port=args.port)