"""Deployment API: the same `decide()` over an ONNX model (numpy + onnxruntime
only, no torch) or a PyTorch checkpoint.

    from minijev import MiniJEV
    model = MiniJEV.load("model.onnx")
    decision = model.decide(state=state, question="Choose the next move",
                            choices=["UP", "DOWN", "LEFT", "RIGHT"])
    decision.choice, decision.confidence, decision.scores
"""
from __future__ import annotations

import json
from functools import lru_cache
from pathlib import Path
from typing import Any, Sequence

import numpy as np

from minijev.keys import STATE_INPUTS
from minijev.tokenizers import ByteTokenizer
from minijev.tasks import Task, get_task
from minijev.types import BooleanDecision, Decision

META_KEY = "minijev"


class Decider:
    """Common front-end: tokenisation, state encoding, Decision objects.

    `run` returns *calibrated* probabilities/logits (logits / T); the raw, uncalibrated
    distribution is softmax(logits * T) and is reported next to it in every Decision.
    """

    task: Task
    tokenizer = ByteTokenizer()
    max_question_len: int = 32
    max_choice_len: int = 16
    temperature: float = 1.0

    # -- to implement: probs, logits, embedding from model-ready arrays
    def _run(self, state: np.ndarray, question_ids: np.ndarray, choice_ids: np.ndarray):
        raise NotImplementedError

    @lru_cache(maxsize=256)
    def _tokens(self, question: str, choices: tuple[str, ...]):
        dyn = self.tokenizer.kind != "byte"
        q = self.tokenizer.encode_batch([question], self.max_question_len, dynamic=dyn)
        c = self.tokenizer.encode_batch(list(choices), self.max_choice_len, dynamic=dyn)[None]
        return q, c

    def encode(self, states: Sequence[Any]) -> np.ndarray:
        enc = self.task.encode_states(states)
        return next(iter(enc.values()))

    def probs(self, states: Sequence[Any], question: str, choices: Sequence[str],
              encoded: np.ndarray | None = None) -> np.ndarray:
        return self.run(states, question, choices, encoded)[0]

    def run(self, states, question, choices, encoded=None):
        q, c = self._tokens(question, tuple(choices))
        x = self.encode(states) if encoded is None else encoded
        B = x.shape[0]
        return self._run(x, np.repeat(q, B, 0), np.repeat(c, B, 0))

    def decide_batch(self, states: Sequence[Any], question: str, choices: Sequence[str],
                     return_embedding: bool = False) -> list[Decision]:
        if len(set(choices)) != len(choices):
            raise ValueError(f"duplicate choices: {list(choices)}")
        probs, logits, emb = self.run(states, question, choices)
        raw = _softmax(logits * self.temperature)
        out = []
        for i in range(len(probs)):
            k = int(probs[i].argmax())
            out.append(Decision(choice=choices[k], index=k, confidence=float(probs[i, k]),
                                scores={ch: float(p) for ch, p in zip(choices, probs[i])},
                                logits=[float(v) for v in logits[i]],
                                embedding=emb[i].tolist() if return_embedding else None,
                                raw_scores={ch: float(p) for ch, p in zip(choices, raw[i])},
                                temperature=self.temperature))
        return out

    def decide(self, state: Any, question: str | None = None, choices: Sequence[str] | None = None,
               return_embedding: bool = False) -> Decision:
        question = question if question is not None else self.task.default_question
        choices = list(choices) if choices is not None else list(self.task.default_choices)
        return self.decide_batch([state], question, choices, return_embedding)[0]

    def choice(self, state: Any, question: str, choices: Sequence[str]) -> str:
        return self.decide(state, question, choices).choice

    def boolean(self, state: Any, question: str, labels: tuple[str, str] = ("yes", "no")) -> BooleanDecision:
        """Yes/no question, answered internally as a 2-choice decision."""
        d = self.decide(state, question, list(labels))
        p_yes = d.scores[labels[0]]
        return BooleanDecision(answer=p_yes >= 0.5, confidence=max(p_yes, 1 - p_yes), p_yes=p_yes,
                               raw_p_yes=d.raw_scores[labels[0]])


def _softmax(x: np.ndarray) -> np.ndarray:
    z = x - x.max(-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(-1, keepdims=True)


class OnnxDecider(Decider):
    def __init__(self, path: str | Path, threads: int = 1, providers: list[str] | None = None):
        import onnxruntime as ort

        so = ort.SessionOptions()
        so.intra_op_num_threads = threads
        so.inter_op_num_threads = 1
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.session = ort.InferenceSession(str(path), so, providers=providers or ["CPUExecutionProvider"])
        meta = json.loads(self.session.get_modelmeta().custom_metadata_map[META_KEY])
        self.meta = meta
        self.task = get_task(meta["task"], meta.get("task_cfg", {}))
        self.max_question_len = meta["max_question_len"]
        self.max_choice_len = meta["max_choice_len"]
        self.state_input = meta["state_input"]
        self.path = str(path)

    def _run(self, state, question_ids, choice_ids):
        probs, logits, emb = self.session.run(None, {self.state_input: state, "question_ids": question_ids,
                                                     "choice_ids": choice_ids})
        return probs, logits, emb


class TorchDecider(Decider):
    def __init__(self, checkpoint: str | Path | None = None, model=None, cfg: dict | None = None,
                 device: str = "cpu", threads: int | None = None, task: str | None = None):
        import torch

        from minijev.choices import EmbeddingCache
        from minijev.tokenizers import make_tokenizer
        from minijev.training.supervised import make_model, make_task_for, task_specs

        if threads:
            torch.set_num_threads(threads)
        task_T = {}
        if model is None:
            ck = torch.load(checkpoint, map_location=device, weights_only=False)
            cfg = ck["config"]
            task_T = ck.get("task_temperatures") or {}
            model = make_model(cfg)
            model.load_state_dict(ck["model"])
        self.specs = {s["key"]: s["task"] for s in task_specs(cfg)}
        self.tokenizer = make_tokenizer(cfg["model"].get("text_encoder"))
        self.torch, self.model, self.cfg, self.device = torch, model.to(device).eval(), cfg, device
        self.max_question_len = cfg["model"].get("max_question_len", 32)
        self.max_choice_len = cfg["model"].get("max_choice_len", 16)
        self.global_temperature = float(model.log_temperature.exp())
        self.task_temperatures = task_T
        self.cache = EmbeddingCache() if model.backbone is not None and model.backbone.frozen else None
        self._make_task = lambda key: make_task_for(cfg, self.specs[key])
        self.use_task(task or next(iter(self.specs)))

    @classmethod
    def from_pretrained(cls, path_or_repo: str | Path, device: str = "cpu", threads: int | None = None,
                        task: str | None = None) -> "TorchDecider":
        """Load a folder written by `minijev export-hf` (local path or Hugging Face repo id)."""
        from minijev.hub import load_pretrained

        model, cfg, task_T = load_pretrained(path_or_repo, device)
        dec = cls(model=model, cfg=cfg, device=device, threads=threads, task=task)
        dec.task_temperatures = task_T
        return dec.use_task(dec.task_key)

    def use_task(self, key: str) -> "TorchDecider":
        """Switch the state format (and per-task temperature) for a multi-task model."""
        self.task_key = key
        self.task = self._make_task(key)
        self.state_key = STATE_INPUTS[self.task.state_encoder]
        self.temperature = float(self.task_temperatures.get(key, self.global_temperature))
        return self

    def _bb(self, texts: list[str], max_len: int):
        torch = self.torch
        dev = self.device

        def compute(missing):
            ids = torch.as_tensor(self.tokenizer.encode_batch(missing, max_len, dynamic=True), device=dev)
            return self.model.backbone(ids)
        tok, mask, pooled = self.cache.get_many(texts, compute)
        return tok.to(dev), mask.to(dev), pooled.to(dev)

    def forward(self, states, question: str, choices: Sequence[str], encoded=None):
        torch = self.torch
        x = self.encode(states) if encoded is None else encoded
        B = x.shape[0]
        state = {self.state_key: torch.as_tensor(x, device=self.device)}
        with torch.inference_mode():
            if self.cache is not None:  # frozen backbone: each question / label is encoded once, ever
                qt, qm, qp = self._bb([question], self.max_question_len)
                _, _, cp = self._bb(list(choices), self.max_choice_len)
                cm = torch.ones(B, len(choices), dtype=torch.bool, device=self.device)
                return self.model(state, question_bb=(qt.expand(B, -1, -1), qm.expand(B, -1), qp.expand(B, -1)),
                                  choice_bb=(cp[None].expand(B, -1, -1), cm))
            q, c = self._tokens(question, tuple(choices))
            return self.model(state, torch.as_tensor(np.repeat(q, B, 0), device=self.device),
                              torch.as_tensor(np.repeat(c, B, 0), device=self.device))

    def run(self, states, question, choices, encoded=None):
        torch = self.torch
        out = self.forward(states, question, choices, encoded)
        logits = out["logits"] / self.temperature
        probs = torch.softmax(logits, -1)
        return probs.cpu().numpy(), logits.cpu().numpy(), out["embedding"].float().cpu().numpy()

    def _run(self, state, question_ids, choice_ids):  # kept for callers passing model-ready arrays
        torch = self.torch
        with torch.inference_mode():
            out = self.model({self.state_key: torch.as_tensor(state, device=self.device)},
                             torch.as_tensor(question_ids, device=self.device),
                             torch.as_tensor(choice_ids, device=self.device))
            logits = out["logits"] / self.temperature
        return torch.softmax(logits, -1).cpu().numpy(), logits.cpu().numpy(), out["embedding"].float().cpu().numpy()

    def score(self, state: Any, question: str) -> float:
        """Continuous 0..1 score from the scalar head (meaningful once trained with `score` targets)."""
        out = self.forward([state], question, ["yes", "no"])
        return float(self.torch.sigmoid(out["score"])[0])


class MiniJEV:
    """Entry point: `MiniJEV.load(path)` for .onnx or .pt files."""

    @staticmethod
    def load(path: str | Path, threads: int = 1, **kw) -> Decider:
        p = Path(path)
        if p.suffix == ".onnx":
            return OnnxDecider(p, threads=threads, **kw)
        if p.suffix == ".pt":
            return TorchDecider(p, threads=threads, **kw)
        return TorchDecider.from_pretrained(path, threads=threads, **kw)  # export-hf folder or hub repo id