"""Serialized decisions over one-state native prefix caches."""

from __future__ import annotations

import copy
import hashlib
import inspect
import json
import time

from .core import direct_messages, softmax
from .direct import PROMPT_VERSION, encode_prompt


def _state_prefix(tokenizer, state) -> list[int]:
    row = {
        "id": "prefix-only",
        "state": state,
        "question": "prefix boundary placeholder",
        "options": [{"id": "yes", "description": "Yes"}, {"id": "no", "description": "No"}],
    }
    turns = direct_messages(row)
    prompt = tokenizer.apply_chat_template(
        turns, tokenize=False, add_generation_prompt=True, enable_thinking=False
    )
    payload = turns[-1]["content"]
    evidence = json.dumps({"evidence": state}, ensure_ascii=False)[:-1]
    if prompt.count(payload) != 1 or not payload.startswith(evidence):
        raise ValueError("Cannot establish a deterministic evidence prefix")
    text = prompt[: prompt.index(payload)] + evidence
    return tokenizer.encode(text, add_special_tokens=False)[:-1]


def _cached_forward(model, inputs):
    parameters = inspect.signature(model.forward).parameters
    if "logits_to_keep" not in parameters and hasattr(model, "get_base_model"):
        parameters = inspect.signature(model.get_base_model().forward).parameters
    if "logits_to_keep" not in parameters:
        raise RuntimeError("Model lacks selective last-position logits")
    return model(**inputs, use_cache=True, return_dict=True, logits_to_keep=1)


class SerialPrefixScorer:
    """Cache the current state, then score independent copied suffix branches."""

    def __init__(self, model, tokenizer, metadata: dict, max_tokens: int = 4096):
        self.model = model
        self.tokenizer = tokenizer
        self.metadata = {**metadata, "serving_config": "native-state-prefix-cache-v1"}
        self.max_tokens = max_tokens
        self.device = next(model.parameters()).device
        self.cache = None
        self.state = None
        self.prefix = None

    def score(self, row: dict) -> dict:
        import torch

        started = time.perf_counter()
        ids, slots, prompt_hash = encode_prompt(self.tokenizer, row, self.max_tokens)
        hit = self.cache is not None and row["state"] == self.state
        prefix = self.prefix if hit else _state_prefix(self.tokenizer, row["state"])
        if not prefix or ids[: len(prefix)] != prefix or len(ids) <= len(prefix):
            raise ValueError("State prefix does not match the full prompt")
        sync = lambda: torch.cuda.synchronize(self.device) if self.device.type == "cuda" else None
        prefill_seconds = 0.0
        self.model.eval()
        with torch.inference_mode():
            if not hit:
                self.cache = self.state = self.prefix = None
                sync()
                mark = time.perf_counter()
                output = _cached_forward(
                    self.model,
                    {
                        "input_ids": torch.tensor([prefix], dtype=torch.long, device=self.device),
                        "attention_mask": torch.ones((1, len(prefix)), dtype=torch.long, device=self.device),
                    },
                )
                self.cache = output.past_key_values
                del output
                sync()
                prefill_seconds = time.perf_counter() - mark
                if self.cache is None or self.cache.get_seq_length() != len(prefix):
                    raise RuntimeError("Invalid native prefix cache")
                self.state, self.prefix = row["state"], prefix
            sync()
            mark = time.perf_counter()
            branch = copy.deepcopy(self.cache)
            sync()
            copy_seconds = time.perf_counter() - mark
            inputs = {
                "input_ids": torch.tensor([ids[len(prefix) :]], dtype=torch.long, device=self.device),
                "attention_mask": torch.ones((1, len(ids)), dtype=torch.long, device=self.device),
                "past_key_values": branch,
            }
            sync()
            mark = time.perf_counter()
            output = _cached_forward(self.model, inputs)
            sync()
            suffix_seconds = time.perf_counter() - mark
            vocabulary = output.logits[0, -1, :].float()
            selected_tensor = vocabulary[slots]
            selected = selected_tensor.cpu().tolist()
            result = {
                "id": row["id"],
                "option_ids": [option["id"] for option in row["options"]],
                "probabilities": softmax(selected),
                "option_logits": selected,
                "answer_token_ids": slots,
                "input_tokens": len(ids),
                "prompt_sha256": prompt_hash,
                "prompt_version": PROMPT_VERSION,
                "model": self.metadata,
                "readout": "native-state-prefix-cache-last-position",
                "probability_status": "conditional option score; uncalibrated as decision confidence",
                "cache_hit": hit,
                "prefix_tokens": len(prefix),
                "prefix_sha256": hashlib.sha256(json.dumps(prefix).encode()).hexdigest(),
                "allowed_token_mass": float(
                    (selected_tensor.logsumexp(-1) - vocabulary.logsumexp(-1)).exp()
                ),
                "full_vocab_argmax_id": int(vocabulary.argmax()),
                "prefill_seconds": prefill_seconds,
                "copy_seconds": copy_seconds,
                "suffix_forward_seconds": suffix_seconds,
            }
            del output, branch
        sync()
        result["forward_seconds"] = prefill_seconds + suffix_seconds
        result["total_seconds"] = time.perf_counter() - started
        return result