"""Minimal OpenAI-compatible chat server on top of the DeepInfer offline engine.

DeepInfer ships offline inference only. This wraps one `LLM` per data-parallel rank
(attention DP + MoE EP, same layout as examples/benchmark.py) behind a FastAPI frontend
so request-replay benchmarks such as InferenceX AgentX (aiperf) can drive it.

    python serve/openai_server.py --model ~/models/DeepSeek-V4.1-Flash-1017 --dp 4 --port 8888

Limitations, by design of the open-source engine:
- No cross-request prefix cache: every request prefills its full prompt.
- Responses are streamed as plain `content`; DSML tool calls are not parsed into `tool_calls`.
"""

from __future__ import annotations

import argparse
import asyncio
import collections
import hashlib
import json
import multiprocessing as mp
import os
import queue
import sys
import threading
import time
import uuid
from pathlib import Path
from typing import TYPE_CHECKING, Any

# Module level so FastAPI can resolve the string annotation `request: Request`.
from fastapi import FastAPI, Request

if TYPE_CHECKING:
    from deepinfer._C.engine import Sequence

REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "examples"))


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True, help="local HF checkpoint directory")
    p.add_argument("--served-model-name", default=None)
    p.add_argument("--dp", type=int, default=4, help="data-parallel ranks on this node (= EP size)")
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8888)
    p.add_argument("--max-model-len", type=int, default=1 << 20)
    p.add_argument("--max-num-seqs", type=int, default=128, help="per DP rank")
    p.add_argument("--max-num-batched-tokens", type=int, default=32768, help="prefill chunk per DP rank")
    p.add_argument("--gpu-memory-utilization", type=float, default=0.9)
    p.add_argument("--yoco-keep-tokens", type=int, default=140, help="0 disables decoder SWA bounded replay")
    p.add_argument("--thinking-mode", default="chat", choices=["chat", "thinking"])
    p.add_argument("--log-interval-steps", type=int, default=200)
    p.add_argument("--hicache-size-gb", type=float, default=32, help="pinned host KV page pool per DP rank; 0 disables the prefix cache")
    p.add_argument("--hicache-state-slots", type=int, default=4096, help="host SWA state slots per DP rank")
    p.add_argument(
        "--routing",
        default="affinity",
        choices=["affinity", "least-loaded"],
        help="affinity: a conversation's turns go to one rank (needed for cache hits)",
    )
    p.add_argument(
        "--prefill-any-interval",
        type=int,
        default=2,
        help="DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL: a rank with a waiting prompt gets a prefill step at most this "
        "many decode steps later even if other ranks are idle (engine default 100 => ~1.4 s TTFT floor)",
    )
    p.add_argument(
        "--eplb-interval",
        type=int,
        default=0,
        help="DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL: rearrange experts every N steps; 0 (default) disables EPLB. "
        "Each rearrangement stalls every rank (~0.3 s on B200/GB300), which under online traffic costs far more than "
        "the balance it buys",
    )
    p.add_argument(
        "--eplb-after-prefill",
        type=int,
        default=0,
        choices=[0, 1],
        help="DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE: also rearrange after every prefill step (only with "
        "--eplb-interval > 0); with agentic traffic nearly every step batch contains a prefill",
    )
    return p.parse_args()


# --------------------------------------------------------------------------- engine side


def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: mp.Queue) -> None:
    """One DP rank: owns an engine and steps it in lockstep with the other ranks."""
    import ctypes
    import signal

    # Die with the frontend instead of lingering as an orphan that holds the GPU.
    ctypes.CDLL("libc.so.6").prctl(1, signal.SIGKILL)  # PR_SET_PDEATHSIG

    from _config import create_llm, setup_common_env

    if args.yoco_keep_tokens > 0:
        os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact"
        os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = str(args.yoco_keep_tokens)
        os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1"
    if args.hicache_size_gb > 0:
        # Keep more SWA history in the state ring so the prefix cache can store states at several
        # 128-boundaries behind the end of a prompt (see HostPrefixCache.state_lookback_blocks).
        os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512")
    os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval))
    setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False)
    if args.eplb_interval > 0:
        os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval)
        os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill)

    llm = create_llm(
        args.model,
        dp_rank,
        ep_size=args.dp,
        env={"DEEPINFER_LOGGING_INTERVAL_STEP": str(args.log_interval_steps)},
        data_parallel_size=args.dp,
        max_num_seqs=args.max_num_seqs,
        max_num_batched_tokens=args.max_num_batched_tokens,
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
        lookahead_tokens=min(1, int(os.getenv("DEEPINFER_USE_NEXTN_DECODE", "5"))),
    )

    import numpy as np

    from deepinfer._C.engine import RequestOutput  # pyright: ignore[reportMissingModuleSource]
    from deepinfer.sampling_params import SamplingParams

    engine = llm.llm_engine
    cache = None
    if args.hicache_size_gb > 0:
        try:
            from deepinfer.hicache import HostPrefixCache
        except ImportError:  # tree without the host prefix cache: serve without it
            print(f"[server] deepinfer.hicache not available, prefix cache disabled (rank {dp_rank})", flush=True)
        else:
            cache = HostPrefixCache(engine, size_bytes=int(args.hicache_size_gb * 2**30), num_state_slots=args.hicache_state_slots)
            engine.pe_do_save_kvcache = cache.on_prefill_step  # write-through at the end of every prefill
    out_q.put(("ready", dp_rank, None))

    def mm_data():
        return dict(image_bytes=[], image_sizes=[]) if engine.multi_modal_processor is not None else None

    def enqueue(rid: str, seq: Sequence, num_cached_tokens: int) -> None:
        # Same as LLMEngine.add_request's _enqueue: the processor must get the sequence's own multimodal dict
        # (overwrite_prompt wrote the text/image segments into it), plus the cached prefix length.
        engine.scheduler.add_sequence(seq)
        if engine.multi_modal_processor is not None:
            engine.multi_modal_processor.add(rid, seq.multi_modal_data, num_cached_tokens)

    pending = collections.deque()  # (rid, seq, prompt_ids, use_cache): created, waiting for a cache decision
    inflight = False
    last_stats = time.monotonic()
    last_logged_lookups = 0
    while True:
        # Requests must be enqueued before the DP-wide has_unfinished_requests all-reduce,
        # so every rank takes the same step / flush / idle branch in the same iteration.
        while True:
            try:
                kind, rid, payload = in_q.get_nowait()
            except queue.Empty:
                break
            if kind == "add":
                try:
                    prompt_ids = np.asarray(payload["prompt_token_ids"], dtype=np.int32)
                    seq = engine.add_request(
                        user_id="0",
                        request_id=rid,
                        sampling_params=SamplingParams(**payload["sampling"]),
                        prompt_token_ids=prompt_ids,
                        multi_modal_data=mm_data(),
                        enqueue=cache is None,
                    )
                    if cache is not None:
                        if seq is None:
                            out_q.put(("error", rid, "request dropped: invalid prompt tokens"))
                        else:
                            pending.append((rid, seq, prompt_ids, payload.get("use_cache", True)))
                except Exception as e:  # noqa: BLE001 - surface to the client, keep the rank alive
                    out_q.put(("error", rid, repr(e)))
            elif kind == "abort":
                if any(p[0] == rid for p in pending):
                    pending = collections.deque(p for p in pending if p[0] != rid)
                else:
                    engine.abort_request(rid)
                if cache is not None:
                    cache.forget(rid)

        # Cache decisions, in arrival order. A hit pre-allocates the sequence's blocks and fills them from the
        # host; if the GPU cannot allocate yet, wait (while something is running that will free memory).
        while pending:
            rid, seq, prompt_ids, use_cache = pending[0]
            hit = cache.try_resume(seq, rid, prompt_ids) if use_cache else 0
            if hit < 0 and (engine.scheduler.get_num_unfinished_seqs() > 0 or inflight):
                break
            pending.popleft()
            enqueue(rid, seq, max(hit, 0))
        if cache is not None:
            cache.poll()
            if time.monotonic() - last_stats > 5:
                last_stats = time.monotonic()
                st = cache.stats_dict()
                out_q.put(("stats", dp_rank, st))
                if st["lookups"] != last_logged_lookups:
                    last_logged_lookups = st["lookups"]
                    print(
                        f"[hicache rank {dp_rank}] lookups {st['lookups']} hits {st['hits']} "
                        f"token hit rate {st['hit_rate_tokens']:.1%} deferred {st['deferred']} stores {st['stores']} "
                        f"pages {st['pages_used']}/{st['pages_total']} states {st['states_used']}/{st['states_total']} "
                        f"evicted {st['evicted_pages']}",
                        flush=True,
                    )

        if engine.has_unfinished_requests():
            outputs = engine.step()
            inflight = True
        elif inflight:
            outputs = engine.flush_inflight_reqs()
            inflight = False
        else:
            time.sleep(0.001)
            continue

        if len(outputs):
            # Pointers are recycled on the next step, so read the fields now.
            batch = [
                (
                    RequestOutput.get_request_id(p),
                    RequestOutput.get_token_id(p),
                    RequestOutput.get_finished(p),
                    RequestOutput.get_finish_reason(p),
                    RequestOutput.get_prompt_len(p),
                )
                for p in outputs
            ]
            if cache is not None:
                for rid, _, finished, _, _ in batch:
                    if finished:
                        cache.forget(rid)
            out_q.put(("tokens", None, batch))


# --------------------------------------------------------------------------- HTTP side


class IncrementalDetokenizer:
    """vLLM-style prefix/read offsets so each delta decodes only a short tail."""

    def __init__(self, tokenizer: Any) -> None:
        self.tok = tokenizer
        self.ids: list[int] = []
        self.num_tokens = 0
        self.prefix_offset = 0
        self.read_offset = 0

    def add(self, token_id: int) -> str:
        self.num_tokens += 1
        if token_id == self.tok.eos_token_id:
            return ""
        self.ids.append(token_id)
        prefix = self.tok.decode(self.ids[self.prefix_offset : self.read_offset], skip_special_tokens=False)
        full = self.tok.decode(self.ids[self.prefix_offset :], skip_special_tokens=False)
        if len(full) > len(prefix) and not full.endswith("�"):
            self.prefix_offset, self.read_offset = self.read_offset, len(self.ids)
            return full[len(prefix) :]
        return ""


class Frontend:
    def __init__(self, args: argparse.Namespace) -> None:
        from transformers import AutoTokenizer

        sys.path.insert(0, str(Path(args.model) / "encoding"))
        from encoding import encode_messages  # the checkpoint's own prompt-format reference

        self.args = args
        self.model_name = args.served_model_name or args.model
        self.encode_messages = encode_messages
        self.tokenizer = AutoTokenizer.from_pretrained(args.model)
        ctx = mp.get_context("spawn")
        self.out_q: mp.Queue = ctx.Queue()
        self.in_qs: list[mp.Queue] = [ctx.Queue() for _ in range(args.dp)]
        # Not daemonic: the engine forks its own helper processes (as examples/_launch.py allows).
        self.procs = [ctx.Process(target=worker_main, args=(args, r, self.in_qs[r], self.out_q), daemon=False) for r in range(args.dp)]
        self.load = [0] * args.dp
        self.owner: dict[str, int] = {}
        self.streams: dict[str, asyncio.Queue] = {}
        self.ready = 0
        self.stats = dict(requests=0, finished=0, aborted=0, errors=0, prompt_tokens=0, completion_tokens=0, encode_fallbacks=0)
        self.cache_stats: dict[int, dict] = {}
        self.loop: asyncio.AbstractEventLoop | None = None

    def start(self) -> None:
        for p in self.procs:
            p.start()

    def dispatch_forever(self) -> None:
        """Thread: fan engine outputs out to the per-request asyncio queues."""
        while True:
            try:
                kind, rid, payload = self.out_q.get(timeout=1.0)
            except queue.Empty:
                dead = [r for r, p in enumerate(self.procs) if not p.is_alive()]
                if dead:
                    print(f"[server] engine rank(s) {dead} died; failing open requests and exiting", flush=True)
                    for rid in list(self.streams):
                        self.loop.call_soon_threadsafe(self._deliver, rid, ("error", f"engine rank died: {dead}"))
                    time.sleep(2)
                    os._exit(1)
                continue
            if kind == "ready":
                self.ready += 1
                print(f"[server] DP rank {rid} ready ({self.ready}/{self.args.dp})", flush=True)
            elif kind == "error":
                self.loop.call_soon_threadsafe(self._deliver, rid, ("error", payload))
            elif kind == "stats":
                self.cache_stats[rid] = payload
            else:
                self.loop.call_soon_threadsafe(self._deliver_batch, payload)

    def _deliver(self, rid: str, item: tuple) -> None:
        q = self.streams.get(rid)
        if q is not None:
            q.put_nowait(item)

    def _deliver_batch(self, batch: list[tuple]) -> None:
        for rid, token_id, finished, reason, prompt_len in batch:
            q = self.streams.get(rid)
            if q is not None:
                q.put_nowait(("token", token_id, finished, reason, prompt_len))

    def render_prompt(self, body: dict) -> list[int]:
        messages = [dict(m) for m in body["messages"]]
        if body.get("tools"):
            if not messages or messages[0].get("role") != "system":
                messages.insert(0, {"role": "system", "content": ""})
            messages[0]["tools"] = body["tools"]
        try:
            prompt = self.encode_messages(messages, thinking_mode=self.args.thinking_mode)
        except Exception:  # noqa: BLE001 - trace messages the reference encoder rejects
            # Keep the request (and its prompt length) instead of failing the benchmark.
            self.stats["encode_fallbacks"] += 1
            parts = ["<｜begin▁of▁sentence｜>"]
            for m in messages:
                content = m.get("content") or ""
                if isinstance(content, list):
                    content = "".join(b.get("text", "") for b in content if isinstance(b, dict))
                if m.get("tool_calls"):
                    content += json.dumps(m["tool_calls"], ensure_ascii=False)
                tag = {"user": "<｜User｜>", "assistant": "<｜Assistant｜>", "tool": "<｜User｜>"}.get(m.get("role"), "")
                parts.append(tag + content)
            if body.get("tools"):
                parts.insert(1, json.dumps(body["tools"], ensure_ascii=False))
            prompt = "".join(parts) + "<｜Assistant｜></think>"
        return self.tokenizer.encode(prompt, add_special_tokens=False)

    @staticmethod
    def affinity_key(body: dict) -> bytes | None:
        """Stable per-conversation key: the first user message (AgentX's cache-bust marker lives there and is
        constant across a trajectory's turns, and forked subagents inherit it)."""
        for m in body.get("messages", []):
            if m.get("role") == "user":
                content = m.get("content") or ""
                if isinstance(content, list):
                    content = "".join(b.get("text", "") for b in content if isinstance(b, dict))
                return hashlib.sha1(content[:4096].encode("utf-8", "ignore")).digest()
        return None

    def submit(self, prompt_ids: list[int], sampling: dict, key: bytes | None = None, use_cache: bool = True) -> tuple[str, asyncio.Queue]:
        rid = uuid.uuid4().hex
        if self.args.routing == "affinity" and key is not None:
            rank = int.from_bytes(key[:4], "little") % self.args.dp
        else:
            rank = min(range(self.args.dp), key=lambda r: self.load[r])
        self.load[rank] += 1
        self.owner[rid] = rank
        q: asyncio.Queue = asyncio.Queue()
        self.streams[rid] = q
        self.stats["requests"] += 1
        self.in_qs[rank].put(("add", rid, {"prompt_token_ids": prompt_ids, "sampling": sampling, "use_cache": use_cache}))
        return rid, q

    def release(self, rid: str, aborted: bool) -> None:
        rank = self.owner.pop(rid, None)
        self.streams.pop(rid, None)
        if rank is None:
            return
        self.load[rank] -= 1
        if aborted:
            self.stats["aborted"] += 1
            self.in_qs[rank].put(("abort", rid, None))
        else:
            self.stats["finished"] += 1


def build_app(fe: Frontend) -> FastAPI:
    from fastapi import FastAPI
    from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse

    app = FastAPI()

    @app.on_event("startup")
    async def _startup() -> None:
        fe.loop = asyncio.get_running_loop()
        threading.Thread(target=fe.dispatch_forever, daemon=True).start()

    @app.get("/health")
    async def health():
        ok = fe.ready == fe.args.dp
        return JSONResponse({"status": "ok" if ok else "loading", "ready_ranks": fe.ready}, status_code=200 if ok else 503)

    @app.get("/v1/models")
    async def models():
        return {"object": "list", "data": [{"id": fe.model_name, "object": "model", "owned_by": "deepinfer"}]}

    # Not at /metrics: aiperf scrapes that path automatically and InferenceX's aggregator
    # rejects metric families it has no backend adapter for.
    @app.get("/deepinfer/metrics")
    async def metrics():
        lines = [f"deepinfer:{k}_total {v}" for k, v in fe.stats.items()]
        lines += [f'deepinfer:num_requests_running{{dp_rank="{r}"}} {n}' for r, n in enumerate(fe.load)]
        for r, st in sorted(fe.cache_stats.items()):
            lines += [f'deepinfer_hicache:{k}{{dp_rank="{r}"}} {v}' for k, v in st.items()]
        return PlainTextResponse("\n".join(lines) + "\n")

    @app.post("/v1/chat/completions")
    async def chat(request: Request):
        if fe.ready != fe.args.dp:
            return JSONResponse({"error": {"message": "engine loading"}}, status_code=503)
        body = await request.json()
        prompt_ids = await asyncio.to_thread(fe.render_prompt, body)
        max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") or 4096
        if len(prompt_ids) + 1 > fe.args.max_model_len:
            return JSONResponse({"error": {"message": f"prompt has {len(prompt_ids)} tokens > max_model_len"}}, status_code=400)
        max_tokens = max(1, min(int(max_tokens), fe.args.max_model_len - len(prompt_ids)))
        sampling = dict(
            temperature=float(body.get("temperature", 1.0)),
            top_p=float(body.get("top_p", 1.0)),
            ignore_eos=bool(body.get("ignore_eos", False)),
            max_tokens=max_tokens,
        )
        if body.get("seed") is not None:
            sampling["seq_seed"] = int(body["seed"])
        fe.stats["prompt_tokens"] += len(prompt_ids)
        include_usage = bool((body.get("stream_options") or {}).get("include_usage"))
        # Non-standard request field: "prefix_cache": false bypasses the host prefix cache lookup (A/B checks).
        rid, q = fe.submit(prompt_ids, sampling, fe.affinity_key(body), use_cache=bool(body.get("prefix_cache", True)))
        created = int(time.time())
        cid = f"chatcmpl-{rid}"
        detok = IncrementalDetokenizer(fe.tokenizer)

        def chunk(delta: dict, finish: str | None = None, usage: dict | None = None) -> str:
            obj = {
                "id": cid,
                "object": "chat.completion.chunk",
                "created": created,
                "model": fe.model_name,
                "choices": [] if usage else [{"index": 0, "delta": delta, "finish_reason": finish}],
            }
            if usage:
                obj["usage"] = usage
            return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"

        async def tokens():
            """Yield (text_delta, finish_reason) until the request finishes."""
            while True:
                item = await q.get()
                if item[0] == "error":
                    raise RuntimeError(item[1])
                _, token_id, finished, reason, _ = item
                text = detok.add(token_id)
                fe.stats["completion_tokens"] += 1
                if finished:
                    yield text, ("length" if detok.num_tokens >= max_tokens else "stop")
                    return
                yield text, None

        usage = lambda: {
            "prompt_tokens": len(prompt_ids),
            "completion_tokens": detok.num_tokens,  # noqa: E731
            "total_tokens": len(prompt_ids) + detok.num_tokens,
        }

        if body.get("stream"):

            async def sse():
                done = False
                try:
                    yield chunk({"role": "assistant", "content": ""})
                    async for text, finish in tokens():  # noqa: B007  (the last finish reason is reported below)
                        if text or finish:
                            yield chunk({"content": text} if text else {}, finish)
                    done = True
                    if include_usage:
                        yield chunk({}, usage=usage())
                    yield "data: [DONE]\n\n"
                finally:
                    fe.release(rid, aborted=not done)

            return StreamingResponse(sse(), media_type="text/event-stream")

        done = False
        try:
            parts, finish = [], None
            async for text, finish in tokens():  # noqa: B007  (the last finish reason is reported below)
                parts.append(text)
            done = True
        finally:
            fe.release(rid, aborted=not done)
        return {
            "id": cid,
            "object": "chat.completion",
            "created": created,
            "model": fe.model_name,
            "choices": [{"index": 0, "message": {"role": "assistant", "content": "".join(parts)}, "finish_reason": finish}],
            "usage": usage(),
        }

    return app


def main() -> None:
    args = parse_args()
    args.model = str(Path(args.model).expanduser().resolve())
    os.environ.setdefault("WORLD_SIZE", "1")
    os.environ.setdefault("RANK", "0")
    fe = Frontend(args)
    fe.start()

    import uvicorn

    uvicorn.run(build_app(fe), host=args.host, port=args.port, log_level="warning", timeout_keep_alive=900)


if __name__ == "__main__":
    main()