"""Same-model baseline: the base Qwen3.5-0.8B (the instruct/chat release deg starts from,
NOT deg's weights) writes the whole JSON object under xgrammar JSON-schema-constrained
greedy decoding (HF transformers generate + xgrammar logits processor), batch 1.

Prompt: chat template, thinking off, :data:`deg.bench.systems.INSTRUCTION` as the system
message, schema + context as the user message. The grammar is compiled from the schema
(``x-generate`` dropped as an annotation) with standard ``", "`` / ``": "`` separators and no
other whitespace; ``compact=True`` uses ``","`` / ``":"`` (fewer tokens, but the 0.8B model
then often writes a lone ``","`` as a string value: an ablation, not the main arm). Grammar compilation is timed separately (``extra.compile_s``) and is not in
``latency_s``; generated tokens count every emitted token except the stop token.
"""

from __future__ import annotations

import copy
import json
import time
from pathlib import Path
from typing import Any

import torch

from deg.bench.score import Prediction
from deg.bench.systems import INSTRUCTION, System, user_message

MAX_NEW_TOKENS = 1536


def grammar_schema(schema: Any, drop: tuple[str, ...] = ("x-generate",)) -> Any:
    def walk(node: Any) -> Any:
        if isinstance(node, dict):
            return {k: walk(v) for k, v in node.items() if k not in drop}
        if isinstance(node, list):
            return [walk(v) for v in node]
        return node
    return walk(copy.deepcopy(schema))


def stop_ids(tok: Any, model_dir: Path) -> list[int]:
    """End-of-turn ids: Qwen's ``<|im_end|>`` when the vocab has it, the tokenizer eos, and
    every ``eos_token_id`` in ``generation_config.json`` (Gemma / LFM / MiniCPM turn ends)."""
    stop: list[int] = []
    im_end = tok.convert_tokens_to_ids("<|im_end|>")
    if isinstance(im_end, int) and im_end != tok.unk_token_id:
        stop.append(im_end)
    cands: list[Any] = [tok.eos_token_id]
    gc = model_dir / "generation_config.json"
    if gc.exists():
        e = json.loads(gc.read_text()).get("eos_token_id")
        cands += e if isinstance(e, list) else [e]
    for t in cands:
        if isinstance(t, int) and t not in stop:
            stop.append(t)
    return stop


def native_context(model_dir: Path) -> int | None:
    """The model's own maximum context (``max_position_embeddings``, text config first)."""
    cfg = json.loads((model_dir / "config.json").read_text())
    for c in (cfg.get("text_config") or {}, cfg):
        if c.get("max_position_embeddings"):
            return int(c["max_position_embeddings"])
    return None


def chat_prompt(tok: Any, rec: dict[str, Any]) -> str:
    """The model's own chat template, thinking off where the template supports it; a
    template without a system role gets the instruction prepended to the user turn."""
    user = user_message(rec["context"], rec["schema"])
    try:
        return tok.apply_chat_template(
            [{"role": "system", "content": INSTRUCTION}, {"role": "user", "content": user}],
            tokenize=False, add_generation_prompt=True, enable_thinking=False)
    except Exception:  # noqa: BLE001 - template rejects a system turn
        return tok.apply_chat_template(
            [{"role": "user", "content": INSTRUCTION + "\n\n" + user}],
            tokenize=False, add_generation_prompt=True, enable_thinking=False)


class ConstrainedSystem(System):
    def __init__(self, model_path: str | None, label: str | None = None,
                 compact: bool = False) -> None:
        import xgrammar as xgr
        from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen3_5ForCausalLM

        from deg.model import base_snapshot

        path = Path(model_path) if model_path else base_snapshot()
        self.tok = AutoTokenizer.from_pretrained(str(path))
        model_type = json.loads((path / "config.json").read_text()).get("model_type", "")
        cls = Qwen3_5ForCausalLM if model_type.startswith("qwen3_5") else AutoModelForCausalLM
        self.lm = cls.from_pretrained(str(path), dtype=torch.bfloat16)
        self.lm.to("cuda").eval()
        vocab = self.lm.get_output_embeddings().weight.shape[0]
        self.stop = stop_ids(self.tok, path)
        self.native_ctx = native_context(path)
        info = xgr.TokenizerInfo.from_huggingface(self.tok, vocab_size=vocab,
                                                  stop_token_ids=self.stop)
        self.compiler = xgr.GrammarCompiler(info, max_threads=8)
        self.xgr = xgr
        self.separators = (",", ":") if compact else (", ", ": ")
        self.name = label or ("qwen3.5-0.8b-xgrammar" + ("-compact" if compact else ""))
        self.info = {"model": str(path),
                     "weights": ("base Qwen3.5-0.8B (chat release), not deg"
                                 if model_path is None else f"{model_type} (as released)"),
                     "decoding": "greedy, xgrammar JSON-schema constrained",
                     "separators": list(self.separators), "native_context": self.native_ctx,
                     "dtype": "bf16", "batch": 1, "runtime": "HF transformers generate"}
        self._warm()

    def _warm(self) -> None:
        rec = {"uid": "warm", "context": "Order AB-1 arrived late.",
               "schema": {"type": "object", "properties": {"id": {"type": "string"}},
                          "required": ["id"]}}
        self.predict(rec)

    def prompt_ids(self, rec: dict[str, Any]) -> list[int]:
        return self.tok(chat_prompt(self.tok, rec), add_special_tokens=False)["input_ids"]

    def compile(self, schema: Any) -> tuple[Any, str]:
        clean = grammar_schema(schema)
        try:
            return self.compiler.compile_json_schema(
                json.dumps(clean), any_whitespace=False, separators=self.separators), "full"
        except Exception:  # noqa: BLE001 - fall back to a schema without string formats
            loose = grammar_schema(clean, drop=("format", "pattern"))
            return self.compiler.compile_json_schema(
                json.dumps(loose), any_whitespace=False, separators=self.separators), "no-format"

    @torch.no_grad()
    def predict(self, rec: dict[str, Any]) -> Prediction:
        from xgrammar.contrib.hf import LogitsProcessor

        t0 = time.perf_counter()
        try:
            grammar, grammar_mode = self.compile(rec["schema"])
        except Exception as exc:  # noqa: BLE001 - record and score as a failure
            return Prediction(rec["uid"], error=f"grammar: {str(exc)[:200]}")
        compile_s = time.perf_counter() - t0
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        ids = self.prompt_ids(rec)
        if self.native_ctx and len(ids) >= self.native_ctx:
            # deterministic: the document does not fit the model's own context; scored 0
            return Prediction(rec["uid"], error=f"context_exceeded: native {len(ids)} prompt "
                                                f"tokens >= {self.native_ctx}",
                              prompt_tokens=len(ids))
        inp = torch.tensor([ids], device="cuda")
        out = self.lm.generate(inp, attention_mask=torch.ones_like(inp),
                               logits_processor=[LogitsProcessor(grammar)],
                               max_new_tokens=MAX_NEW_TOKENS, do_sample=False,
                               eos_token_id=self.stop, pad_token_id=self.stop[0])
        torch.cuda.synchronize()
        dt = time.perf_counter() - t1
        new = out[0, len(ids):].tolist()
        while new and new[-1] in self.stop:
            new.pop()
        text = self.tok.decode(new, skip_special_tokens=True)
        extra = {"compile_s": round(compile_s, 4), "grammar": grammar_mode}
        try:
            value = json.loads(text)
        except json.JSONDecodeError as exc:
            hit_max = len(new) >= MAX_NEW_TOKENS - 1
            return Prediction(rec["uid"], error=f"json: {exc}"[:200] + (" (max tokens)"
                                                                       if hit_max else ""),
                              generated_tokens=len(new), prompt_tokens=len(ids), latency_s=dt,
                              extra={**extra, "text": text[-300:]})
        return Prediction(rec["uid"], value=value, generated_tokens=len(new),
                          prompt_tokens=len(ids), latency_s=dt, extra=extra)


class LlamaCppConstrainedSystem(System):
    """The same baseline on llama.cpp: a GGUF model on a llama.cpp server (CUDA), the
    model's own chat template (from its HF snapshot, ``model_path``; default the base
    Qwen3.5-0.8B), greedy, with the server's JSON-schema -> GBNF grammar.

    ``max_new_tokens`` and :meth:`prompt` are hooks for subclasses (the extraction
    specialists in :mod:`deg.bench.systems.extractors`).

    Context: a prompt longer than the model's native context is ``context_exceeded:
    native``; one that fits the model but not the server's ``n_ctx`` (sized to the GPU
    memory left beside Clef) is ``context_exceeded: server`` (a harness limit, reported
    apart). Both score 0."""

    max_new_tokens = MAX_NEW_TOKENS

    def __init__(self, url: str | None, label: str | None = None,
                 model_path: str | None = None, quant: str | None = None) -> None:
        from transformers import AutoTokenizer

        from deg.model import base_snapshot

        self.url = (url or "http://127.0.0.1:8792").rstrip("/")
        path = Path(model_path) if model_path else base_snapshot()
        self.tok = AutoTokenizer.from_pretrained(str(path))
        self.native_ctx = native_context(path)
        self.server_ctx = self._server_ctx()
        self.name = label or "qwen3.5-0.8b-llamacpp-gbnf"
        self.info = {"model": (str(path) if model_path else
                               "Qwen3.5-0.8B bf16 GGUF (same base weights, not deg)"),
                     "quantization": quant, "native_context": self.native_ctx,
                     "server_n_ctx": self.server_ctx,
                     "decoding": "greedy, llama.cpp json_schema -> GBNF grammar",
                     "runtime": "llama.cpp server (CUDA, GPU0), batch 1"}

    def _server_ctx(self) -> int | None:
        import urllib.request
        try:
            with urllib.request.urlopen(self.url + "/props", timeout=30) as r:
                props = json.loads(r.read())
            return int(props["default_generation_settings"]["n_ctx"])
        except Exception:  # noqa: BLE001 - older servers: rely on the HTTP error instead
            return None

    def prompt(self, rec: dict[str, Any]) -> str:
        return chat_prompt(self.tok, rec)

    def predict(self, rec: dict[str, Any]) -> Prediction:
        import urllib.error
        import urllib.request

        prompt = self.prompt(rec)
        bos = self.tok.bos_token
        if bos and prompt.startswith(bos):  # llama.cpp adds BOS itself (GGUF add_bos_token)
            prompt = prompt[len(bos):]
        n = len(self.tok(prompt, add_special_tokens=False)["input_ids"])
        if self.native_ctx and n >= self.native_ctx:
            return Prediction(rec["uid"], error=f"context_exceeded: native {n} prompt tokens "
                                                f">= {self.native_ctx}", prompt_tokens=n)
        if self.server_ctx and n >= self.server_ctx:
            return Prediction(rec["uid"], error=f"context_exceeded: server n_ctx {n} prompt "
                                                f"tokens >= {self.server_ctx} (GPU memory)",
                              prompt_tokens=n)
        last = ""
        for mode, schema in (("full", grammar_schema(rec["schema"])),
                             ("no-format", grammar_schema(rec["schema"],
                                                          drop=("x-generate", "format",
                                                                "pattern")))):
            body = {"prompt": prompt, "n_predict": self.max_new_tokens, "temperature": 0.0,
                    "top_k": 1, "json_schema": schema, "cache_prompt": False}
            req = urllib.request.Request(self.url + "/completion",
                                         data=json.dumps(body).encode(),
                                         headers={"Content-Type": "application/json"})
            t0 = time.perf_counter()
            try:
                with urllib.request.urlopen(req, timeout=3600) as r:
                    resp = json.loads(r.read())
            except urllib.error.HTTPError as exc:
                last = f"HTTP {exc.code}: {exc.read().decode(errors='replace')[:200]}"
                if "context" in last.lower():
                    return Prediction(rec["uid"], error=f"context_exceeded: server {last}"[:240],
                                      prompt_tokens=n)
                continue
            except (OSError, ValueError) as exc:
                return Prediction(rec["uid"], error=f"server: {exc}"[:200])
            dt = time.perf_counter() - t0
            text = resp.get("content", "")
            gen, ptok = resp.get("tokens_predicted"), resp.get("tokens_evaluated")
            extra = {"grammar": mode, "server_timings": resp.get("timings")}
            try:
                value = json.loads(text)
            except json.JSONDecodeError as exc:
                return Prediction(rec["uid"], error=f"json: {exc}"[:200], generated_tokens=gen,
                                  prompt_tokens=ptok, latency_s=dt,
                                  extra={**extra, "text": text[-300:]})
            return Prediction(rec["uid"], value=value, generated_tokens=gen, prompt_tokens=ptok,
                              latency_s=dt, extra=extra)
        return Prediction(rec["uid"], error=f"grammar: {last}"[:240])