"""Real checkpoint inference, with bounded logit and attention storage."""
from __future__ import annotations

import os
import threading
import time
import uuid
from collections import OrderedDict

import numpy as np
import torch
from jinja2 import TemplateError
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Qwen/Qwen3-0.6B"
REVISION = "c1899de289a04d12100db370d81485cdf75e47ca"
GEMMA_ID = "ThomasTheMaker/gm3-270m-TinyGSM-reasoning"
MODELS = {
    MODEL_ID: dict(name="Qwen3-0.6B", revision=REVISION, native_thinking=True),
    GEMMA_ID: dict(name="Gemma 3 270M · TinyGSM reasoning",
                   revision="aa034f2832467400d4e0a09df9fae0c517d41579", native_thinking=False),
}
MAX_TOKENS = int(os.environ.get("QWEN_MAX_TOKENS", "512"))


class Runtime:
    def __init__(self, model_id=MODEL_ID):
        self.model_id = model_id
        self.spec = MODELS[model_id]
        self.config = None
        self.phase = "loading"
        self.message = f"Loading {self.spec['name']}…"
        self.tokenizer = None
        self.model = None
        self.device = "cpu"
        self.lock = threading.Lock()
        self.tokenizer_lock = threading.Lock()
        self.sessions = OrderedDict()
        self.active = None

    def load(self):
        try:
            torch.set_num_threads(min(6, os.cpu_count() or 4))
            self.config = AutoConfig.from_pretrained(self.model_id, revision=self.spec["revision"])
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, revision=self.spec["revision"])
            self.message = f"Loading the {self.spec['name']} checkpoint…"
            requested = os.environ.get("QWEN_DEVICE", "auto")
            self.device = ("mps" if torch.backends.mps.is_available() else "cpu") if requested == "auto" else requested
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_id, revision=self.spec["revision"], config=self.config, dtype=torch.float32,
                attn_implementation="eager",
            ).to(self.device).eval()
            self.phase = "ready"
            self.message = f"{self.spec['name']} · {self.device.upper()} · float32 · eager attention"
        except Exception as exc:
            self.phase = "error"
            self.message = f"Could not load the model: {exc}"
            import traceback
            traceback.print_exc()

    def metadata(self):
        config = self.config
        return dict(model=self.model_id, name=self.spec["name"], revision=self.spec["revision"],
                    native_thinking=self.spec["native_thinking"],
                    layers=config.num_hidden_layers if config else None,
                    heads=config.num_attention_heads if config else None,
                    kv_heads=config.num_key_value_heads if config else None,
                    vocab_size=config.vocab_size if config else None,
                    attention_scale=getattr(config, "query_pre_attn_scalar", config.head_dim) if config else None,
                    layer_types=getattr(config, "layer_types", None) if config else None,
                    sliding_window=getattr(config, "sliding_window", None) if config else None)

    def status(self):
        result = dict(phase=self.phase, message=self.message, tokenizer_ready=self.tokenizer is not None,
                      device=self.device, max_tokens=MAX_TOKENS, active=self.active, **self.metadata())
        if self.tokenizer:
            result["special_tokens"] = [dict(id=i, text=t.content, special=t.special)
                                        for i, t in self.tokenizer.added_tokens_decoder.items()
                                        if i < self.config.vocab_size and not t.content.startswith("<unused")
                                        and (t.special or t.content.startswith(("<", "[")))]
        return result

    def prepare(self, payload):
        if not self.tokenizer:
            raise ValueError("The tokenizer is still loading. Try again in a few seconds.")
        if payload.mode == "chat":
            if payload.thinking and not self.spec["native_thinking"]:
                raise ValueError("This Gemma checkpoint has no native reasoning switch. It generates learned Reasoning / Solution text.")
            if not payload.messages or not any(m.role == "user" for m in payload.messages):
                raise ValueError("Add at least one user message.")
            if payload.messages[-1].role != "user":
                raise ValueError("The last message must have the user role before a new response. Use Raw for an arbitrary prefix.")
            kwargs = {"enable_thinking": payload.thinking} if self.spec["native_thinking"] else {}
            try:
                text = self.tokenizer.apply_chat_template(
                    [m.model_dump() for m in payload.messages], tokenize=False,
                    add_generation_prompt=True, **kwargs,
                )
            except TemplateError as exc:
                raise ValueError(f"Invalid conversation for this model's chat template: {exc}") from exc
        else:
            text = payload.text
        with self.tokenizer_lock:
            ids = self.tokenizer.encode(text, add_special_tokens=False)
        if not ids:
            raise ValueError("The context is empty. Enter text or choose a template.")
        if any(i >= self.config.vocab_size for i in ids):
            raise ValueError("The input contains a tokenizer marker outside this text checkpoint's vocabulary. Remove unsupported image or other modality markers.")
        return text, ids

    def token(self, token_id):
        tid = int(token_id)
        return dict(id=tid, text=self.tokenizer.decode([tid], clean_up_tokenization_spaces=False),
                    piece=self.tokenizer.convert_ids_to_tokens(tid), special=tid in self.tokenizer.all_special_ids)

    def append_ids(self, sid, text=None):
        """Append one real token; the existing context IDs must never change."""
        previous = self.session(sid)["response"]
        if text is None:
            return [previous["predictions"][-1]["ids"][-1]]
        with self.tokenizer_lock:
            ids = self.tokenizer.encode(text, add_special_tokens=False)[:1]
        if not ids:
            raise ValueError("Enter a token or text to append. Spaces are preserved.")
        if any(tid >= self.config.vocab_size for tid in ids):
            raise ValueError("The text contains a tokenizer marker outside this text checkpoint's vocabulary.")
        return ids

    @torch.inference_mode()
    def analyze(self, payload, prefix=None, *, input_count=None, token_sources=None):
        if self.phase != "ready":
            raise ValueError(self.message)
        if not self.lock.acquire(blocking=False):
            raise BlockingIOError("The model is already computing. Wait for the current run to finish.")
        started = time.perf_counter()
        try:
            text, ids = self.prepare(payload) if prefix is None else prefix
            prefix_count = len(ids)
            if input_count is None:
                input_count = prefix_count
            sources = list(token_sources) if token_sources is not None else ["input"] * prefix_count
            if prefix_count + payload.max_new_tokens > MAX_TOKENS:
                raise ValueError(f"Received {prefix_count} context tokens + {payload.max_new_tokens} continuation tokens. The inspector limit is {MAX_TOKENS}; input has not been truncated. Shorten the context or continuation.")
            forward_message = f"Forward pass: {self.config.num_hidden_layers} layers × {self.config.num_attention_heads} heads…"
            self.active = "Generating the continuation…" if payload.max_new_tokens else forward_message
            input_ids = torch.tensor([ids], device=self.device)
            if payload.max_new_tokens:
                generated = self.model.generate(
                    input_ids, attention_mask=torch.ones_like(input_ids),
                    max_new_tokens=payload.max_new_tokens, do_sample=False,
                    pad_token_id=self.tokenizer.eos_token_id, use_cache=True,
                )
                ids = generated[0].tolist()
                input_ids = generated
                sources.extend(["predicted"] * (len(ids) - prefix_count))
            self.active = forward_message
            output = self.model.model(input_ids, output_hidden_states=True, output_attentions=True, use_cache=False)
            if output.attentions is None or any(a is None for a in output.attentions):
                raise RuntimeError("The model did not return attention. Eager attention is required.")
            # Float16 is only used to retain attention maps. Computation and lens are float32.
            # Transfer float32 first, then cast in NumPy into independently owned memory.
            # Fusing MPS→CPU and float16 conversion can corrupt strided attention tensors.
            attention = [a[0].cpu().contiguous().numpy().astype(np.float16, copy=True)
                         for a in output.attentions]
            for matrix in attention:
                sums = matrix.astype(np.float32).sum(axis=-1)
                if (not np.isfinite(matrix).all() or not np.allclose(sums, 1, atol=1e-3)
                        or np.count_nonzero(np.triu(matrix, 1))):
                    raise RuntimeError("An attention map failed the normalization or causal mask check.")
            top_ids, top_probs = [], []
            for layer in range(self.model.config.num_hidden_layers):
                self.active = f"Logit lens: layer {layer + 1} / {self.config.num_hidden_layers}…"
                # HF returns the last hidden state after final RMSNorm; do not normalize twice.
                hidden = output.hidden_states[layer + 1]
                if layer != self.model.config.num_hidden_layers - 1:
                    hidden = self.model.model.norm(hidden)
                layer_ids, layer_probs = [], []
                for chunk in hidden[0].split(32):
                    logits = self.model.lm_head(chunk).float()
                    softcap = getattr(self.config, "final_logit_softcapping", None)
                    if softcap is not None:
                        logits = torch.tanh(logits / softcap) * softcap
                    values, indices = logits.topk(5, dim=-1)
                    # Center first: early Gemma states can produce large logits, where
                    # subtracting an uncentered logsumexp loses float32 precision.
                    maximum = values[:, :1]
                    probabilities = (values - maximum - (logits - maximum).logsumexp(dim=-1, keepdim=True)).exp()
                    layer_ids.extend(indices.cpu().tolist())
                    layer_probs.extend(probabilities.cpu().tolist())
                top_ids.append(layer_ids)
                top_probs.append(layer_probs)
            del output
            sid = uuid.uuid4().hex
            vocab_ids = {tid for layer in top_ids for row in layer for tid in row}
            token_dict = {str(tid): self.token(tid) for tid in vocab_ids}
            response = dict(
                id=sid, context=text, tokens=[self.token(tid) for tid in ids], input_count=input_count,
                **self.metadata(), device=self.device, max_tokens=MAX_TOKENS,
                predictions=[dict(ids=[r[0] for r in layer], probabilities=[r[0] for r in probs])
                             for layer, probs in zip(top_ids, top_probs)],
                vocabulary=token_dict, elapsed=round(time.perf_counter() - started, 2),
                token_sources=sources, added_count=len(ids) - input_count,
                continuation_text=self.tokenizer.decode(ids[input_count:], clean_up_tokenization_spaces=False),
                generated_text=self.tokenizer.decode([tid for tid, source in zip(ids, sources) if source == "predicted"], clean_up_tokenization_spaces=False),
                generated_count=sources.count("predicted"), manual_count=sources.count("manual"),
                stopped_on_eos=bool(len(ids) > input_count and ids[-1] == self.tokenizer.eos_token_id),
                stop_token=self.token(ids[-1])["text"] if len(ids) > input_count else None,
                thinking=payload.thinking if payload.mode == "chat" and self.spec["native_thinking"] else None,
            )
            self.sessions[sid] = dict(attention=attention, top_ids=top_ids, top_probs=top_probs,
                                      response=response)
            # Bound memory; even at 512 tokens, at most ~470 MB of attention is retained.
            while len(self.sessions) > 2:
                self.sessions.popitem(last=False)
            return response
        finally:
            self.active = None
            self.lock.release()

    def session(self, sid):
        session = self.sessions.get(sid)
        if session is None:
            raise KeyError("This result has been evicted from memory. Run the analysis again.")
        return session

    def inspect(self, sid, layer, query):
        s = self.session(sid)
        weights = s["attention"][layer][:, query, :].astype(np.float32)
        return dict(top=[dict(**self.token(tid), probability=p) for tid, p in
                         zip(s["top_ids"][layer][query], s["top_probs"][layer][query])],
                    attention=weights.tolist())

    def attention_layers(self, sid, query, head="mean"):
        """Read one query across layers, without another model forward pass."""
        session = self.session(sid)
        weights = []
        for matrix in session["attention"]:
            if head == "mean":
                row = matrix[:, query, :].astype(np.float32).mean(axis=0)
            else:
                row = matrix[head, query, :].astype(np.float32)
            weights.append(row.tolist())
        return dict(query=query, head=head, weights=weights)