"""Per-question specialist heads on top of the Pluto backbone."""

from __future__ import annotations

import json
import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file, save_file
from transformers import AutoConfig, AutoModel, AutoTokenizer

from .prompt import option_texts
from .schema import KIND_TO_ID, Question, normalize_state, parse_questions


class TaskSpecialist(nn.Module):
    def __init__(
        self,
        qids: dict[str, int],
        n_classes: dict[str, int],
        kinds: dict[str, str],
        backbone_name: str | None = None,
        backbone_config=None,
        vocab_size: int | None = None,
    ):
        super().__init__()
        if backbone_config is None:
            attn = "eager" if backbone_name and "deberta" in backbone_name.lower() else "sdpa"
            self.backbone = AutoModel.from_pretrained(backbone_name, attn_implementation=attn)
        else:
            attn = "eager" if getattr(backbone_config, "model_type", "") in ("deberta", "deberta-v2") else "sdpa"
            self.backbone = AutoModel.from_config(backbone_config, attn_implementation=attn)
        if vocab_size is not None:
            self.backbone.resize_token_embeddings(vocab_size)
        if next(self.backbone.parameters()).dtype == torch.float16:
            self.backbone = self.backbone.float()
        h = self.backbone.config.hidden_size
        self.qids = dict(qids)
        self.n_classes = dict(n_classes)
        self.kinds = dict(kinds)
        self.heads = nn.ModuleDict({q: nn.Linear(h, n) for q, n in n_classes.items()})
        self.register_buffer("_qid_index", torch.tensor([qids[q] for q in n_classes], dtype=torch.long), persistent=False)

    def forward(self, input_ids, attention_mask, qid_ids: torch.Tensor) -> torch.Tensor:
        h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0]
        maxc = max(self.n_classes.values())
        out = torch.full((h.shape[0], maxc), -1e9, device=h.device)
        for q, idx in self.qids.items():
            m = qid_ids == idx
            if m.any():
                out[m, : self.n_classes[q]] = self.heads[q](h[m]).float()
        return out.float()


def context_text(state: str, q: Question) -> str:
    instruction = q.instructions
    if q.kind == "score":
        instruction += " Options are ordered from lowest to highest."
    opts = "\n".join(f"{i}. {t}" for i, t in enumerate(option_texts(q)))
    return (
        f"State:\n{normalize_state(state)}\n\n"
        f"Question:\n{instruction}\n\n"
        f"Options:\n{opts}"
    )


def save_specialist(model: TaskSpecialist, tokenizer, out_dir: str, extra: dict) -> None:
    os.makedirs(out_dir, exist_ok=True)
    save_file({k: v.detach().cpu() for k, v in model.state_dict().items()}, os.path.join(out_dir, "model.safetensors"))
    tokenizer.save_pretrained(out_dir)
    model.backbone.config.to_json_file(os.path.join(out_dir, "backbone_config.json"))
    cfg = {
        "format": "pluto-task-1",
        "qids": model.qids,
        "n_classes": model.n_classes,
        "kinds": model.kinds,
        "vocab_size": model.backbone.config.vocab_size,
        **extra,
    }
    with open(os.path.join(out_dir, "pluto_task_config.json"), "w", encoding="utf-8") as f:
        json.dump(cfg, f, indent=2)


def load_specialist(path: str, device: str | torch.device = "cpu") -> tuple[TaskSpecialist, dict]:
    cfg = json.load(open(os.path.join(path, "pluto_task_config.json"), encoding="utf-8"))
    backbone_cfg = AutoConfig.for_model(
        **json.load(open(os.path.join(path, "backbone_config.json"), encoding="utf-8"))
    )
    model = TaskSpecialist(
        cfg["qids"], cfg["n_classes"], cfg["kinds"], backbone_config=backbone_cfg, vocab_size=cfg["vocab_size"]
    )
    state = load_file(os.path.join(path, "model.safetensors"))
    prefix = cfg.get("tensor_prefix")
    if prefix and state and all(k.startswith(prefix) for k in state):
        state = {k[len(prefix):]: v for k, v in state.items()}
    elif not prefix and state:
        base_prefix = os.path.basename(os.path.normpath(path)) + "."
        if all(k.startswith(base_prefix) for k in state):
            state = {k[len(base_prefix):]: v for k, v in state.items()}
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing or unexpected:
        raise RuntimeError(
            f"checkpoint {path!r} did not load cleanly: missing={list(missing)[:4]} unexpected={list(unexpected)[:4]}"
        )
    model.to(device).eval()
    return model, cfg


class TaskSpecialistEngine:
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, device: str | None = None, **kwargs) -> "TaskSpecialistEngine":
        return cls(pretrained_model_name_or_path, device=device)

    def __init__(self, path: str, device: str | None = None):
        if device is None:
            device = "cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else "cpu"
        self.path = path
        self.device = torch.device(device)
        self.model, self.config = load_specialist(path, device=self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model_name = self.config.get("model_name", "pluto-typed")
        self.temperatures = self.config.get("calibration", {}).get("temperature", {})
        self.max_len = int(self.config.get("training", {}).get("max_len", 512))

    @property
    def info(self) -> dict:
        return {
            "name": self.model_name,
            "path": self.path,
            "device": str(self.device),
            "qids": sorted(self.config["qids"]),
            "temperatures": self.temperatures,
            "params": int(sum(p.numel() for p in self.model.parameters())),
        }

    def _resolve_qid(self, name: str) -> str:
        if name in self.config["qids"]:
            return name
        for cq in self.config["qids"]:
            if name.endswith("_" + cq) or cq.endswith("_" + name):
                return cq
        raise ValueError(f"question {name!r} is not supported by {self.model_name}")

    @torch.no_grad()
    def _forward_strings(self, texts: list[str], qids: list[str]) -> np.ndarray:
        enc = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=self.max_len,
            return_tensors="pt",
        )
        input_ids = enc["input_ids"].to(self.device)
        mask = enc["attention_mask"].to(self.device)
        resolved = [self._resolve_qid(q) for q in qids]
        qid_ids = torch.tensor([self.config["qids"][q] for q in resolved], dtype=torch.long, device=self.device)
        ctx = (
            torch.autocast(device_type="cuda", dtype=torch.bfloat16)
            if self.device.type == "cuda"
            else __import__("contextlib").nullcontext()
        )
        with ctx:
            logits = self.model(input_ids, mask, qid_ids)
        logits = logits.float().cpu()
        temps = torch.tensor(
            [self.temperatures.get(k, 1.0) for k in ("choice", "noul", "score")], dtype=torch.float32
        )
        kinds = torch.tensor([KIND_TO_ID[self.config["kinds"][q]] for q in resolved])
        probs = F.softmax(logits / temps[kinds].unsqueeze(-1), dim=-1)
        ns = torch.tensor([self.config["n_classes"][q] for q in resolved])
        probs = probs.masked_fill(torch.arange(probs.shape[1]).unsqueeze(0) >= ns.unsqueeze(1), 0.0)
        probs = probs / probs.sum(-1, keepdim=True).clamp_min(1e-9)
        return probs.numpy()

    def predict(self, state, questions, max_len: int | None = None, alpha: float | None = None) -> dict:
        qs = parse_questions(questions) if not isinstance(questions, list) else questions
        texts, qids = [], []
        for q in qs:
            if q.name not in self.config["qids"]:
                self._resolve_qid(q.name)
            texts.append(context_text(state, q))
            qids.append(q.name)
        probs = self._forward_strings(texts, qids)
        answers = {}
        escalate = False
        for i, q in enumerate(qs):
            p = probs[i, : q.n_options]
            top = int(np.argmax(p))
            conf = float(p[top])
            item = {
                "type": q.kind,
                "confidence": conf,
                "accept_probability": conf,
                "escalate": conf < 0.5,
                "probabilities": {k: float(v) for k, v in zip(q.keys, p)},
            }
            if q.kind == "choice":
                item["choice"] = q.keys[top]
            elif q.kind == "noul":
                item["noul"] = float(p[1])
            else:
                item["score"] = float(np.dot(np.arange(q.n_options), p))
                item["score_normalized"] = item["score"] / max(1, q.n_options - 1)
            answers[q.name] = item
            escalate = escalate or item["escalate"]
        return {
            "answers": answers,
            "escalate": escalate,
            "routing": {"model": self.model_name, "device": str(self.device)},
        }

    def predict_many(self, items, batch_size: int = 64, max_len: int | None = None) -> list[dict]:
        items = list(items)
        texts = [context_text(s, q) for s, q in items]
        qids = [q.name for _, q in items]
        out = []
        for start in range(0, len(items), batch_size):
            probs = self._forward_strings(texts[start : start + batch_size], qids[start : start + batch_size])
            for row, (_, q) in enumerate(items[start : start + batch_size]):
                p = probs[row, : q.n_options]
                top = int(np.argmax(p))
                value = q.keys[top] if q.kind == "choice" else (float(p[1]) if q.kind == "noul" else float(np.dot(np.arange(q.n_options), p)))
                out.append(
                    {
                        "index": start + row,
                        "top": top,
                        "value": value,
                        "probs": p.astype(np.float32),
                        "confidence": float(p[top]),
                        "accept_probability": float(p[top]),
                    }
                )
        return out