"""OpenAI/vLLM-compatible serving shim for the DiffusionGemma engine.

The published deployment is vLLM plus its structured decision server. This shim speaks the small part of vLLM's
HTTP API that server uses — chat/completions with ``vllm_xargs`` diffusion requests — so the unmodified upstream
server runs on top of the MLX engine. That keeps the decision layer (prompt, templates, canvas, slot reads,
adaptive thinking) bit-for-bit the published one instead of a re-implementation.

Diffusion xargs understood (as sent by runtime/structured_server.py):
  diffusion_seed_canvas   initial canvas token ids
  diffusion_canvas_length canvas width
  diffusion_max_steps     denoising steps
  diffusion_read_only     return logprobs at every canvas position, do not sample
  diffusion_pinned        canvas positions held at their seed values
  diffusion_constrained   the unembedding is read over ``logprob_token_ids`` only
  diffusion_samples       number of draws
and ``djev_bidir`` = flat [a1, b1, a2, b2, ...] prompt ranges that attend bidirectionally (patch_vllm_bidir.py).
"""

from __future__ import annotations

import argparse
import time
from typing import Any

import mlx.core as mx
import numpy as np
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel

from .engine import CANVAS_LENGTH, Engine

app = FastAPI()
ENGINE: Engine | None = None
OPTS: dict[str, Any] = {"canvas": CANVAS_LENGTH, "max_prompt": 8192, "model": "super-decider"}


class ChatRequest(BaseModel):
    model: str | None = None
    messages: list[dict] | None = None
    prompt: Any = None
    max_tokens: int | None = None
    logprobs: Any = None
    top_logprobs: int | None = None
    logprob_token_ids: list[int] | None = None
    return_tokens_as_token_ids: bool = False
    chat_template_kwargs: dict | None = None
    vllm_xargs: dict | None = None
    stop_token_ids: list[int] | None = None
    seed: int | None = None
    temperature: float | None = None
    stream: bool = False


def _tok(t: int, as_id: bool) -> str:
    return f"token_id:{t}" if as_id else ENGINE.tokenizer.decode([t])


def _block_ids(n_prompt: int, ranges) -> mx.array | None:
    """Per-token bidirectional group ids from the flat ``djev_bidir`` ranges, -1 outside them."""
    if not ranges:
        return None
    ids = np.full(n_prompt, -1, dtype=np.int32)
    for k, i in enumerate(range(0, len(ranges) - 1, 2)):
        a, b = int(ranges[i]), int(ranges[i + 1])
        ids[max(a, 0) : min(b, n_prompt)] = k
    return mx.array(ids)


def _usage(prompt_tokens: int, completion_tokens: int) -> dict:
    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": prompt_tokens + completion_tokens,
    }


def _read(req: ChatRequest, prompt_ids: list[int]):
    """A diffusion read: label logits at every canvas position, plus the canvas tokens themselves."""
    x = req.vllm_xargs or {}
    canvas = [int(t) for t in x.get("diffusion_seed_canvas") or []]
    canvas_len = int(x.get("diffusion_canvas_length") or OPTS["canvas"])
    steps = int(x.get("diffusion_max_steps") or 1)
    pinned_pos = [int(p) for p in x.get("diffusion_pinned") or []]
    samples = int(x.get("diffusion_samples") or 1)
    label_ids = [int(t) for t in (req.logprob_token_ids or [])]
    if not canvas:
        raise ValueError("a diffusion request needs diffusion_seed_canvas")
    canvas = (canvas + [0] * canvas_len)[:canvas_len]
    block_ids = _block_ids(len(prompt_ids), x.get("djev_bidir"))
    pinned = {p: canvas[p] for p in pinned_pos if 0 <= p < canvas_len}
    probs = ENGINE.read(
        prompt_ids,
        canvas,
        list(range(canvas_len)),
        label_ids,
        block_ids=block_ids,
        steps=steps,
        pinned=pinned,
        seed=int(req.seed or 0),
        samples=samples,
    )
    return canvas, label_ids, probs


def _top_rows(probs_row: np.ndarray, label_ids, as_id: bool, k: int):
    order = np.argsort(-probs_row)[: max(k, 1)]
    return [
        {"token": _tok(label_ids[i], as_id), "logprob": float(probs_row[i]), "bytes": None}
        for i in order
    ]


def _read_content(canvas, label_ids, probs, as_id: bool, top_k: int):
    """One logprobs entry per canvas position (``logprobs.content`` / ``top_logprobs``)."""
    rows = []
    for pos, tok in enumerate(canvas):
        tops = _top_rows(probs[pos], label_ids, as_id, top_k) if label_ids else []
        best = tops[0]["logprob"] if tops else 0.0
        rows.append(
            {"token": _tok(tok, as_id), "logprob": best, "bytes": None, "top_logprobs": tops}
        )
    return rows


def _token_logprobs(tokens: list[int], lps, as_id: bool, chat: bool):
    """vLLM's legacy per-token logprobs object (``content`` for chat, ``tokens``/``token_logprobs`` for completions).

    ``top_logprobs`` stays empty: the model reads a label set, not a token ranking, so nothing here asks for it.
    """
    if lps is None:
        return None
    vals = [float(v) for v in lps]
    if chat:
        return {
            "content": [
                {"token": _tok(t, as_id), "logprob": v, "bytes": None, "top_logprobs": []}
                for t, v in zip(tokens, vals)
            ],
            "refusal": None,
        }
    return {
        "tokens": [_tok(t, as_id) for t in tokens],
        "token_logprobs": vals,
        "top_logprobs": [{} for _ in tokens],
        "text_offset": list(range(len(tokens))),
    }


def _generate(req: ChatRequest, prompt_ids: list[int]):
    x = req.vllm_xargs or {}
    want = req.logprobs is not None
    out = ENGINE.generate(
        prompt_ids,
        max_new_tokens=req.max_tokens,
        max_denoising_steps=int(x.get("diffusion_max_steps") or 0) or None,
        stop_ids=req.stop_token_ids,
        block_ids=_block_ids(len(prompt_ids), x.get("djev_bidir")),
        token_logprobs=want,
    )
    gen, lps = out if want else (out, None)
    if req.stop_token_ids:
        for k, t in enumerate(gen):
            if t in req.stop_token_ids:
                gen = gen[: k + 1]
                break
    if req.max_tokens:
        gen, lps = gen[: req.max_tokens], (lps[: req.max_tokens] if lps else lps)
    return gen, lps


@app.get("/v1/models")
@app.get("/health")
def models():
    return {"object": "list", "data": [{"id": OPTS["model"], "object": "model", "owned_by": "local"}]}


def _handle(req: ChatRequest, chat: bool):
    if chat:
        prompt_ids = ENGINE.chat_ids(
            req.messages or [],
            bool((req.chat_template_kwargs or {}).get("enable_thinking", False)),
        )
    else:
        prompt_ids = [int(t) for t in (req.prompt or [])]
    if len(prompt_ids) > OPTS["max_prompt"]:
        # A real 400: vLLM answers 400 here, and a 200 body without choices makes the
        # decision layer raise KeyError('choices') instead of reporting the length.
        return JSONResponse(
            status_code=400,
            content={
                "error": {
                    "code": 400,
                    "message": f"state too long: {len(prompt_ids)} tokens, this build serves "
                    f"{OPTS['max_prompt']} (raise --max-prompt; attention cost grows quadratically)",
                    "type": "BadRequestError",
                }
            },
        )
    as_id = bool(req.return_tokens_as_token_ids)
    diffusion = bool((req.vllm_xargs or {}).get("diffusion_seed_canvas"))
    created = int(time.time())
    if diffusion:
        canvas, label_ids, probs = _read(req, prompt_ids)
        # One choice per draw, the way vLLM answers a multi-sample request.
        choices = []
        for k in range(probs.shape[0]):
            content = _read_content(canvas, label_ids, probs[k], as_id, int(req.top_logprobs or 20))
            choice = {
                "index": k,
                "text": "",
                "finish_reason": "stop",
                "logprobs": {
                    "tokens": [c["token"] for c in content],
                    "token_logprobs": [c["logprob"] for c in content],
                    "top_logprobs": [{t["token"]: t["logprob"] for t in c["top_logprobs"]} for c in content],
                    "text_offset": list(range(len(content))),
                },
            }
            if chat:
                choice["message"] = {"role": "assistant", "content": ""}
                choice["logprobs"] = {"content": content, "refusal": None}
            choices.append(choice)
        body = {
            "id": "djev-local",
            "object": "chat.completion" if chat else "text_completion",
            "created": created,
            "model": OPTS["model"],
            "choices": choices,
            "usage": _usage(len(prompt_ids), len(canvas)),
            "system_fingerprint": "super-decider-metal",
        }
        return body

    gen, lps = _generate(req, prompt_ids)
    text = ENGINE.tokenizer.decode(gen)
    choice = {"index": 0, "text": text, "finish_reason": "stop"}
    if chat:
        choice["message"] = {"role": "assistant", "content": text}
    choice["logprobs"] = _token_logprobs(gen, lps, as_id, chat)
    body = {
        "id": "djev-local",
        "object": "chat.completion" if chat else "text_completion",
        "created": created,
        "model": OPTS["model"],
        "choices": [choice],
        "usage": _usage(len(prompt_ids), len(gen)),
        "system_fingerprint": "super-decider-metal",
    }
    return body


@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest):
    return _handle(req, chat=True)


@app.post("/v1/completions")
async def completions(req: ChatRequest):
    return _handle(req, chat=False)


def main():
    global ENGINE
    ap = argparse.ArgumentParser(description="super-decider-metal engine server (vLLM-compatible subset)")
    ap.add_argument("--model", default="artifacts/mlx-nvfp4")
    ap.add_argument("--name", default="super-decider")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8000)
    ap.add_argument("--canvas", type=int, default=CANVAS_LENGTH)
    ap.add_argument("--max-prompt", type=int, default=8192)
    args = ap.parse_args()
    OPTS.update(canvas=args.canvas, max_prompt=args.max_prompt, model=args.name)
    ENGINE = Engine(args.model, canvas_length=args.canvas)
    import uvicorn

    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")


if __name__ == "__main__":
    main()