"""Decision Index kit engine (`--engine matilda_decision_engine:MatildaDecisionEngine`, repository root on PYTHONPATH).

    PYTHONPATH=. python -m decision_index pipeline --engine matilda_decision_engine:MatildaDecisionEngine \
        --option checkpoint=. --out runs/Matilda-Decision-v1-Mini

Images go in as `engine(state, questions, images=[...])`. Options: checkpoint, device, max_image_pixels
(default 262,144), max_tokens, token_budget, batch_size.
"""

import hashlib
from pathlib import Path
from typing import Any

import torch
from decision_index.engines import Engine, Unsupported

from matilda_decision_runtime import MAX_IMAGE_PIXELS, CapacityError, DecisionModel, decide


class MatildaDecisionEngine(Engine):  # type: ignore[misc]
    name = "matilda-decision"
    latency = "In-process request wall time: prompt rendering, image decoding, tokenization and one forward pass per question batch."

    def __init__(self, checkpoint: str = ".", device: str | None = None, max_image_pixels: int = MAX_IMAGE_PIXELS,
                 max_tokens: int = 131072, token_budget: int = 131072, batch_size: int = 64, **options: Any) -> None:
        super().__init__(**options)
        self.model = DecisionModel(checkpoint, device=device, max_image_pixels=int(max_image_pixels))
        self.max_tokens, self.token_budget, self.batch_size = int(max_tokens), int(token_budget), int(batch_size)
        config = Path(checkpoint) / "decision_config.json"
        self.label = str(self.model.config.get("name", "matilda-decision"))
        self.provenance = {"checkpoint": str(Path(checkpoint).resolve()), "revision": self.model.revision,
                           "checkpoint_config_sha256": hashlib.sha256(config.read_bytes()).hexdigest(),
                           "attention": self.model.attention_mode, "temperature": self.model.temperature,
                           "max_image_pixels": self.model.max_image_pixels, "max_tokens": self.max_tokens}

    def runtime(self) -> dict[str, object]:
        if not torch.cuda.is_available():
            return {"device": "cpu", "torch": torch.__version__}
        properties = torch.cuda.get_device_properties(0)
        return {"device": properties.name, "memory_gib": round(properties.total_memory / 2**30), "torch": torch.__version__}

    def synchronize(self) -> None:
        if torch.cuda.is_available():
            torch.cuda.synchronize()

    def __call__(self, state: object, questions: dict[str, dict[str, Any]], images: Any = None) -> tuple[dict[str, object], None]:
        for question in questions.values():
            if question["type"] not in ("choice", "noul"):
                raise Unsupported(f"question type {question['type']!r}")
        try:
            distributions, input_tokens = decide(self.model, state, questions, images=list(images or ()),
                                                 max_tokens=self.max_tokens, token_budget=self.token_budget,
                                                 batch_size=self.batch_size)
        except CapacityError as error:
            raise Unsupported(str(error)) from error
        answers = {key: self._answer(questions[key], values) for key, values in distributions.items()}
        return {"model": self.label, "answers": answers, "usage": {"input_tokens": input_tokens}}, None

    @staticmethod
    def _answer(question: dict[str, Any], values: list[float]) -> dict[str, object]:
        """The kit's answer shapes: noul -> P(true); choice -> argmax key plus every option's probability."""
        if question["type"] == "noul":
            return {"type": "noul", "noul": values[1]}
        keys = list(question["criteria"])
        return {"type": "choice", "choice": keys[max(range(len(keys)), key=values.__getitem__)],
                "probabilities": dict(zip(keys, values, strict=True))}