"""The inference engine behind the server: decisions, generation, and decide-then-generate.

One model in memory. Decisions prefill the state once with the base (adapter off), then run each
question suffix with the adapter on against that cache. Generation runs with the adapter off.
`auto` decides first and, when the router flags the question, decodes from the same state cache,
so escalation never prefills the state a second time.

Requests use TypeSafe-style question types: noul (yes/no), choice (named options, at most 16)
and score (ordered levels, at most 6).
"""

import json
import threading
from pathlib import Path

import numpy as np
import torch

from dmm import lora
from dmm.artifacts import read_router, router_file, weights_file
from dmm.cache import StateCache
from dmm.head import DecisionHead, option_logits
from dmm.router import Router, confidence_features
from dmm.template import Template, chat_text, decision_text


def to_item(state: str, name: str, q: dict) -> dict:
    """A TypeSafe-style question as one of our decision items."""
    kind = q["type"]
    text = str(q.get("instructions") or "").strip()
    if kind == "noul":
        return {
            "id": name,
            "kind": "yesno",
            "state": state,
            "question": text,
            "options": ["no", "yes"],
            "answer": None,
        }
    if kind == "choice":
        names = list(q["criteria"])
        if not 1 <= len(names) <= 16:
            raise ValueError(f"{name}: choice takes 1 to 16 options, got {len(names)}")
        opts = [n if not q["criteria"][n] else f"{n}: {q['criteria'][n]}" for n in names]
        return {
            "id": name,
            "kind": "choice",
            "state": state,
            "question": text,
            "options": opts,
            "answer": None,
            "_names": names,
        }
    if kind == "score":
        levels = list(q["criteria"])
        if not 2 <= len(levels) <= 6:
            raise ValueError(f"{name}: score takes 2 to 6 levels, got {len(levels)}")
        legend = "; ".join(f"{i} = {lv}" for i, lv in enumerate(levels))
        return {
            "id": name,
            "kind": "score",
            "state": state,
            "question": f"{text} ({legend})",
            "options": [str(i) for i in range(len(levels))],
            "answer": None,
        }
    raise ValueError(f"{name}: unknown question type {kind!r}")


class Engine:
    def __init__(
        self,
        model,
        tokenizer,
        head: DecisionHead,
        temperatures: dict,
        router: Router | None = None,
        router_threshold: float = 1.0,
    ):
        self.model, self.tok, self.head = model, tokenizer, head
        self.tpl = Template(tokenizer)
        self.temperatures, self.router, self.threshold = temperatures, router, router_threshold
        self.pad = tokenizer.pad_token_id or tokenizer.eos_token_id
        self.lock = threading.Lock()  # one request on the model at a time
        self.model.eval()

    @classmethod
    def from_artifacts(cls, model, tokenizer, cfg: dict, ckpt_dir: str, router_dir: str | None):
        from train.train_decision import load_checkpoint, prepare

        head, _ = prepare(model, tokenizer, cfg)
        load_checkpoint(model, head, weights_file(ckpt_dir))
        temps = json.loads(Path(f"{ckpt_dir}/temperatures.json").read_text())
        router, threshold = None, 1.0
        if router_dir:
            router = Router(model.config.hidden_size)
            router.load_state_dict(read_router(router_file(router_dir)))
            threshold = json.loads(Path(f"{router_dir}/operating_point.json").read_text())[
                "router_threshold"
            ]
        return cls(model, tokenizer, head.float(), temps, router, threshold)

    # Decisions ---------------------------------------------------------------------------

    @torch.no_grad()
    def _decide_on(self, cache: StateCache, item: dict) -> dict:
        h = cache.run_suffix(self.tpl.encode(decision_text(item)))
        z = option_logits(self.head(h[None].to(self.head.proj.weight.device)), [item])[0]
        z = z.float() / self.temperatures.get(item["kind"], 1.0)
        p = torch.softmax(z, -1).cpu().numpy().astype(np.float64)
        score = None
        if self.router is not None:
            hf = h.float().cpu().numpy()[None]
            score = float(self.router.score(hf, confidence_features([p]))[0])
        return {
            "probs": p,
            "router_score": score,
            "escalate": score is not None and score >= self.threshold,
        }

    def decide(self, state: str, questions: dict[str, dict]) -> dict:
        items = {name: to_item(state, name, q) for name, q in questions.items()}
        with self.lock:
            cache = StateCache(self.model, self.tpl.state_ids(next(iter(items.values()))))
            raw = {name: self._decide_on(cache, it) for name, it in items.items()}
        return {name: self._answer(items[name], r) for name, r in raw.items()}

    @staticmethod
    def _answer(item: dict, r: dict) -> dict:
        p = r["probs"]
        out = {
            "type": {"yesno": "noul", "choice": "choice", "score": "score"}[item["kind"]],
            "router_score": r["router_score"],
            "needs_generation": r["escalate"],
        }
        if item["kind"] == "yesno":
            out |= {
                "noul": float(p[1]),
                "probabilities": {"false": float(p[0]), "true": float(p[1])},
            }
        elif item["kind"] == "choice":
            names = item["_names"]
            out |= {
                "choice": names[int(np.argmax(p))],
                "probabilities": {n: float(x) for n, x in zip(names, p, strict=True)},
            }
        else:
            out |= {
                "score": float(np.dot(p, np.arange(len(p)))),
                "probabilities": {str(i): float(x) for i, x in enumerate(p)},
            }
        return out

    # Generation --------------------------------------------------------------------------

    @torch.no_grad()
    def _generate(
        self, ids: list[int], max_new_tokens: int, cache: StateCache | None = None, streamer=None
    ):
        x = torch.tensor([ids], device=self.model.device)
        kw = {"past_key_values": cache.kv} if cache is not None else {}
        if streamer is not None:
            kw["streamer"] = streamer
        with lora.disabled(self.model):
            out = self.model.generate(
                x,
                attention_mask=torch.ones_like(x),
                do_sample=False,
                max_new_tokens=max_new_tokens,
                pad_token_id=self.pad,
                **kw,
            )
        return out[0, len(ids) :].tolist()

    def _stream(self, ids: list[int], max_new_tokens: int, cache: StateCache | None = None):
        """Generate as _generate does, yielding {"event": "text"} pieces as they decode.

        Returns the generated token ids (the generator's return value). Generation runs in a
        thread feeding a streamer; the thread is always joined, even if the caller stops early,
        so the model and the state cache are never touched by two requests at once.
        """
        from transformers import TextIteratorStreamer

        streamer = TextIteratorStreamer(self.tok, skip_prompt=True, skip_special_tokens=True)
        result: dict = {}

        def run():
            try:
                result["tokens"] = self._generate(ids, max_new_tokens, cache, streamer)
            except BaseException as e:  # surface in the caller's thread
                result["error"] = e
                streamer.end()

        t = threading.Thread(target=run, daemon=True)
        t.start()
        try:
            for piece in streamer:
                if piece:
                    yield {"event": "text", "text": piece}
        finally:
            t.join()
        if "error" in result:
            raise result["error"]
        return result["tokens"]

    def chat_stream(self, messages: list[dict], max_new_tokens: int = 512):
        """chat(), streamed: {"event": "text"} pieces, then {"event": "done", "tokens": [...]}."""
        ids = self.tok.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=True, return_dict=True
        )["input_ids"]
        with self.lock:
            tokens = yield from self._stream(ids, max_new_tokens)
        yield {"event": "done", "tokens": tokens}

    def auto_stream(self, state: str, name: str, question: dict, max_new_tokens: int = 512):
        """auto(), streamed. Yields {"event": "decision", ...} as soon as the quick decision and
        the router's call are known, then {"event": "text"} reasoning pieces when escalated, then
        {"event": "done", ...} carrying exactly what auto() returns."""
        item = to_item(state, name, question)
        with self.lock:
            cache = StateCache(self.model, self.tpl.state_ids(item))
            r = self._decide_on(cache, item)
            out = {
                "decision": self._answer(item, r),
                "escalated": r["escalate"],
                "state_prefills": 1,
            }
            yield {"event": "decision", **out}
            if r["escalate"]:
                ids = cache.ids + self.tpl.encode(chat_text(item))
                tokens = yield from self._stream(ids, max_new_tokens, cache)
                out["generation_tokens"] = tokens
                out["generation"] = self.tok.decode(tokens, skip_special_tokens=True)
                cache.kv.crop(len(cache))
        yield {"event": "done", **out}

    def chat(self, messages: list[dict], max_new_tokens: int = 512) -> list[int]:
        ids = self.tok.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=True, return_dict=True
        )["input_ids"]
        with self.lock:
            return self._generate(ids, max_new_tokens)

    def auto(self, state: str, name: str, question: dict, max_new_tokens: int = 512) -> dict:
        """Decide; when the router flags it, decode from the same state cache."""
        item = to_item(state, name, question)
        with self.lock:
            cache = StateCache(self.model, self.tpl.state_ids(item))
            r = self._decide_on(cache, item)
            out = {
                "decision": self._answer(item, r),
                "escalated": r["escalate"],
                "state_prefills": 1,
            }
            if r["escalate"]:
                ids = cache.ids + self.tpl.encode(chat_text(item))
                out["generation_tokens"] = self._generate(ids, max_new_tokens, cache)
                out["generation"] = self.tok.decode(
                    out["generation_tokens"], skip_special_tokens=True
                )
                cache.kv.crop(len(cache))
        return out