"""The unified local server: decisions, generation, and decide-then-generate from one cache.

    uv run --group serve python -m serve.server --artifacts /path/to/artifacts --port 8008

POST /v1/systemone        typed decisions with probabilities and the router's flag
POST /v1/chat/completions base-model generation, adapter off
POST /v1/auto             one question: decide, and decode from the state cache if flagged

Plain PyTorch, one request on the model at a time: correctness before speed.
"""

import argparse
import time
import uuid
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse
from pydantic import BaseModel, Field

from dmm.engine import Engine


class SystemOneRequest(BaseModel):
    state: str
    questions: dict[str, dict] = Field(min_length=1)
    model: str = "dev-4b"


class ChatRequest(BaseModel):
    messages: list[dict]
    max_tokens: int = 512
    model: str = "dev-4b"


class AutoRequest(BaseModel):
    state: str
    question: dict
    name: str = "q"
    max_tokens: int = 512


def create_app(engine: Engine) -> FastAPI:
    app = FastAPI(title="Dev-4B")

    @app.get("/", response_class=HTMLResponse)
    def playground():
        """A small browser page for trying the three endpoints by hand."""
        return (Path(__file__).parent / "playground.html").read_text()

    @app.post("/v1/systemone")
    def systemone(req: SystemOneRequest):
        t0 = time.perf_counter()
        try:
            answers = engine.decide(req.state, req.questions)
        except ValueError as e:
            raise HTTPException(422, str(e)) from e
        return {
            "model": req.model,
            "answers": answers,
            "latency_ms": round((time.perf_counter() - t0) * 1000, 1),
        }

    @app.post("/v1/chat/completions")
    def chat(req: ChatRequest):
        tokens = engine.chat(req.messages, req.max_tokens)
        text = engine.tok.decode(tokens, skip_special_tokens=True)
        return {
            "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
            "object": "chat.completion",
            "model": req.model,
            "tokens": tokens,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": text},
                    "finish_reason": "stop",
                }
            ],
            "usage": {"completion_tokens": len(tokens)},
        }

    @app.post("/v1/auto")
    def auto(req: AutoRequest):
        try:
            return engine.auto(req.state, req.name, req.question, req.max_tokens)
        except ValueError as e:
            raise HTTPException(422, str(e)) from e

    return app


def main():
    import uvicorn
    import yaml

    from dmm.load import load

    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/stage6.yaml")
    ap.add_argument("--artifacts", required=True, help="dir with best.pt and temperatures.json")
    ap.add_argument("--router", default=None, help="dir with router.pt and operating_point.json")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8008)
    a = ap.parse_args()
    cfg = yaml.safe_load(open(a.config))
    base = yaml.safe_load(open(cfg["base_config"]))
    model, tok = load(base["model"], device=a.device)
    engine = Engine.from_artifacts(model, tok, cfg, a.artifacts, a.router or a.artifacts)
    uvicorn.run(create_app(engine), host=a.host, port=a.port)


if __name__ == "__main__":
    main()