"""Single-forward-pass typed-decision inference for the demo Space.

Takes a `state` (free-form JSON, optionally containing embedded images) and a `questions` dict
(noul / choice / score primitives) and returns calibrated probabilities -- without generating any
text. The model never explains itself; it reads out logits for a small, fixed set of answer
tokens after one forward pass over the rendered prompt.

This is a standalone reimplementation scoped to exactly what this demo needs. It does not import
or vendor any third-party inference harness -- just `transformers`/`torch` plus the small amount
of prompt-rendering and calibration arithmetic below.
"""
from __future__ import annotations

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

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import AutoModelForCausalLM, AutoModelForImageTextToText, AutoProcessor, AutoTokenizer

REPO = "nutrientdocs/nutrient-document-decision-private"
SUBFOLDER = "merged"

SYSTEM_PROMPT = (
    "You are a document decision model. You read the Evidence and answer each Criterion "
    "by choosing exactly one of the listed options. You never explain. You answer with "
    "the single option label only."
)
IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"


# --------------------------------------------------------------------------- model + calibration

_model = None
_tokenizer = None
_processor = None
_multimodal = False
_calibration: dict[str, Any] = {"temperature": {"noul": 1.0, "choice": 1.0, "score": 1.0}, "head": None}


def load(device: str = "cuda"):
    global _model, _tokenizer, _processor, _multimodal, _calibration
    token = os.environ.get("HF_TOKEN")
    _tokenizer = AutoTokenizer.from_pretrained(REPO, subfolder=SUBFOLDER, token=token)
    try:
        _processor = AutoProcessor.from_pretrained(REPO, subfolder=SUBFOLDER, token=token)
        _multimodal = True
    except Exception as e:
        print(f"[serve] AutoProcessor failed to load, falling back to text-only: {e!r}", flush=True)
        _processor = None
        _multimodal = False
    loader = AutoModelForImageTextToText if _multimodal else AutoModelForCausalLM
    _model = loader.from_pretrained(REPO, subfolder=SUBFOLDER, dtype=torch.bfloat16, device_map=device, token=token)
    _model.eval()
    cal_path = hf_hub_download(REPO, "calibration.json", token=token)
    _calibration = json.load(open(cal_path))


# --------------------------------------------------------------------------- image handling

def _is_image_ref(x: Any) -> bool:
    return isinstance(x, dict) and x.get("type") == "image" and isinstance(x.get("source"), str)


def _has_images(state: Any) -> bool:
    if _is_image_ref(state):
        return True
    if isinstance(state, dict):
        return any(_has_images(v) for v in state.values())
    if isinstance(state, list):
        return any(_has_images(v) for v in state)
    return False


def _load_image(source: str) -> Image.Image:
    if source.startswith("data:"):
        header, _, payload = source.partition(",")
        blob = base64.b64decode(payload)
    elif source.startswith(("http://", "https://")):
        import httpx

        blob = httpx.get(source, timeout=30, follow_redirects=True).content
    else:
        blob = open(source, "rb").read()
    return Image.open(io.BytesIO(blob)).convert("RGB")


def _split_images(state: Any) -> tuple[Any, list[Image.Image]]:
    """Replace every {"type": "image", "source": ...} with a placeholder; collect the images
    in document order."""
    images: list[Image.Image] = []

    def walk(x):
        if _is_image_ref(x):
            images.append(_load_image(x["source"]))
            return IMAGE_PLACEHOLDER
        if isinstance(x, dict):
            return {k: walk(v) for k, v in x.items()}
        if isinstance(x, list):
            return [walk(v) for v in x]
        return x

    return walk(state), images


def _render_text(x: Any) -> str:
    if x is None:
        return ""
    if isinstance(x, str):
        return x
    return json.dumps(x, ensure_ascii=False, indent=2)


# --------------------------------------------------------------------------- prompt construction

def _prefix(state: Any) -> str:
    body = f"# Evidence\n{_render_text(state)}\n\n"
    return f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\n{body}"


def _branch_tail(body: str) -> str:
    return f"{body}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"


def _options_block(instructions: Any, labelled: list[tuple[str, str]], ask: str) -> str:
    lines = [f"# Criterion\n{_render_text(instructions)}\n", "# Options"]
    for label, desc in labelled:
        lines.append(f"{label}. {desc}" if desc else f"{label}.")
    lines.append(f"\n{ask}\n")
    return "\n".join(lines)


def _build_branch(qid: str, q: dict) -> tuple[str, list[str], list[Any]]:
    """Returns (branch_text, label_tokens_in_prompt_order, keys_in_same_order)."""
    kind = q["type"]
    if kind == "noul":
        criteria = q.get("criteria") or {}
        yes_text = _render_text(criteria.get("true")) or "The statement is true."
        no_text = _render_text(criteria.get("false")) or "The statement is false."
        pair = [("A", "yes: " + yes_text), ("B", "no: " + no_text)]
        body = _options_block(q["instructions"], pair, "Respond with only the letter.")
        return _branch_tail(body), ["A", "B"], [True, False]

    if kind == "choice":
        keys = list(q["criteria"].keys())
        labels = [LETTERS[i] for i in range(len(keys))]
        labelled = []
        for lab, k in zip(labels, keys):
            desc = _render_text(q["criteria"][k])
            labelled.append((lab, f"{k}: {desc}" if desc else k))
        body = _options_block(q["instructions"], labelled, "Respond with only the letter of the best option.")
        return _branch_tail(body), labels, keys

    if kind == "score":
        keys = list(range(len(q["criteria"])))
        labels = [LETTERS[i] for i in range(len(keys))]
        labelled = [(lab, f"(level {k} of {len(keys) - 1}) {_render_text(q['criteria'][k])}") for lab, k in zip(labels, keys)]
        body = _options_block(q["instructions"], labelled, "Respond with only the letter of the level that best matches.")
        return _branch_tail(body), labels, keys

    raise ValueError(f"unknown question type {kind!r}")


# --------------------------------------------------------------------------- calibration

def _features(kind: str, logits: np.ndarray, state_tokens: int) -> np.ndarray:
    z = logits - logits.max()
    p = np.exp(z) / np.exp(z).sum()
    n = len(z)
    ent = float(-(p * np.log(np.clip(p, 1e-12, 1))).sum() / math.log(max(n, 2)))
    top = np.sort(z)[::-1]
    margin = float(top[0] - top[1]) if n > 1 else 0.0
    return np.array([
        1.0, kind == "noul", kind == "choice", kind == "score",
        math.log(n), math.log(max(state_tokens, 1)) / 10.0, ent, min(margin, 20.0) / 10.0,
    ])


def _temperature(kind: str, logits: np.ndarray, state_tokens: int) -> float:
    head = _calibration.get("head")
    if head is not None:
        f = _features(kind, logits, state_tokens)
        return float(np.exp(np.clip(f @ np.asarray(head), math.log(0.2), math.log(20.0))))
    return float(_calibration.get("temperature", {}).get(kind, 1.0))


def _softmax(logits: np.ndarray, temperature: float) -> np.ndarray:
    z = logits / max(temperature, 1e-6)
    z = z - z.max()
    e = np.exp(z)
    return e / e.sum()


def _confidence(p: np.ndarray) -> float:
    n = len(p)
    if n <= 1:
        return 1.0
    p = np.clip(p, 1e-12, 1.0)
    h = -(p * np.log(p)).sum()
    return float(max(0.0, min(1.0, 1.0 - h / math.log(n))))


# --------------------------------------------------------------------------- forward pass

def _label_id(label: str) -> int:
    ids = _tokenizer.encode(label, add_special_tokens=False)
    if len(ids) != 1:
        raise ValueError(f"label {label!r} is not a single token")
    return ids[0]


@torch.inference_mode()
def _forward(prefix_text: str, images: list[Image.Image], branch_text: str) -> np.ndarray:
    """One full forward pass over prefix+branch; returns the last-token logits (fp32 numpy)."""
    text = prefix_text + branch_text
    if images and not _multimodal:
        raise RuntimeError(
            "This request includes an image, but the multimodal processor failed to load "
            "at startup (see the [serve] log line) -- refusing to silently answer from text alone."
        )
    if images:
        inputs = _processor(text=[text], images=images, return_tensors="pt")
    else:
        inputs = _tokenizer(text, return_tensors="pt")
    inputs = {k: v.to(_model.device) for k, v in inputs.items()}
    out = _model(**inputs, use_cache=False)
    return out.logits[0, -1].float().cpu().numpy()


def answer(state: Any, questions: dict[str, dict]) -> dict:
    """Mirrors the SystemOneResponse shape: {qid: {type, ...}}"""
    rendered_state, images = _split_images(state)
    prefix_text = _prefix(rendered_state)
    state_tokens = len(_tokenizer.encode(prefix_text, add_special_tokens=False))

    results = {}
    for qid, q in questions.items():
        branch_text, labels, keys = _build_branch(qid, q)
        row = _forward(prefix_text, images, branch_text)
        restricted = np.array([row[_label_id(lb)] for lb in labels])
        kind = q["type"]
        temp = _temperature(kind, restricted, state_tokens)
        probs = _softmax(restricted, temp)
        key_probs = dict(zip(keys, probs.tolist()))

        if kind == "noul":
            results[qid] = {"type": "noul", "noul": round(key_probs[True], 6)}
        elif kind == "choice":
            p = {str(k): round(v, 6) for k, v in key_probs.items()}
            best = max(p, key=p.get)
            results[qid] = {
                "type": "choice", "choice": best, "probabilities": p,
                "confidence": round(_confidence(np.array(list(p.values()))), 6),
            }
        else:  # score
            levels = sorted(key_probs.keys())
            pvec = np.array([key_probs[i] for i in levels])
            score = float((pvec * np.array(levels)).sum())
            results[qid] = {
                "type": "score", "score": round(score, 6),
                "legend": {str(i): _render_text(q["criteria"][i]) for i in levels},
                "probabilities": {str(i): round(float(key_probs[i]), 6) for i in levels},
                "confidence": round(_confidence(pvec), 6),
            }
    return results