"""OpenAI-compatible shim for DSpark + Gemma4-12B on hiro (RTX 5090).

FP8 target (torchao float8 dynamic activation/weight) + BF16 draft head. Wraps DeepSeek DeepSpec's
DSpark speculative generate loop (Gemma4DSparkEvaluator.generate_one_sample) behind
/v1/chat/completions. Runs from ~/DeepSpec (imports `deepspec`) in its .venv.

- FP8 target ~13 GB + BF16 draft ~7 GB, plus a windowed KV cache -> 128k fits in 26.6 GB and 256k in
  28.7 GB, both fully in-VRAM on the 32 GB card (no host spilling).
- Long context needs DeepSpec's patched base_evaluator (chunked prefill + rolling hidden-state window)
  and windowed_cache.py (SpecSlidingLayer): sliding-attn layers window to 1024, so 256k KV is ~5 GB
  instead of ~90 GB. Validated: windowed cache is logit-exact vs the full forward past 1024 tokens.
- bsz=1 in the DSpark loop -> requests are serialized under a lock (fine for a single-tenant gemma).
- Subclasses the evaluator to load the FP8 target; the long-context patches live in the DeepSpec repo.
"""
import os
import time
import threading
import queue
import uuid
import io
import base64
import urllib.request

os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29500")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
# persist inductor's compiled/autotuned kernels off /tmp (which is wiped on reboot) so restarts
# reuse them instead of re-running the ~30-min max-autotune compile. fx_graph_cache is on by default.
os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", os.path.expanduser("~/DeepSpec/.inductor_cache"))
# long-context (128k-256k) defaults: window the draft's context + chunk the prefill so the whole run
# fits in the 32 GB card (128k=26.6GB, 256k=28.7GB, both in-VRAM). Implemented in DeepSpec's patched
# base_evaluator (chunked prefill + rolling hidden-state window) and windowed_cache (SpecSlidingLayer).
os.environ.setdefault("DSPARK_DRAFT_CTX_WINDOW", "16384")
os.environ.setdefault("DSPARK_PREFILL_CHUNK", "2048")

import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
# NOTE: torch._dynamo.config.caching_precompile (frontend guard cache) is INCOMPATIBLE with the
# torchao fp8 target — serializing guards calls empty_like on a Float8Tensor, which torchao doesn't
# implement (NotImplementedError). It crashes the cache save. So we can't cache the dynamo frontend;
# the ~200s warmup tracing stays. (AOTInductor would avoid dynamo entirely — separate effort.)
from types import SimpleNamespace
from typing import Any, List, Optional
import json as _json
from PIL import Image
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer, TorchAoConfig
from torchao.quantization import (
    Float8DynamicActivationFloat8WeightConfig, Float8WeightOnlyConfig, PerRow)
from torchao.quantization.quantize_.common.kernel_preference import KernelPreference
from deepspec.eval.base_evaluator import assert_no_final_target_layer, resolve_stop_token_ids
from deepspec.eval.dspark import Gemma4DSparkEvaluator

TARGET = os.environ.get("DSPARK_TARGET", "google/gemma-4-12B-it")
DRAFT = os.environ.get("DSPARK_DRAFT", "deepseek-ai/dspark_gemma4_12b_block7")
MODEL_NAME = os.environ.get("DSPARK_MODEL_NAME", "gemma4-dspark")
# gemma4 reasons in a `<|channel>thought\n ... <channel|>` channel; enable_thinking leaves it open so
# the model reasons, then emits the answer after the <channel|> token (id 101). We expose the reasoning
# as `reasoning_content` and the answer as `content`. Toggle with DSPARK_THINKING=0.
THINKING = os.environ.get("DSPARK_THINKING", "1") == "1"
_CH_CLOSE = 101  # <channel|> — separates the thought channel from the answer


class FP8Gemma4DSparkEvaluator(Gemma4DSparkEvaluator):
    """Same as the stock evaluator but loads the TARGET in torchao FP8 (draft stays BF16)."""
    def build_models(self):
        fp8 = os.environ.get("DSPARK_TARGET_DTYPE", "fp8").lower() == "fp8"
        if fp8:
            recipe = os.environ.get("DSPARK_FP8_RECIPE", "dynamic").lower()
            if recipe == "weightonly":
                # fp8 WEIGHTS (13 GB, memory-bandwidth win), bf16 activations -> no per-forward
                # activation-quant overhead (which slowed dynamic fp8 below bf16 in DSpark's small forwards)
                fp8_cfg = Float8WeightOnlyConfig()
            else:
                # dynamic activation fp8, native torch._scaled_mm path (AUTO tries a broken cutlass kernel here)
                fp8_cfg = Float8DynamicActivationFloat8WeightConfig(
                    granularity=PerRow(), kernel_preference=KernelPreference.TORCH)
            target = AutoModelForCausalLM.from_pretrained(
                self.args.target_name_or_path,
                quantization_config=TorchAoConfig(fp8_cfg),
                dtype=torch.bfloat16,
                device_map={"": self.device},
                attn_implementation=self.EVAL_ATTN_IMPLEMENTATION,
            ).eval()
        else:
            target = AutoModelForCausalLM.from_pretrained(
                self.args.target_name_or_path,
                dtype=torch.bfloat16,
                attn_implementation=self.EVAL_ATTN_IMPLEMENTATION,
            ).to(self.device).eval()
        if os.environ.get("DSPARK_COMPILE", "0") == "1":
            # let inductor fuse the fp8 activation-quant + _scaled_mm into one kernel (no cudagraphs:
            # the DSpark loop has dynamic control flow); dynamic shapes for the varying verify/prefill M
            target = torch.compile(target, mode="max-autotune-no-cudagraphs", dynamic=True)
        draft = self.draft_model_cls.from_pretrained(
            self.args.draft_name_or_path,
            dtype=torch.bfloat16,
            attn_implementation=self.EVAL_ATTN_IMPLEMENTATION,
        ).to(self.device).eval()
        assert_no_final_target_layer(target, draft.target_layer_ids)
        tokenizer = AutoTokenizer.from_pretrained(self.args.target_name_or_path)
        return target, draft, tokenizer


_args = SimpleNamespace(
    target_name_or_path=TARGET, draft_name_or_path=DRAFT,
    max_new_tokens=512, temperature=1.0, confidence_threshold=0.0,
    tensorboard_dir=None, step=None, seed=980406, tasks=[],
)
print(f"[dspark] loading target(FP8)={TARGET} + draft(BF16)={DRAFT} ...", flush=True)
_T0 = time.time()
EV = FP8Gemma4DSparkEvaluator(0, _args)
print(f"[dspark] TIMING model+quant load: {time.time()-_T0:.1f}s", flush=True)
EV.confidence_head_recorder = None  # metrics recorder is only started inside evaluate(); we bypass it
TOK = EV.tokenizer
# The multimodal processor (image + text). gemma-4-12B is a VLM; DSpark's target processes the image in
# the prefill (pixel_values), the KV cache carries it, and the draft speculates over image-aware hidden
# states — so vision rides the same speculative loop as text. Only loaded once here.
PROC = AutoProcessor.from_pretrained(TARGET)
STOP = resolve_stop_token_ids(EV.target_model, TOK)  # gemma real eos set (e.g. [1,106,50]), not a guess
_LOCK = threading.Lock()

if os.environ.get("DSPARK_COMPILE", "0") == "1":
    # torch 2.11 mega-cache: the portable "saved compiled version". Load it before compiling so the
    # warmup is a near-instant cache hit (full graph codegen + autotune), not a ~30-min recompile.
    _MEGA = os.path.expanduser("~/DeepSpec/.dspark_megacache.bin")
    if os.path.exists(_MEGA):
        try:
            _t = time.time()
            with open(_MEGA, "rb") as _f:
                torch.compiler.load_cache_artifacts(_f.read())
            print(f"[dspark] TIMING mega-cache load: {time.time()-_t:.1f}s", flush=True)
        except Exception as _e:
            print("[dspark] mega-cache load failed:", _e, flush=True)
    print("[dspark] warming up torch.compile ...", flush=True)
    for _i, _wp in enumerate(["Hello there.", "def fib(n):", "Explain gravity briefly."]):
        _enc = TOK.apply_chat_template([{"role": "user", "content": _wp}],
                                       add_generation_prompt=True, return_tensors="pt")
        _ids = (_enc["input_ids"] if hasattr(_enc, "keys") else _enc).to(EV.device)
        EV.args.max_new_tokens = 48
        _t = time.time()
        with torch.no_grad():
            EV.generate_one_sample(input_ids=_ids, stop_token_ids=STOP)
        print(f"[dspark] TIMING warmup[{_i}]: {time.time()-_t:.1f}s", flush=True)
    # persist the full compiled graph so future starts skip the compile entirely
    try:
        _art = torch.compiler.save_cache_artifacts()
        _blob = _art[0] if isinstance(_art, tuple) else _art
        if _blob:
            with open(_MEGA, "wb") as _f:
                _f.write(_blob)
            print(f"[dspark] saved compile mega-cache ({len(_blob)//1024} KB)", flush=True)
    except Exception as _e:
        print("[dspark] mega-cache save failed:", _e, flush=True)
    print("[dspark] warmup done", flush=True)

print(f"[dspark] ready. VRAM={torch.cuda.memory_allocated()/1e9:.1f}GB stop={STOP}", flush=True)

app = FastAPI()


class Msg(BaseModel):
    role: str
    content: Any  # str, or OpenAI multimodal list: [{"type":"text",...}, {"type":"image_url",...}]


def _load_image(url: str) -> Image.Image:
    """Load a PIL image from a data: URI or an http(s) URL."""
    if url.startswith("data:"):
        return Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))).convert("RGB")
    return Image.open(io.BytesIO(urllib.request.urlopen(url, timeout=20).read())).convert("RGB")


def _build_chat(messages):
    """Turn request messages into (chat-for-template, has_images). Accepts plain string content or an
    OpenAI multimodal content list (text parts + image_url/image parts), which become gemma image parts."""
    chat, has_images = [], False
    for m in messages:
        c = m.content
        if isinstance(c, str):
            chat.append({"role": m.role, "content": c})
            continue
        parts = []
        for part in (c or []):
            t = part.get("type") if isinstance(part, dict) else None
            if t == "text":
                parts.append({"type": "text", "text": part.get("text", "")})
            elif t in ("image_url", "image"):
                url = part.get("image_url", {}).get("url") if t == "image_url" else (part.get("image") or part.get("url"))
                parts.append({"type": "image", "image": _load_image(url)})
                has_images = True
            elif isinstance(part, str):
                parts.append({"type": "text", "text": part})
        chat.append({"role": m.role, "content": parts})
    return chat, has_images


class ChatReq(BaseModel):
    model: Optional[str] = None
    messages: List[Msg]
    max_tokens: Optional[int] = 512
    temperature: Optional[float] = 0.7
    stream: Optional[bool] = False


@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_NAME, "target": TARGET, "draft": DRAFT}


@app.get("/v1/models")
def models():
    return {"object": "list", "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "deepspec-dspark"}]}


def _split_think(ids, thinking):
    """Split generated token ids into (reasoning, content). With thinking on, the model emits the
    thought channel then the <channel|> token (101) then the answer; before 101 appears it's all
    reasoning. The leading `thought` label (regular token after the special <|channel>) is stripped."""
    if not thinking:
        return "", TOK.decode(ids, skip_special_tokens=True)
    if _CH_CLOSE in ids:
        i = ids.index(_CH_CLOSE)
        r = TOK.decode(ids[:i], skip_special_tokens=True)
        c = TOK.decode(ids[i + 1:], skip_special_tokens=True)
    else:
        r, c = TOK.decode(ids, skip_special_tokens=True), ""
    if r.startswith("thought"):
        r = r[len("thought"):]
    return r.strip("\n"), c


@app.post("/v1/chat/completions")
def chat(req: ChatReq):
    # thinking on by default; a `*-nothink` model alias (routed to this same backend) turns it off
    thinking = THINKING and not (req.model or "").endswith("-nothink")
    chat_msgs, has_images = _build_chat(req.messages)
    if has_images:
        # multimodal: the processor emits input_ids + pixel_values/mm_token_type_ids/image_position_ids;
        # the extra keys ride to the DSpark prefill (prefill_mm) so the target embeds the image and the
        # draft speculates over image-aware hidden states — vision on the same speculative loop as text.
        enc = PROC.apply_chat_template(
            chat_msgs, add_generation_prompt=True, tokenize=True, return_dict=True,
            return_tensors="pt", enable_thinking=thinking,
        ).to(EV.device)
        ids = enc["input_ids"]
        prefill_mm = {k: enc[k] for k in enc if k not in ("input_ids", "attention_mask")}
    else:
        ids = TOK.apply_chat_template(
            chat_msgs, add_generation_prompt=True, enable_thinking=thinking,
            return_tensors="pt", return_dict=True,
        )["input_ids"].to(EV.device)
        prefill_mm = None
    cid = "chatcmpl-" + uuid.uuid4().hex[:12]
    created = int(time.time())
    model = req.model or MODEL_NAME
    max_new = int(req.max_tokens or 512)
    if thinking:
        # gemma4 reasons deeply; keep the client's max_tokens as the ANSWER budget and add a SEPARATE,
        # GENEROUS reasoning ceiling on top, so nothing gets truncated mid-thought. This is a ceiling, not
        # a forced length — simple prompts stop early on their own, so being generous costs them nothing;
        # only genuinely hard prompts spend it. (We have the context for it — don't cap thinking small.)
        max_new += int(os.environ.get("DSPARK_THINK_BUDGET", "16384"))
    temp = max(float(req.temperature if req.temperature is not None else 1.0), 0.05)

    if req.stream:
        # Real streaming: run the (blocking) DSpark loop in a thread; its stream_callback pushes the
        # cumulative generated ids onto a queue as speculative decoding accepts them; the SSE generator
        # decodes, splits thought/answer, and emits reasoning_content + content deltas.
        q: "queue.Queue" = queue.Queue()

        def _run():
            try:
                with _LOCK:
                    EV.args.max_new_tokens = max_new
                    EV.args.temperature = temp
                    with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]), torch.no_grad():
                        EV.generate_one_sample(
                            input_ids=ids, stop_token_ids=STOP, prefill_mm=prefill_mm,
                            stream_callback=lambda t: q.put(t[0].tolist()),
                        )
            except Exception as e:  # surface generation errors to the stream instead of hanging
                q.put(("__error__", str(e)))
            finally:
                q.put(None)

        threading.Thread(target=_run, daemon=True).start()

        def _chunk(delta, finish=None):
            return "data: " + _json.dumps({
                "id": cid, "object": "chat.completion.chunk", "created": created, "model": model,
                "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
            }) + "\n\n"

        def _sse():
            yield _chunk({"role": "assistant"})
            last_r = last_c = ""
            while True:
                item = q.get()
                if item is None:
                    break
                if isinstance(item, tuple) and item[0] == "__error__":
                    yield _chunk({"content": f"\n[error: {item[1]}]"})
                    break
                r, c = _split_think(item, thinking)
                if len(r) > len(last_r):
                    yield _chunk({"reasoning_content": r[len(last_r):]})
                    last_r = r
                if len(c) > len(last_c):
                    yield _chunk({"content": c[len(last_c):]})
                    last_c = c
            yield _chunk({}, finish="stop")
            yield "data: [DONE]\n\n"

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

    with _LOCK:
        EV.args.max_new_tokens = max_new
        EV.args.temperature = temp
        t = time.time()
        with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]), torch.no_grad():
            res = EV.generate_one_sample(input_ids=ids, stop_token_ids=STOP, prefill_mm=prefill_mm)
        dt = time.time() - t
    gen = res.output_ids[0, res.num_input_tokens:].tolist()
    reasoning, content = _split_think(gen, thinking)
    al = res.acceptance_lengths
    msg = {"role": "assistant", "content": content}
    if reasoning:
        msg["reasoning_content"] = reasoning
    return {
        "id": cid, "object": "chat.completion", "created": created, "model": model,
        "choices": [{"index": 0, "message": msg, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": res.num_input_tokens, "completion_tokens": res.num_output_tokens,
                  "total_tokens": res.num_input_tokens + res.num_output_tokens},
        "dspark": {"mean_accept_len": round(sum(al) / len(al), 2) if al else 0,
                   "verify_passes": len(al), "gen_seconds": round(dt, 2),
                   "tokens_per_sec": round(res.num_output_tokens / dt, 1) if dt > 0 else 0},
    }