"""HTTP server for the Clef-flash decision model.

Clef is a decision model and not a chat model. It reads one state plus a schema
of typed questions and returns one probability per allowed option. Consumers
call POST /v1/systemone and never /v1/chat/completions.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import sys
import traceback
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, AsyncIterator

import torch
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from safetensors.torch import load_file
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration

# uvicorn configures its own loggers only. Without this call the root logger
# stays at WARNING and the startup lines below never reach the pod log.
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
LOG = logging.getLogger("clef-serve")

MODEL_PATH = os.environ.get("MODEL_PATH", "/mnt/models/clef-flash")
DTYPE = os.environ.get("DTYPE", "float16")
MAX_LEN = int(os.environ.get("MAX_LEN", "65536"))
PORT = int(os.environ.get("PORT", "8080"))

# HF_HOME stays untouched. The k8s manifest sets it, and the loader must reuse
# that cache path instead of downloading a second copy of an 18.8 GB checkpoint.

_TORCH_DTYPES = {
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
}


def _resolve_dtype() -> torch.dtype:
    """Map the DTYPE env string to a torch dtype. Fail fast on a typo."""
    try:
        return _TORCH_DTYPES[DTYPE]
    except KeyError:
        raise ValueError(
            f"The server does not support DTYPE={DTYPE!r}. Use one of {sorted(_TORCH_DTYPES)}."
        ) from None


def _load_runtime() -> tuple[Any, Any, list[str], Any]:
    """Load backbone, joint head and processor. Runs in a worker thread."""
    # The released code imports its sibling modules by bare name. The model
    # directory is therefore the only place on sys.path that resolves them.
    sys.path.insert(0, MODEL_PATH)

    # This import must follow the sys.path.insert above. joint_schema_model.py
    # ships inside MODEL_PATH and not in site-packages. An import before that
    # insert raises ModuleNotFoundError. A local import binds the name in this
    # function scope only, so the caller must receive the callable as a value.
    from joint_schema_model import (
        ClefModel,
        JointSchemaHead,
        collate_records,
        encode_record,
        systemone,
    )

    dtype = _resolve_dtype()
    LOG.info("Loading backbone. path=%s dtype=%s", MODEL_PATH, DTYPE)

    # load_release_model passes device_map={"": str(device)} and so pins the
    # whole backbone on one card. This loader asks accelerate to shard instead.
    backbone = Qwen3_5ForConditionalGeneration.from_pretrained(
        MODEL_PATH,
        dtype=dtype,
        device_map="auto",
        low_cpu_mem_usage=True,
    )

    # The head must sit on the device that produces the final hidden state,
    # because ClefModel applies it to the output of the last layer.
    head_cfg = json.loads((Path(MODEL_PATH) / "joint_head_config.json").read_text())
    head = JointSchemaHead(**head_cfg)
    head.load_state_dict(load_file(str(Path(MODEL_PATH) / "joint_head.safetensors")))
    # The output embeddings carry the last layer's device. Without them, take
    # the last parameter, which lives there too.
    output_embeddings = backbone.get_output_embeddings()
    if output_embeddings is not None:
        head_device = output_embeddings.weight.device
    else:
        head_device = next(reversed(list(backbone.parameters()))).device
    head.to(device=head_device, dtype=dtype)

    model = ClefModel(backbone, head).eval()
    processor = AutoProcessor.from_pretrained(MODEL_PATH)

    # hf_device_map maps layer names to ints for CUDA shards. Report it so a
    # human can see whether one card or two carried the weights.
    raw_map = getattr(backbone, "hf_device_map", None) or {}
    if not raw_map:
        raw_map = {"": str(next(backbone.parameters()).device)}
    devices = list(
        dict.fromkeys(
            f"cuda:{name}" if isinstance(name, int) else str(name)
            for name in raw_map.values()
        )
    )
    LOG.info(
        "Model ready. dtype=%s devices=%s device_map=%s head_device=%s",
        DTYPE,
        devices,
        raw_map,
        head_device,
    )
    return model, processor, devices, systemone


@asynccontextmanager
async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
    """Load the model on startup without blocking the event loop."""
    # A 9.4B load takes minutes. asyncio.to_thread keeps /health answerable
    # during that window, so the startup probe sees 503 and not a dead socket.
    model, processor, devices, systemone_fn = await asyncio.to_thread(_load_runtime)
    app.state.model = model
    app.state.processor = processor
    app.state.devices = devices
    # The route handler runs at module scope and cannot see the local import
    # inside _load_runtime. Publish the callable here, or the handler raises
    # NameError on the first request.
    app.state.systemone = systemone_fn
    app.state.lock = asyncio.Lock()
    LOG.info("Runtime published. systemone_callable=%s", systemone_fn is not None)
    yield
    app.state.model = None
    app.state.processor = None
    app.state.systemone = None


def create_app() -> FastAPI:
    app = FastAPI(title="clef-flash", lifespan=_lifespan)

    @app.get("/health")
    async def health() -> JSONResponse:
        # A ready flag alone is not enough. The pod must not take traffic until
        # the weights and the joint head are both resident.
        if getattr(app.state, "model", None) is None:
            return JSONResponse({"status": "loading"}, status_code=503)
        return JSONResponse(
            {
                "status": "ok",
                "dtype": DTYPE,
                "devices": app.state.devices,
            },
            status_code=200,
        )

    @app.post("/v1/systemone")
    async def systemone_route(request: Request) -> Any:
        model = getattr(app.state, "model", None)
        if model is None:
            return JSONResponse({"error": "model is not loaded"}, status_code=503)

        body = await request.json()
        processor = app.state.processor

        def _forward() -> dict:
            # inference_mode is thread-local. Enter it inside this call, or the
            # worker thread records no autograd state saving and still builds
            # the graph.
            with torch.inference_mode():
                return app.state.systemone(model, processor, body, max_length=MAX_LEN)

        # Clef runs one prefill with shared buffers and a shared joint head.
        # Two calls at once corrupt each other, so serialize them here.
        async with app.state.lock:
            try:
                result = await asyncio.to_thread(_forward)
            except ValueError as exc:
                # A malformed request body is the caller's fault.
                return JSONResponse({"error": str(exc)}, status_code=400)
            except Exception as exc:
                LOG.error("systemone failed\n%s", traceback.format_exc())
                return JSONResponse(
                    {"error": type(exc).__name__, "detail": str(exc)},
                    status_code=500,
                )
        return result

    return app


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app:create_app", factory=True, host="0.0.0.0", port=PORT)