"""Matilda Decision for 🤗 Transformers (`trust_remote_code=True`).

`AutoModel.from_pretrained(repo, trust_remote_code=True)` loads the repository through its own runtime
(`matilda_decision_runtime.py`: bidirectional prompt attention and the 255-option decision readout) and returns a
model with `system_one(state=..., questions={...}, images=[...])`. Without `trust_remote_code` the repository loads
as the plain, causal `Qwen3_5Model` backbone, which does not give the model's decisions.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any

import torch
from transformers import PreTrainedModel
from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config

try:
    from .matilda_decision_runtime import MAX_IMAGE_PIXELS, DecisionModel, system_one
except ImportError:
    from matilda_decision_runtime import MAX_IMAGE_PIXELS, DecisionModel, system_one

HUB_OPTIONS = ("revision", "cache_dir", "token", "local_files_only", "force_download")


class MatildaDecisionModel(PreTrainedModel):
    """A Matilda Decision checkpoint behind `system_one(state=..., questions={...}, images=[...])`."""

    config_class = Qwen3_5Config
    base_model_prefix = "decision"
    _no_split_modules: list[str] = []

    def __init__(self, config: Qwen3_5Config):
        super().__init__(config)
        self.runtime: DecisionModel | None = None
        self.post_init()

    def _init_weights(self, module: Any) -> None:
        """Every weight comes from the checkpoint."""

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, *model_args: Any,
                        config: Qwen3_5Config | None = None, **kwargs: Any) -> MatildaDecisionModel:
        """Load a Hub repository or a local download.

        `device` (or `device_map` naming one device) picks the device, default cuda:0 when a GPU is visible.
        `max_image_pixels` caps each image's resized size (default 262,144 = 0.26 MP, as trained).
        Numerics are fixed by the checkpoint (BF16 backbone, FP32 readout).
        """
        hub = {key: kwargs.pop(key) for key in HUB_OPTIONS if key in kwargs}
        device = kwargs.pop("device", None)
        device_map = kwargs.pop("device_map", None)
        if isinstance(device_map, dict):
            if set(device_map) != {""}:
                raise ValueError("Matilda Decision runs on one device: pass a device name or {'': device}")
            device_map = device_map[""]
        if device is None and device_map not in (None, "auto"):
            device = f"cuda:{device_map}" if isinstance(device_map, int) else str(device_map)
        max_image_pixels = int(kwargs.pop("max_image_pixels", MAX_IMAGE_PIXELS))
        for key in ("dtype", "torch_dtype"):
            if kwargs.pop(key, None) not in (None, "auto", "bfloat16", torch.bfloat16):
                raise ValueError(f"{key}: numerics are fixed by the checkpoint (BF16 backbone, FP32 readout)")
        local = Path(os.fspath(pretrained_model_name_or_path)).expanduser()
        if not local.is_dir():
            from huggingface_hub import snapshot_download

            local = Path(snapshot_download(str(pretrained_model_name_or_path),
                                           **{key: value for key, value in hub.items() if value is not None}))
        if config is None:
            config = cls.config_class.from_pretrained(str(local))
        model = cls(config)
        model.runtime = DecisionModel(local, device=device, max_image_pixels=max_image_pixels)
        model.name_or_path = str(pretrained_model_name_or_path)
        model.eval()
        return model

    def _require(self) -> DecisionModel:
        if self.runtime is None:
            raise RuntimeError("Load the model with from_pretrained")
        return self.runtime

    @property
    def decision_config(self) -> dict[str, Any]:
        return self._require().config

    def system_one(self, *, state: Any, questions: dict[str, Any], images: list[Any] | None = None,
                   batch_size: int = 64, token_budget: int = 131072, max_tokens: int = 131072) -> dict[str, Any]:
        """Typed choice / noul / score answers about one state: `{"model", "answers", "usage"}`.

        `images`: any number of images every question sees (PIL images, local paths, http(s) URLs or base64
        `data:image/...` URLs), placed before the text and resized to at most 0.26 MP each. Nothing is truncated:
        a request longer than `max_tokens` raises an error.
        """
        return system_one(self._require(), state=state, questions=questions, images=images,
                          max_tokens=max_tokens, token_budget=token_budget, batch_size=batch_size)

    def forward(self, state: Any = None, questions: dict[str, Any] | None = None,
                images: list[Any] | None = None) -> dict[str, Any]:
        return self.system_one(state=state, questions=questions or {}, images=images)

    def to(self, *args: Any, **kwargs: Any) -> MatildaDecisionModel:
        device, dtype, _, _ = torch._C._nn._parse_to(*args, **kwargs)
        if dtype is not None:
            raise TypeError("Numerics are fixed by the checkpoint; only the device can change")
        if device is not None:
            runtime = self._require()
            runtime.to(device)
            runtime.device_name = str(device)
        return self

    def train(self, mode: bool = True) -> MatildaDecisionModel:
        if mode:
            raise RuntimeError("Matilda Decision runs in inference mode only")
        return super().train(False)

    def save_pretrained(self, *args: Any, **kwargs: Any) -> None:
        raise NotImplementedError("The repository itself is the package; copy it with huggingface_hub.snapshot_download")