"""Load a DecisionNode Edge model directory and answer requests.

    from decisionnode_edge import EdgeModel

    model = EdgeModel.load("path/to/decisionnode-edge")
    response = model.decide({"state": "...", "questions": {"route": {"type": "choice", "criteria": {...}}}})

The model reads the prompt once per question and never generates text. Its answer is read at the
<|answer|> position: the last hidden state there is scored against a small readout table
(readout.safetensors), one row per answer code, and only the rows of the question's codes are read.
Calibration (calibration.json) turns those scores into probabilities.

The backbone is read by torch (transformers) or by a faster backend (backends.py: vLLM, OpenVINO, MLX,
llama.cpp); the readout and calibration are this module's in every case.
"""

from __future__ import annotations

import base64
import io
import json
import logging
import os
from typing import Any

import torch
from safetensors.torch import load_file

from . import answers, backends, request as requests
from .template import Template

CONFIG_FILE = "decisionnode_edge.json"

log = logging.getLogger("decisionnode_edge")


def load_backbone(path: str, device: str | None, dtype: str | torch.dtype | None) -> tuple[Any, bool]:
    """The transformers backbone on `device` (default CUDA when present, else the CPU) in `dtype` (default bfloat16 on
    CUDA, float32 elsewhere), and whether it runs under bf16 autocast."""
    from transformers import AutoModelForImageTextToText

    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    if dtype is None:
        dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32
    dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype
    backbone = AutoModelForImageTextToText.from_pretrained(path, dtype=dtype).to(device).eval()
    # bf16 weights run under bf16 autocast, the reference numerics: some operations then compute in float32
    return backbone, dtype == torch.bfloat16


class EdgeModel:
    def __init__(self, backbone: Any, tokenizer: Any, readout: torch.Tensor, config: dict, calibration: dict, path: str, autocast: bool,
                 runner: Any = None, torch_options: dict | None = None):  # fmt: skip
        self.backbone = backbone  # the transformers model; None until torch is needed when a faster backend reads the prompts
        self.tokenizer = tokenizer
        self.readout = readout
        self.config = config
        self.calibration = calibration
        self.path = path
        self.autocast = autocast
        self.runner = runner  # the backend object (backends.py), or None: torch reads every prompt
        self.torch_options = torch_options or {}  # device and dtype for torch when it is loaded on first use
        self.template = Template(tokenizer, bytes.fromhex(config["template"]["code_seed_hex"]).decode("utf-8"))
        if self.template.first_id != config["readout"]["first_token_id"]:
            raise ValueError("the tokenizer and the readout table disagree on the ids of the added tokens")
        self.limits = config["limits"]
        self._images = None

    @classmethod
    def load(cls, path: str, device: str | None = None, dtype: str | torch.dtype | None = None, backend: str | None = None,
             artifacts: str | None = None, options: dict | None = None) -> "EdgeModel":  # fmt: skip
        """Load the model directory.

        `backend`: "auto" (the default, or DECISIONNODE_EDGE_BACKEND), "torch", "vllm", "openvino", "mlx" or "llamacpp".
        Auto takes the fastest that is installed and can run here (backends.ORDER). `artifacts`: the directory the build
        tool wrote for openvino or llamacpp (or DECISIONNODE_EDGE_ARTIFACTS); for either asked for by name it defaults to
        the model directory when that holds a build (the GGUF that ships with the model). `options`: the backend's settings (the
        README; DECISIONNODE_EDGE_OPTIONS as JSON is read under them).

        `device` defaults to CUDA when present, else the CPU, and `dtype` to bfloat16 on CUDA (the reference
        precision) and float32 elsewhere; both are torch's, also when torch only reads the prompts a
        faster backend cannot take."""
        from transformers import AutoTokenizer

        with open(os.path.join(path, CONFIG_FILE), encoding="utf-8") as f:
            config = json.load(f)
        with open(os.path.join(path, config["calibration"]), encoding="utf-8") as f:
            calibration = json.load(f)
        tokenizer = AutoTokenizer.from_pretrained(path)
        table = load_file(os.path.join(path, config["readout"]["file"]))[config["readout"]["tensor"]]
        artifacts = artifacts or os.environ.get(backends.ENV_ARTIFACTS) or None
        if artifacts is None and (backend or os.environ.get(backends.ENV_BACKEND) or "").strip().lower() in ("openvino", "llamacpp") and backends.read_build(path):
            artifacts = path  # a build that ships in the model folder (the GGUF), for a backend asked for by name
        options = {**backends.env_options(), **(options or {})}
        name = backends.choose(backend, device, artifacts, options)
        log.info("the %s backend reads the prompts", name)
        if name == "torch":
            backbone, autocast = load_backbone(path, device, dtype)
            return cls(backbone, tokenizer, table.float().to(backbone.device), config, calibration, path, autocast)
        runner = backends.load(name, path, config, artifacts, options, device)
        # the readout and calibration stay here, in float32 on the CPU
        return cls(None, tokenizer, table.float(), config, calibration, path, False, runner=runner, torch_options={"device": device, "dtype": dtype})

    @property
    def device(self) -> torch.device:
        return self.readout.device

    @property
    def backend(self) -> str:
        """The backend that reads the prompts ("torch" also reads, on any backend, the pictures and the prompts longer than
        the backend takes)."""
        return self.runner.name if self.runner is not None else "torch"

    # -- public -----------------------------------------------------------

    def decide(self, body: dict) -> dict:
        """Answer a `/v1/decide` request body. Returns {"model", "answers", "usage"} as the hosted API does."""
        parsed = requests.parse(body)
        probs, input_tokens = self._predict(parsed)
        out = {}
        for q in parsed.questions:
            p = probs[q.id]
            if q.type == "choice":
                out[q.id] = answers.choice_answer(q.keys, p)
            elif q.type == "score":
                out[q.id] = answers.score_answer(q.levels, p)
            elif q.type == "truth":
                out[q.id] = answers.truth_answer(p)
            else:
                out[q.id] = answers.number_answer(q.grid, p)
        return {"model": self.config["name"], "answers": out, "usage": {"input_tokens": input_tokens, "output_tokens": 0}}

    def predict(self, body: dict) -> dict[str, list[float]]:
        """The unrounded calibrated probabilities of every question: per option or level in request order, [yes, no]
        for truth, per grid value for number."""
        return self._predict(requests.parse(body))[0]

    # -- internals --------------------------------------------------------

    def _predict(self, parsed: requests.Request) -> tuple[dict[str, list[float]], int]:
        images = self._load_images(parsed.images) if parsed.images else None
        image_ids = [image["id"] for image in parsed.images]
        out: dict[str, list[float]] = {}
        state_tokens, total = None, 0
        for q in parsed.questions:
            answered = []
            for row in q.rows:
                encoded = self.template.encode(row, parsed.state, image_ids, images[2] if images else None)
                if len(encoded.input_ids) > self.limits["max_input_tokens"]:
                    raise requests.RequestError(f"question {q.id!r} needs {len(encoded.input_ids):,} tokens; the maximum context length is {self.limits['max_input_tokens']:,} tokens",
                                                error_type="max_tokens_exceeded")  # fmt: skip
                logits = self._logits(encoded, images)
                answered.append(answers.calibrated(logits, row["type"], self.calibration, bool(image_ids)))
                # the state and its pictures count once, every question's own tokens on top (as the hosted usage counts)
                total += len(encoded.input_ids) - (encoded.state_tokens if state_tokens is not None else 0)
                state_tokens = encoded.state_tokens
            if q.type == "number":
                out[q.id] = [1.0] if not q.rows else requests.compose(q.grid, answered)
            else:
                out[q.id] = answered[0] if answered else [1.0]
        return out, total

    @torch.inference_mode()
    def _logits(self, encoded: Any, images: tuple | None) -> list[float]:
        if self.runner is not None and images is None and self.runner.accepts(len(encoded.input_ids)):
            h = self.runner.hidden(encoded.input_ids).to(self.device, torch.float32)
        else:
            h = self._torch_hidden(encoded, images).to(self.device)
        scores = h @ self.readout.t()  # one score per added token, never the vocabulary
        rows = torch.tensor(encoded.class_ids, device=self.device) - self.template.first_id
        return scores[rows].tolist()

    def _torch_hidden(self, encoded: Any, images: tuple | None) -> torch.Tensor:
        """The hidden state at the <|answer|> position (the last token) from the transformers backbone, float32."""
        if self.backbone is None:
            log.warning("loading torch for the prompts the %s backend does not take (pictures, longer inputs)", self.backend)
            self.backbone, self.autocast = load_backbone(self.path, self.torch_options.get("device"), self.torch_options.get("dtype"))
        device = self.backbone.device
        input_ids = torch.tensor([encoded.input_ids], device=device)
        extra = {}
        if images is not None:
            extra = {
                "pixel_values": images[0].to(device),
                "image_grid_thw": images[1].to(device),
                "mm_token_type_ids": (input_ids == self.template.image_pad).long(),
            }
        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=self.autocast):
            hidden = self.backbone.model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids), use_cache=False, **extra).last_hidden_state
        return hidden[0, -1].float()

    def _load_images(self, images: list[dict]) -> tuple[torch.Tensor, torch.Tensor, dict[str, int]]:
        """(pixel values, grid, placeholder tokens per image id) of a request's pictures: EXIF orientation applied,
        RGB, the longest side at most max_image_side, then the backbone's own image processor."""
        from PIL import Image, ImageOps
        from transformers import AutoImageProcessor

        if self._images is None:
            self._images = AutoImageProcessor.from_pretrained(self.path)
        max_side = int(self.limits["max_image_side"])
        fitted = []
        for image in images:
            source = image["path"] if image.get("path") is not None else io.BytesIO(base64.b64decode(image["data"]))
            try:
                with Image.open(source) as im:
                    im = ImageOps.exif_transpose(im).convert("RGB")
                    if max(im.size) > max_side:
                        scale = max_side / max(im.size)
                        im = im.resize((max(1, round(im.width * scale)), max(1, round(im.height * scale))), Image.BICUBIC)
                    fitted.append(im.copy())
            except (OSError, ValueError) as error:
                raise requests.RequestError(f"Image {image['id']!r} could not be read: {error}") from None
        out = self._images(images=fitted, return_tensors="pt")
        merge = int(getattr(self._images, "merge_size", 2))
        counts = (out["image_grid_thw"].prod(dim=-1) // merge**2).tolist()
        return out["pixel_values"], out["image_grid_thw"], {image["id"]: int(n) for image, n in zip(images, counts)}