"""Native Qwen vision, independent candidate paths, and a shared scalar head.

Public base weights remain external inputs. Checkpoints store the adapter,
head, processor, and immutable backbone revision.
"""

import copy
import hashlib
import json
from collections import OrderedDict
from pathlib import Path

import torch
from peft import LoraConfig, PeftModel, get_peft_model, get_peft_model_state_dict
from PIL import Image, ImageOps
from safetensors.torch import load_file, save_file
from torch import nn
from transformers import AutoProcessor, DynamicCache, Qwen3_5Model

from .calibration import UNFITTED, checkpoint_calibration
from .head import TARGETS, CandidateHead
from .schema import validate

MODEL = "Qwen/Qwen3.5-0.8B-Base"
REVISION = "dc7cdfe2ee4154fa7e30f5b51ca41bfa40174e68"
PROMPT_VERSION = "branchone-multimodal-segments-v2"


class BranchOne(nn.Module):
    def __init__(
        self,
        *,
        device="cuda",
        model_id=MODEL,
        revision=REVISION,
        max_length=8192,
        image_tokens=256,
        max_rows=1,
        batch_token_budget=8192,
        adapter_path=None,
        head_path=None,
        trainable=False,
        lora_rank=8,
        processor_path=None,
        hub_options=None,
    ):
        super().__init__()
        if not 2 <= max_length <= 32768:
            raise ValueError("Configured context must be between 2 and 32768 tokens")
        self.model_id, self.revision = model_id, revision
        self.max_length, self.image_tokens = max_length, image_tokens
        self.max_rows, self.batch_token_budget = max_rows, batch_token_budget
        self.temperature = 1.0
        self.calibration_status = UNFITTED
        self.cache_validated = False
        self.processor = AutoProcessor.from_pretrained(
            processor_path or model_id, revision=None if processor_path else revision
        )
        self.tokenizer = self.processor.tokenizer
        self.backbone, self.loading_info = Qwen3_5Model.from_pretrained(
            model_id,
            revision=revision,
            dtype=torch.float16,
            attn_implementation="sdpa",
            output_loading_info=True,
            **(hub_options or {}),
        )
        if any(
            self.loading_info.get(k)
            for k in ["missing_keys", "mismatched_keys", "error_msgs"]
        ):
            raise RuntimeError(
                f"Incomplete native multimodal weight load: {self.loading_info}"
            )
        self.backbone.requires_grad_(False)
        patch = self.backbone.config.vision_config.patch_size
        merge = self.backbone.config.vision_config.spatial_merge_size
        self.processor.image_processor.size = {
            "shortest_edge": min(64, image_tokens) * (patch * merge) ** 2,
            "longest_edge": image_tokens * (patch * merge) ** 2,
        }
        self.adapter_audit = {"warm_started": bool(adapter_path)}
        if adapter_path:
            self.backbone.language_model = PeftModel.from_pretrained(
                self.backbone.language_model, adapter_path, is_trainable=trainable
            )
            source = load_file(str(Path(adapter_path) / "adapter_model.safetensors"))
            actual = get_peft_model_state_dict(self.backbone.language_model)
            if source.keys() != actual.keys():
                raise RuntimeError("Warm-start adapter tensor names differ")
            error = max(
                float((actual[k].float().cpu() - v.float()).abs().max())
                for k, v in source.items()
            )
            if error != 0:
                raise RuntimeError(f"Warm-start adapter changed: {error}")
            self.adapter_audit.update(tensors=len(source), max_error=error)
        elif lora_rank:
            config = LoraConfig(
                task_type="FEATURE_EXTRACTION",
                r=lora_rank,
                lora_alpha=2 * lora_rank,
                lora_dropout=0.05,
                target_modules=TARGETS,
            )
            self.backbone.language_model = get_peft_model(
                self.backbone.language_model, config
            )
            if not trainable:
                self.backbone.language_model.requires_grad_(False)
        self.lora_rank = lora_rank
        self.head = CandidateHead(hidden=self.backbone.config.text_config.hidden_size)
        if head_path:
            self.head.load_state_dict(load_file(str(head_path)), strict=True)
        self._vision_cache = OrderedDict()
        self._prepared_cache = OrderedDict()
        self.to(device)
        self.train(trainable)

    @property
    def device(self):
        return next(self.backbone.parameters()).device

    def train(self, mode=True):
        super().train(mode)
        self.backbone.visual.eval()  # frozen vision must stay deterministic
        if mode:
            self.temperature, self.calibration_status = 1.0, UNFITTED
        return self

    def enable_gradient_checkpointing(self):
        self.backbone.language_model.gradient_checkpointing_enable(
            gradient_checkpointing_kwargs={"use_reentrant": False}
        )

    @staticmethod
    def _lru_put(cache, key, value, limit=64):
        cache[key] = value
        cache.move_to_end(key)
        while len(cache) > limit:
            cache.popitem(last=False)

    @staticmethod
    def _image(value):
        if isinstance(value, Image.Image):
            im = ImageOps.exif_transpose(value).convert("RGB")
        elif isinstance(value, (str, Path)):
            if "://" in str(value):
                raise ValueError(
                    "Pass a local image path; URL fetching is not supported"
                )
            with Image.open(value) as opened:
                im = ImageOps.exif_transpose(opened).convert("RGB")
        else:
            raise ValueError("image must be a local path or a PIL image")
        content_hash = hashlib.sha256(str(im.size).encode() + im.tobytes()).hexdigest()
        return im, content_hash

    def prepare(self, request):
        """Build CPU tensors, enforce expanded token limits, and keep IDs external."""
        validate(request)
        if "images" in request:
            raise ValueError(
                "Use the singular image field; exactly one image is supported"
            )
        image_value = request.get("image")
        image_hash = None
        if image_value is not None:
            im, image_hash = self._image(image_value)
        state = "State:\n" + request["state"] + "\n"
        key = hashlib.sha256(
            json.dumps(
                {
                    "state": state,
                    "question": request["question"],
                    "type": request.get("type", "choice"),
                    "image": image_hash,
                    "candidates": [c["description"] for c in request["candidates"]],
                    "max_length": self.max_length,
                    "image_tokens": self.image_tokens,
                },
                sort_keys=True,
            ).encode()
        ).hexdigest()
        if key in self._prepared_cache:
            self._prepared_cache.move_to_end(key)
            return self._prepared_cache[key]

        def enc(text):
            return self.tokenizer.encode(text, add_special_tokens=False)

        pixels = grid = None
        if image_hash:
            text = (
                "State:\n<|vision_start|><|image_pad|><|vision_end|>\n"
                + request["state"]
                + "\n"
            )
            processed = self.processor(
                text=text, images=im, add_special_tokens=False, return_tensors="pt"
            )
            state_ids = processed["input_ids"][0].tolist()
            pixels, grid = processed["pixel_values"], processed["image_grid_thw"]
            visual_count = state_ids.count(self.backbone.config.image_token_id)
            if not 0 < visual_count <= self.image_tokens:
                raise ValueError(
                    f"Visual token count {visual_count} exceeds budget {self.image_tokens}"
                )
        else:
            state_ids, visual_count = enc(state), 0
        prefix = state_ids + enc(
            "Task: "
            + request.get("type", "choice")
            + "\nQuestion:\n"
            + request["question"]
            + "\n"
        )
        suffixes = [
            enc("Candidate:\n" + c["description"] + "\nDecision:")
            for c in request["candidates"]
        ]
        paths = [prefix + suffix for suffix in suffixes]
        lengths = [len(p) for p in paths]
        if max(lengths) > self.max_length:
            raise ValueError(
                f"Expanded branch has {max(lengths)} tokens; limit {self.max_length}; no truncation"
            )
        result = {
            "paths": paths,
            "prefix": prefix,
            "suffixes": suffixes,
            "pixel_values": pixels,
            "image_grid_thw": grid,
            "image_hash": image_hash,
            "token_counts": lengths,
            "visual_tokens": visual_count,
            "image_size": list(im.size) if image_hash else None,
        }
        self._lru_put(self._prepared_cache, key, result)
        return result

    @torch.no_grad()
    def visual_features(self, prepared, reuse=True):
        if not prepared["image_hash"]:
            return None
        key = prepared["image_hash"]
        if reuse and key in self._vision_cache:
            self._vision_cache.move_to_end(key)
            return self._vision_cache[key].to(self.device)
        out = self.backbone.get_image_features(
            prepared["pixel_values"].to(self.device),
            prepared["image_grid_thw"].to(self.device),
            return_dict=True,
        )
        features = torch.cat(out.pooler_output, dim=0).detach()
        if reuse:
            self._lru_put(self._vision_cache, key, features.cpu())
        return features

    def _inputs(self, paths, prepared, features, *, pad_to_bucket=True):
        lengths = torch.tensor([len(p) for p in paths], device=self.device)
        width = int(lengths.max())
        if pad_to_bucket:
            width = ((width + 63) // 64) * 64
        pad_id = self.tokenizer.pad_token_id
        if pad_id is None:
            pad_id = self.tokenizer.eos_token_id
        ids = torch.full(
            (len(paths), width), pad_id, dtype=torch.long, device=self.device
        )
        for i, path in enumerate(paths):
            ids[i, : len(path)] = torch.tensor(path, device=self.device)
        mask = torch.arange(width, device=self.device)[None, :] < lengths[:, None]
        embeds = self.backbone.get_input_embeddings()(ids)
        if features is not None:
            is_image = ids == self.backbone.config.image_token_id
            assert int(is_image.sum()) == len(paths) * features.shape[0]
            embeds = embeds.masked_scatter(
                is_image[..., None], features.repeat(len(paths), 1).to(embeds.dtype)
            )
            grid = prepared["image_grid_thw"].to(self.device).repeat(len(paths), 1)
            positions, _ = self.backbone.get_rope_index(
                ids,
                mm_token_type_ids=is_image.long(),
                image_grid_thw=grid,
                attention_mask=mask.long(),
            )
        else:
            positions = torch.arange(width, device=self.device)[None, :].expand(
                len(paths), -1
            )
            positions = positions.masked_fill(~mask, 0)
        return {
            "inputs_embeds": embeds,
            "attention_mask": mask.long(),
            "position_ids": positions,
        }, lengths

    def representations(self, request, *, reuse_image=True, max_rows=None):
        prepared = self.prepare(request)
        features = self.visual_features(prepared, reuse=reuse_image)
        paths = prepared["paths"]
        buckets = {}
        for i, p in enumerate(paths):
            buckets.setdefault(((len(p) + 63) // 64) * 64, []).append(i)
        hidden = [None] * len(paths)
        for width, indices in sorted(buckets.items()):
            batch_size = max(
                1, min(max_rows or self.max_rows, self.batch_token_budget // width)
            )
            for start in range(0, len(indices), batch_size):
                group = indices[start : start + batch_size]
                inputs, lengths = self._inputs(
                    [paths[i] for i in group], prepared, features
                )
                with torch.autocast(
                    "cuda", dtype=torch.float16, enabled=self.device.type == "cuda"
                ):
                    out = self.backbone.language_model(
                        **inputs, use_cache=False
                    ).last_hidden_state
                leaves = out[
                    torch.arange(len(group), device=self.device), lengths - 1
                ].float()
                for i, leaf in zip(group, leaves):
                    hidden[i] = leaf
        return torch.stack(hidden), prepared

    def forward(
        self, request, *, return_features=False, reuse_image=True, max_rows=None
    ):
        hidden, prepared = self.representations(
            request, reuse_image=reuse_image, max_rows=max_rows
        )
        valid = torch.ones((1, len(hidden)), dtype=torch.bool, device=self.device)
        logits, features = self.head(hidden[None], valid, return_features=True)
        return (logits[0], features[0], prepared) if return_features else logits[0]

    @torch.inference_mode()
    def native_reference(self, request):
        """Native full multimodal forward with matching padding for integration checks."""
        prepared = self.prepare(request)
        hidden = []
        for path in prepared["paths"]:
            width = ((len(path) + 63) // 64) * 64
            pad_id = self.tokenizer.pad_token_id
            if pad_id is None:
                pad_id = self.tokenizer.eos_token_id
            ids = torch.tensor(
                [path + [pad_id] * (width - len(path))], device=self.device
            )
            mask = (torch.arange(width, device=self.device)[None] < len(path)).long()
            kwargs = {"input_ids": ids, "attention_mask": mask, "use_cache": False}
            if prepared["image_hash"]:
                kwargs.update(
                    pixel_values=prepared["pixel_values"].to(self.device),
                    image_grid_thw=prepared["image_grid_thw"].to(self.device),
                    mm_token_type_ids=(
                        ids == self.backbone.config.image_token_id
                    ).long(),
                )
            hidden.append(
                self.backbone(**kwargs).last_hidden_state[0, len(path) - 1].float()
            )
        h = torch.stack(hidden)[None]
        return self.head(
            h, torch.ones(h.shape[:2], dtype=torch.bool, device=self.device)
        )[0]

    @torch.inference_mode()
    def cached_logits(self, request):
        """Experimental native-cache fork; callers must validate before serving."""
        prepared = self.prepare(request)
        features = self.visual_features(prepared)
        full, _ = self._inputs(
            prepared["paths"], prepared, features, pad_to_bucket=False
        )
        prefix_length = len(prepared["prefix"])
        positions = full["position_ids"]
        prefix_pos = (
            positions[:, :1, :prefix_length]
            if positions.ndim == 3
            else positions[:1, :prefix_length]
        )
        out = self.backbone.language_model(
            inputs_embeds=full["inputs_embeds"][:1, :prefix_length],
            attention_mask=full["attention_mask"][:1, :prefix_length],
            position_ids=prefix_pos,
            past_key_values=DynamicCache(config=self.backbone.config.text_config),
            use_cache=True,
        )
        prefix_cache = out.past_key_values
        hidden = []
        for i, path in enumerate(prepared["paths"]):
            length = len(path)
            pos = (
                positions[:, i : i + 1, prefix_length:length]
                if positions.ndim == 3
                else positions[i : i + 1, prefix_length:length]
            )
            branch = self.backbone.language_model(
                inputs_embeds=full["inputs_embeds"][i : i + 1, prefix_length:length],
                attention_mask=full["attention_mask"][i : i + 1, :length],
                position_ids=pos,
                past_key_values=copy.deepcopy(prefix_cache),
                use_cache=True,
            )
            hidden.append(branch.last_hidden_state[0, -1].float())
        if prefix_cache.get_seq_length() != prefix_length:
            raise RuntimeError("Shared prefix cache was mutated")
        h = torch.stack(hidden)[None]
        return self.head(
            h, torch.ones(h.shape[:2], dtype=torch.bool, device=self.device)
        )[0]

    @torch.inference_mode()
    def predict(self, request, *, cached=False, threshold=None):
        self.eval()
        if cached and not self.cache_validated:
            raise ValueError(
                "Prefix cache has not passed checkpoint-specific validation"
            )
        logits = self.cached_logits(request) if cached else self(request)
        p = (logits / self.temperature).softmax(-1).cpu()
        ids = [c["id"] for c in request["candidates"]]
        selected = int(p.argmax())
        prepared = self.prepare(request)
        result = {
            "selected_id": ids[selected],
            "candidate_ids": ids,
            "probabilities": p.tolist(),
            "logits": logits.cpu().tolist(),
            "temperature": self.temperature,
            "calibration_status": self.calibration_status,
            "model_version": "branchone-0.3.0-independent",
            "context_limit": self.max_length,
            "branch_token_counts": prepared["token_counts"],
            "visual_tokens": prepared["visual_tokens"],
            "image_size": prepared["image_size"],
            "action": "defer"
            if threshold is None
            else ("accept" if float(p.max()) >= threshold else "defer"),
            "threshold": threshold,
            "deferral_status": "no validated policy"
            if threshold is None
            else "caller-supplied policy",
        }
        if request.get("type") == "boolean":
            result.update(
                p_true=float(p[ids.index("true")]), value=ids[selected] == "true"
            )
        if request.get("type") == "score":
            result["expected_score"] = sum(
                float(x) * c["level"] for x, c in zip(p, request["candidates"])
            )
        return result

    def save(self, path):
        """Save trainable weights without carrying calibration to a new payload."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)
        self.backbone.language_model.save_pretrained(path / "adapter")
        save_file(
            {
                k: v.detach().cpu().contiguous()
                for k, v in self.head.state_dict().items()
            },
            str(path / "head.safetensors"),
        )
        self.processor.save_pretrained(path / "processor")
        meta = {
            "model_id": self.model_id,
            "revision": self.revision,
            "max_length": self.max_length,
            "image_tokens": self.image_tokens,
            "max_rows": self.max_rows,
            "batch_token_budget": self.batch_token_budget,
            "lora_rank": self.lora_rank,
            "prompt_version": PROMPT_VERSION,
            "temperature": 1.0,
            "calibration_status": UNFITTED,
            "cache_validated": False,
        }
        (path / "model_config.json").write_text(json.dumps(meta, indent=2))
        (path / "calibration.json").unlink(missing_ok=True)

    @classmethod
    def load(cls, path, *, device="cuda", trainable=False, hub_options=None):
        path = Path(path)
        meta = json.loads((path / "model_config.json").read_text())
        if meta.pop("prompt_version") != PROMPT_VERSION:
            raise ValueError("Incompatible prompt version")
        temperature, status = checkpoint_calibration(path, meta, trainable=trainable)
        meta.pop("temperature")
        meta.pop("calibration_status")
        meta.pop("cache_validated", None)
        model = cls(
            **meta,
            device=device,
            trainable=trainable,
            adapter_path=path / "adapter",
            head_path=path / "head.safetensors",
            processor_path=path / "processor",
            hub_options=hub_options,
        )
        model.temperature, model.calibration_status = temperature, status
        return model