"""Score complete candidate records in one model call."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Sequence

import torch


@dataclass(frozen=True, slots=True)
class Row:
    index: int
    candidate_id: str
    targets: tuple[int, ...]
    positions: tuple[int, ...]
    prefix_tokens: int
    input_tokens: int


@dataclass(frozen=True, slots=True)
class CompiledJoint:
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    position_ids: torch.Tensor
    logits_to_keep: torch.Tensor
    rows: tuple[Row, ...]
    candidate_order: tuple[str, ...]
    padded_tokens: int
    logit_vectors: int


def _contains(values: tuple[int, ...], needle: tuple[int, ...]) -> bool:
    return len(needle) <= len(values) and any(
        values[i:i + len(needle)] == needle for i in range(len(values) - len(needle) + 1)
    )


def compile_joint(prefix: Sequence[int], candidate_ids: Sequence[str],
                  suffixes: Sequence[Sequence[int]], terminal: Sequence[int], *,
                  pad_token_id: int, max_rows: int, max_padded_tokens: int,
                  max_logit_vectors: int, max_input_tokens: int) -> CompiledJoint:
    prefix = tuple(prefix); terminal = tuple(terminal)
    ids = tuple(candidate_ids); suffixes = tuple(tuple(v) for v in suffixes)
    if not prefix or not terminal or any(type(v) is not int or v < 0 for v in prefix + terminal):
        raise ValueError("prefix and terminal must contain nonnegative token IDs")
    if not ids or len(ids) != len(set(ids)) or any(not isinstance(v, str) or not v for v in ids):
        raise ValueError("candidate IDs must be unique nonempty strings")
    if len(ids) != len(suffixes) or len(suffixes) != len(set(suffixes)):
        raise ValueError("candidate suffixes must be unique and align with candidate IDs")
    if any(prefix[-n:] == terminal[:n] for n in range(1, min(len(prefix), len(terminal)) + 1)):
        raise ValueError("prefix overlaps the declared terminal boundary")
    for suffix in suffixes:
        if (len(suffix) < len(terminal) or suffix[-len(terminal):] != terminal
                or _contains(suffix[:-len(terminal)], terminal)):
            raise ValueError("every candidate must contain the common terminal exactly at its end")
    for i, left in enumerate(suffixes):
        for j, right in enumerate(suffixes):
            if i != j and len(left) < len(right) and right[:len(left)] == left:
                raise ValueError("candidate tokenizations prefix-overlap")
    if any(any(type(v) is not int or v < 0 for v in suffix) for suffix in suffixes):
        raise ValueError("candidate suffixes must contain nonnegative token IDs")
    if len(suffixes) > max_rows:
        raise ValueError(f"compiled rows {len(suffixes)} exceed max_rows {max_rows}")
    sequences = tuple(prefix + suffix for suffix in suffixes)
    if max(map(len,sequences)) > max_input_tokens:
        raise ValueError(f"joint request exceeds max_input_tokens {max_input_tokens}; no truncation performed")
    width = max(map(len, sequences)); padded = len(sequences) * width
    if padded > max_padded_tokens:
        raise ValueError(f"compiled padded tokens {padded} exceed max_padded_tokens {max_padded_tokens}")
    positions = sorted({p for seq in sequences for p in range(len(prefix) - 1, len(seq) - 1)})
    vectors = len(sequences) * len(positions)
    if vectors > max_logit_vectors:
        raise ValueError(f"compiled logit vectors {vectors} exceed max_logit_vectors {max_logit_vectors}")
    input_ids = torch.full((len(sequences), width), pad_token_id, dtype=torch.long)
    mask = torch.zeros_like(input_ids)
    rows = []
    for i, (candidate_id, suffix, sequence) in enumerate(zip(ids, suffixes, sequences, strict=True)):
        input_ids[i, :len(sequence)] = torch.tensor(sequence)
        mask[i, :len(sequence)] = 1
        rows.append(Row(i, candidate_id, suffix,
                        tuple(range(len(prefix) - 1, len(sequence) - 1)), len(prefix), len(sequence)))
    position_ids = mask.cumsum(-1) - 1; position_ids.masked_fill_(mask == 0, 0)
    return CompiledJoint(input_ids, mask, position_ids, torch.tensor(positions), tuple(rows), ids,
                         padded, vectors)


def score_joint(model_forward: Callable[..., Any], compiled: CompiledJoint, *, device: Any) -> dict[str, Any]:
    expected_positions = compiled.attention_mask.cumsum(-1) - 1
    expected_positions.masked_fill_(compiled.attention_mask == 0, 0)
    if not torch.equal(compiled.position_ids, expected_positions):
        raise ValueError("compiled position IDs mutated before forward")
    keep = [int(value) for value in compiled.logits_to_keep.tolist()]
    with torch.inference_mode():
        output = model_forward(input_ids=compiled.input_ids.to(device),
            attention_mask=compiled.attention_mask.to(device),
            position_ids=compiled.position_ids.to(device), logits_to_keep=keep,
            use_cache=False, return_dict=True)
    logits = output.logits
    if logits.ndim != 3 or tuple(logits.shape[:2]) != (len(compiled.rows), len(keep)):
        raise ValueError("native sparse logits do not align to compiled joint positions")
    logits = logits.float()
    if not bool(torch.isfinite(logits).all()):
        raise ValueError("model logits contain non-finite values")
    columns = {int(p): i for i, p in enumerate(compiled.logits_to_keep.tolist())}
    scores = []
    for row in compiled.rows:
        terms = []
        for position, target in zip(row.positions, row.targets, strict=True):
            distribution = logits[row.index, columns[position]]
            if target >= distribution.numel(): raise ValueError("target token is outside vocabulary")
            terms.append(distribution[target] - torch.logsumexp(distribution, -1))
        scores.append(torch.stack(terms).sum())
    scores_t = torch.stack(scores); probs = torch.exp(scores_t - torch.logsumexp(scores_t, 0))
    selected = int(torch.argmax(scores_t))
    return {"selected_index": selected, "probabilities": probs.cpu().tolist(),
            "log_scores": scores_t.cpu().tolist(),
            "row_map": [{"row_index": r.index, "candidate_id": r.candidate_id,
                         "prediction_positions": list(r.positions)} for r in compiled.rows],
            "accounting": {"backbone_calls": 1, "physical_rows": len(compiled.rows),
                           "padded_tokens": compiled.padded_tokens,
                           "logit_vectors": compiled.logit_vectors, "truncated": False},
            "calibrated": False}