"""Standalone runtime for the merged Pluto-2.1 decision model."""

from __future__ import annotations

import hashlib
import json
import os
import random
from pathlib import Path
from typing import Any, Mapping, Sequence

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

SYSTEM_PROMPT = (
    "Choose exactly one supplied decision token. Return only that token. "
    "Do not explain the answer."
)
DECISION_TYPES = frozenset({"choice", "score", "noul"})


def _canonical_json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)


def _compact_criteria(mapped_criteria: Sequence[Mapping[str, Any]]):
    parsed_descriptions = []
    description_fields = None
    for criterion in mapped_criteria:
        try:
            parsed = json.loads(str(criterion["d"]))
        except json.JSONDecodeError:
            parsed = None
        if not isinstance(parsed, dict) or not all(isinstance(key, str) for key in parsed):
            description_fields = None
            break
        fields = sorted(parsed)
        if description_fields is None:
            description_fields = fields
        elif fields != description_fields:
            description_fields = None
            break
        parsed_descriptions.append(parsed)
    if description_fields is not None and len(parsed_descriptions) == len(mapped_criteria):
        rows = [
            [criterion["t"], criterion["o"], criterion["n"], *(parsed[field] for field in description_fields)]
            for criterion, parsed in zip(mapped_criteria, parsed_descriptions, strict=True)
        ]
        return ["t", "o", "n", *description_fields], rows
    return ["t", "o", "n", "d"], [
        [criterion["t"], criterion["o"], criterion["n"], criterion["d"]]
        for criterion in mapped_criteria
    ]


def render_decision_prompt(
    *,
    state: Any,
    instructions: str,
    candidates: Sequence[Mapping[str, Any]],
    decision_type: str,
    decision_tokens: Sequence[str],
    seed: int,
    row_id: str,
    prompt_schema: str,
):
    candidate_count = len(candidates)
    if not 2 <= candidate_count <= 255:
        raise ValueError("Pluto-2.1 supports 2 to 255 candidates")
    if decision_type not in DECISION_TYPES:
        raise ValueError("decision_type must be choice, score, or noul")
    if len(decision_tokens) < candidate_count:
        raise ValueError("not enough decision tokens")

    presentation_order = list(range(candidate_count))
    digest = hashlib.sha256(f"{seed}:{row_id}:candidate-order".encode()).digest()
    random.Random(int.from_bytes(digest[:8], "big")).shuffle(presentation_order)
    mapped = []
    candidate_to_slot = {}
    for slot, candidate_index in enumerate(presentation_order):
        candidate = candidates[candidate_index]
        candidate_id = str(candidate["id"])
        candidate_to_slot[candidate_id] = slot
        mapped.append(
            {
                "t": decision_tokens[slot],
                "o": candidate_index,
                "n": str(candidate["label"]),
                "d": candidate.get("description") or "",
            }
        )
    fields, criteria = _compact_criteria(mapped)
    content = _canonical_json(
        {
            "schema": prompt_schema,
            "state": state,
            "question": {
                "instructions": instructions,
                "type": decision_type,
                "criteria_fields": fields,
                "criteria": criteria,
            },
        }
    )
    return content, presentation_order, candidate_to_slot


class Pluto21:
    def __init__(self, path: str, device: str | None = "cuda"):
        self.path = Path(path)
        with open(self.path / "pluto21_config.json", encoding="utf-8") as handle:
            self.config = json.load(handle)
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            dtype=torch.bfloat16,
            device_map=device,
        ).eval()
        tokenizer_path = self.path / "tokenizer"
        self.tokenizer = AutoTokenizer.from_pretrained(
            str(tokenizer_path if tokenizer_path.is_dir() else path)
        )
        self.tokenizer.padding_side = "left"
        self.decision_tokens = [
            f"<|decision_{index:03d}|>"
            for index in range(int(self.config["decision_token_count"]))
        ]
        ids = [self.tokenizer.convert_tokens_to_ids(token) for token in self.decision_tokens]
        if any(token_id is None for token_id in ids):
            raise RuntimeError("tokenizer is missing the decision tokens")
        self.decision_token_ids = torch.tensor(ids, dtype=torch.long)

    def predict(
        self,
        *,
        state: Any,
        instructions: str,
        candidates: Sequence[Mapping[str, Any]],
        decision_type: str = "choice",
        seed: int = 0,
        row_id: str = "0",
    ) -> dict[str, Any]:
        content, presentation_order, candidate_to_slot = render_decision_prompt(
            state=state,
            instructions=instructions,
            candidates=candidates,
            decision_type=decision_type,
            decision_tokens=self.decision_tokens,
            seed=seed,
            row_id=row_id,
            prompt_schema=self.config["prompt_schema"],
        )
        prompt = self.tokenizer.apply_chat_template(
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": content},
            ],
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=self.config.get("decision_prompt_enable_thinking", False),
        )
        encoded = self.tokenizer(str(prompt), return_tensors="pt", truncation=False)
        device = self.model.device
        ids = encoded["input_ids"].to(device)
        mask = encoded["attention_mask"].to(device)
        with torch.inference_mode():
            logits = self.model(input_ids=ids, attention_mask=mask).logits[0, -1]
            logits = logits.index_select(0, self.decision_token_ids.to(device))
        slot_logits = logits[: len(candidates)].float()
        slot_probs = torch.softmax(slot_logits, dim=-1)
        probabilities = torch.empty_like(slot_probs)
        for slot, candidate_index in enumerate(presentation_order):
            probabilities[candidate_index] = slot_probs[slot]
        return {
            "content": content,
            "prompt": str(prompt),
            "candidate_to_slot": candidate_to_slot,
            "slot_logits": slot_logits.detach().cpu().tolist(),
            "probabilities": probabilities.detach().cpu().tolist(),
        }


__all__ = ["Pluto21", "render_decision_prompt"]