"""Experimental Qwen3.5 multimodal KEV scorer with one visual prefix per request.

This is deliberately separate from the stable text DecisionModel. It uses the
official processor and the model's native hybrid cache/RoPE continuation; it
never materializes vocabulary logits or shares mutable caches across requests.
"""
from __future__ import annotations

import copy
import json
import math
import threading
from contextlib import nullcontext
from pathlib import Path

import torch
from PIL import Image

from .visual_prompt import SPECIAL, build_content


class VisualModelError(ValueError):
    """An image/question cannot be represented by this scorer safely."""


class VisualDecisionModel:
    def __init__(self, processor, model, head, device: str = "cuda:0", chunk_size: int = 64,
                 max_input_tokens: int | None = None):
        if chunk_size < 1:
            raise ValueError("chunk_size must be positive")
        if max_input_tokens is not None and max_input_tokens < 1:
            raise ValueError("max_input_tokens must be positive when set")
        self.processor, self.model, self.head = processor, model, head
        self.device, self.chunk_size = str(device), int(chunk_size)
        self.max_input_tokens = None if max_input_tokens is None else int(max_input_tokens)
        self.lock = threading.Lock()  # Qwen3.5 stores rope_deltas on its backbone
        self.visual_prefix_forwards = 0
        self.question_forwards = 0
        self.last_usage = {}

    @classmethod
    def from_run(cls, run: str | Path, device: str = "cuda:0", chunk_size: int = 64,
                 max_input_tokens: int | None = None):
        """Load one private Kaggle-style adapter/head bundle; no model is uploaded."""
        from peft import (LoraConfig, get_peft_model, get_peft_model_state_dict,
                          load_peft_weights, set_peft_model_state_dict)
        from transformers import AutoModelForImageTextToText, AutoProcessor
        from .model import PointerHead

        root = Path(run)
        report = json.loads((root / "report.json").read_text(encoding="utf-8"))
        if report.get("success") is not True:
            raise VisualModelError("visual checkpoint report is not successful")
        name, revision = report["model"], report["revision"]
        dtype = torch.float16 if str(device).startswith("cuda") else torch.float32
        processor = AutoProcessor.from_pretrained(name, revision=revision)
        base = AutoModelForImageTextToText.from_pretrained(name, revision=revision, dtype=dtype,
                                                           device_map=str(device))
        base.requires_grad_(False)
        config = json.loads((root / "adapter" / "adapter_config.json").read_text(encoding="utf-8"))
        targets = [module_name for module_name, module in base.named_modules()
                   if module_name.startswith("model.language_model.") and isinstance(module, torch.nn.Linear)]
        if len(targets) != 186 or any("visual" in item or "merger" in item for item in targets):
            raise VisualModelError(f"language-only LoRA target set drifted: {len(targets)} modules")
        model = get_peft_model(base, LoraConfig(r=int(config["r"]), lora_alpha=int(config["lora_alpha"]),
                                                lora_dropout=float(config["lora_dropout"]), bias="none",
                                                target_modules=targets))
        source = load_peft_weights(str(root / "adapter"), device="cpu")
        expected = get_peft_model_state_dict(model)
        if set(source) != set(expected) or any(source[key].shape != expected[key].shape for key in source):
            raise VisualModelError("visual adapter does not strictly cover the language-only LoRA target set")
        set_peft_model_state_dict(model, source)
        loaded = get_peft_model_state_dict(model)
        if any(not torch.equal(loaded[key].cpu(), source[key].to(loaded[key].dtype)) for key in source):
            raise VisualModelError("visual adapter changed during strict load")
        head = PointerHead(int(base.config.text_config.hidden_size)).to(device, dtype=torch.float32)
        state = torch.load(root / "pointer_head.pt", map_location="cpu", weights_only=True)
        head.load_state_dict(state, strict=True)
        model.eval().requires_grad_(False)
        head.eval().requires_grad_(False)
        return cls(processor, model, head, device, chunk_size, max_input_tokens)

    def _backbone(self):
        base = self.model.get_base_model() if hasattr(self.model, "get_base_model") else self.model
        backbone = getattr(base, "model", None)
        if backbone is None or not hasattr(backbone, "rope_deltas"):
            raise VisualModelError("expected Qwen3.5 multimodal backbone with native RoPE state")
        return backbone

    def _prepared(self, rec, images):
        if not isinstance(rec, dict) or not isinstance(rec.get("questions"), list) or not rec["questions"]:
            raise VisualModelError("visual record needs one or more questions")
        if not isinstance(images, list) or not 1 <= len(images) <= 4 or any(not isinstance(im, Image.Image) for im in images):
            raise VisualModelError("visual scorer needs 1..4 PIL images")
        tok = self.processor.tokenizer
        tokens = {x: tok.convert_tokens_to_ids(x) for x in SPECIAL + ("<|image_pad|>",)}
        if any(not isinstance(v, int) or v < 0 for v in tokens.values()):
            raise VisualModelError("Qwen3.5 tokenizer is missing a required delimiter")
        prepared = []
        for question in rec["questions"]:
            content = build_content(rec.get("state"), question, len(images))
            prompt = tok.apply_chat_template([{"role": "user", "content": content}],
                                             tokenize=False, add_generation_prompt=False)
            batch = self.processor(text=[prompt], images=images, return_tensors="pt", padding=True)
            row = batch["input_ids"][0].tolist()
            counts = {x: row.count(tid) for x, tid in tokens.items()}
            if ([counts[x] for x in SPECIAL] != [1, 1, len(question["options"]), len(question["options"]), 1]
                    or counts["<|image_pad|>"] < len(images)):
                raise VisualModelError(f"visual delimiter counts are invalid: {counts}")
            qpos = row.index(tokens[SPECIAL[1]])
            spos = row.index(tokens[SPECIAL[0]])
            options = [i - qpos for i, value in enumerate(row) if value == tokens[SPECIAL[3]]]
            decision = row.index(tokens[SPECIAL[4]]) - qpos
            image_positions = [i for i, value in enumerate(row) if value == tokens["<|image_pad|>"]]
            if not (max(image_positions) < spos < qpos and all(0 <= o < decision for o in options)):
                raise VisualModelError("visual prefix/question boundaries are invalid")
            prepared.append((batch, qpos, options, decision))
        first, prefix_len, _, _ = prepared[0]
        for batch, qpos, _, _ in prepared[1:]:
            if qpos != prefix_len or not torch.equal(batch["input_ids"][:, :qpos], first["input_ids"][:, :prefix_len]):
                raise VisualModelError("visual prefix differs across questions")
            for key in ("pixel_values", "image_grid_thw", "mm_token_type_ids"):
                if key in first or key in batch:
                    if key not in first or key not in batch:
                        raise VisualModelError(f"visual processor changed {key} across questions")
                    left = first[key][:, :prefix_len] if key == "mm_token_type_ids" else first[key]
                    right = batch[key][:, :qpos] if key == "mm_token_type_ids" else batch[key]
                    if not torch.equal(left, right):
                        raise VisualModelError(f"visual processor changed {key} across questions")
        return prepared, tokens

    def _autocast(self):
        return torch.autocast("cuda", dtype=torch.float16) if self.device.startswith("cuda") else nullcontext()

    def probs(self, rec: dict, images: list[Image.Image]) -> list[torch.Tensor]:
        """Score isolated questions; one vision encode and one immutable prefix cache."""
        with self.lock, torch.inference_mode(), self._autocast():
            prepared, tokens = self._prepared(rec, images)
            first, prefix_len, _, _ = prepared[0]
            input_tokens = prefix_len + sum(batch["input_ids"].shape[1] - qpos for batch, qpos, _, _ in prepared)
            if self.max_input_tokens is not None and input_tokens > self.max_input_tokens:
                raise VisualModelError(f"request uses {input_tokens} tokens; serving limit is {self.max_input_tokens}")
            backbone = self._backbone()
            first = {key: value.to(self.device) for key, value in first.items()}
            ids = first["input_ids"]
            image_positions = (ids[0] == tokens["<|image_pad|>"]).nonzero().flatten()
            initial_end = max(min(self.chunk_size, prefix_len), int(image_positions[-1]) + 1)
            if initial_end > 8192:
                raise VisualModelError("image token region exceeds the bounded first chunk")
            backbone.rope_deltas = None
            cache = None
            start = 0
            prefix_chunks = 0
            while start < prefix_len:
                end = initial_end if start == 0 else min(start + self.chunk_size, prefix_len)
                if start == 0:
                    kwargs = {key: (value[:, :end] if key in ("input_ids", "attention_mask", "mm_token_type_ids") else value)
                              for key, value in first.items()}
                    out = backbone(**kwargs, use_cache=True)
                    self.visual_prefix_forwards += 1
                else:
                    out = backbone(input_ids=ids[:, start:end], past_key_values=cache, use_cache=True)
                cache = out.past_key_values
                if cache is None or cache.get_seq_length() != end:
                    raise VisualModelError("native prefix cache length mismatch")
                prefix_chunks += 1
                del out
                start = end
            if cache.get_seq_length() != prefix_len:
                raise VisualModelError("native prefix did not consume the full state")
            prefix_cache = cache
            answers = []
            for batch, qpos, option_offsets, decision_offset in prepared:
                branch = batch["input_ids"][:, qpos:].to(self.device)
                branch_cache = copy.deepcopy(prefix_cache)
                needed = set(option_offsets + [decision_offset])
                hidden = {}
                for start in range(0, branch.shape[1], self.chunk_size):
                    end = min(start + self.chunk_size, branch.shape[1])
                    out = backbone(input_ids=branch[:, start:end], past_key_values=branch_cache, use_cache=True)
                    branch_cache = out.past_key_values
                    self.question_forwards += 1
                    for offset in needed:
                        if start <= offset < end:
                            hidden[offset] = out.last_hidden_state[0, offset - start].float()
                    del out
                if set(hidden) != needed:
                    raise VisualModelError("native branch omitted a pointer position")
                logits = self.head(hidden[decision_offset], torch.stack([hidden[x] for x in option_offsets]))
                if not torch.isfinite(logits).all():
                    raise VisualModelError("non-finite visual pointer logits")
                answers.append(torch.softmax(logits.float(), dim=-1).cpu())
            self.last_usage = {"input_tokens": input_tokens,
                               "state_tokens": prefix_len, "images": len(images),
                               "visual_prefix_forwards": 1, "visual_prefix_chunks": prefix_chunks,
                               "question_forwards": len(prepared)}
            return answers