"""One decision request -> probability distributions (same scoring as the kev benchmark engine).

Every question of a request is scored against the same state in token-budgeted, length-sorted
batches, one forward pass per batch. Nothing is truncated: a request that does not fit (too long,
too many options, or out of GPU memory) raises CapacityError, which the server maps to HTTP 413.
"""

import json
from collections.abc import Mapping, Sequence
from typing import Any, cast

import torch

from maincode_jev_serve.model import MAX_OPTIONS, DecisionModel, options
from maincode_jev_serve.types import DecisionInput, ImageInput


class CapacityError(ValueError):
    pass


def length_estimate(row: Mapping[str, Any]) -> int:
    """Cheap prompt-length estimate used only to order and pack batches (same as kev.train)."""
    measured = row.get("source", {}).get("input_tokens")
    if isinstance(measured, int) and not isinstance(measured, bool):
        return measured + 16
    return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192 + 512 * len(row.get("images", []))


def microbatches(rows: Sequence[dict[str, Any]], batch_size: int, token_budget: int) -> list[list[dict[str, Any]]]:
    result: list[list[dict[str, Any]]] = []
    pending: list[dict[str, Any]] = []
    longest = 0
    for row in rows:
        length = length_estimate(row)
        if pending and (len(pending) == batch_size or max(longest, length) * (len(pending) + 1) > token_budget):
            result.append(pending)
            pending, longest = [], 0
        pending.append(row)
        longest = max(longest, length)
    if pending:
        result.append(pending)
    return result


def decide(model: DecisionModel, state: object, questions: Mapping[str, Mapping[str, Any]], *, temperature: float,
           max_tokens: int, token_budget: int, batch_size: int, images: Sequence[ImageInput] = ()) -> tuple[dict[str, list[float]], int]:
    """Normalized probabilities per question key (in option order) and the input tokens used."""
    rows: list[dict[str, Any]] = []
    for key, question in questions.items():
        count = len(options(cast(Any, question))[0])
        if count > MAX_OPTIONS:
            raise CapacityError(f"at most {MAX_OPTIONS} options per choice question are supported ({count} given)")
        row: dict[str, Any] = {"state": state, "question": dict(question), "id": key, "source": {}}
        if images:
            row["images"] = list(images)
        rows.append(row)
    distributions: dict[str, list[float]] = {}
    input_tokens = 0
    for batch in microbatches(sorted(rows, key=length_estimate), batch_size, token_budget):
        try:
            prepared = model.prepare(cast(list[DecisionInput], batch), max_length=max_tokens)
        except ValueError as error:
            if "token limit" in str(error):
                raise CapacityError(f"request exceeds the maximum context length of {max_tokens} tokens") from error
            raise
        input_tokens += prepared.input_tokens
        try:
            with torch.inference_mode():
                probabilities = (model(prepared) / temperature).softmax(-1).float().cpu().tolist()
        except torch.OutOfMemoryError as error:
            del prepared
            torch.cuda.empty_cache()
            raise CapacityError("request exceeds the maximum context length that fits in GPU memory") from error
        for item, values, count in zip(batch, probabilities, prepared.counts, strict=True):
            total = sum(values[:count])
            distributions[item["id"]] = [value / total for value in values[:count]]
    return {key: distributions[key] for key in questions}, input_tokens