"""Versioned evaluation protocols and canonical comparison support.

The reportable token-level comparison is deliberately narrower than the raw
prediction artifact: it uses answer-token indices ``t >= 1`` with binary G/H
labels and valid finite scores. Token zero, uncertain labels, and invalid score
slots remain in artifacts so exclusions are auditable.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
from enum import Enum
import hashlib
import json
from math import isfinite
from typing import Iterable, Sequence


PREDICTION_SCHEMA_VERSION = 2


class ProtocolSuite(str, Enum):
    """Top-level evaluation suites; each method also records a protocol id."""

    CAUSAL_MATCHED = "causal_matched"
    METHOD_FAITHFUL = "method_faithful"
    OFFLINE = "offline"
    LEGACY_FAIR_V0 = "legacy_fair_v0"


@dataclass(frozen=True)
class EvaluationProtocol:
    """Machine-readable method protocol bound into predictions and reports."""

    suite: str
    protocol_id: str
    unit: str
    information_horizon: str
    label_semantics: str
    token_zero_policy: str
    uncertain_policy: str
    prediction_schema_version: int = PREDICTION_SCHEMA_VERSION

    def to_json_dict(self) -> dict:
        return asdict(self)


@dataclass(frozen=True, order=True)
class SupportEntry:
    """One canonical token in the common comparison population."""

    row_index: int
    row_id: str | None
    token_index: int
    token_id: int
    label: int

    def to_json_dict(self) -> dict:
        return {
            "row_index": int(self.row_index),
            "row_id": self.row_id,
            "token_index": int(self.token_index),
            "token_id": int(self.token_id),
            "label": int(self.label),
        }


def build_comparison_mask(
    labels: Sequence[int] | None,
    *,
    token_indices: Sequence[int] | None = None,
    score_valid_mask: Sequence[bool] | None = None,
    exclude_token_zero: bool = True,
) -> list[bool]:
    """Build the reportable G/H support mask without dropping artifact slots."""

    if labels is None:
        return []
    n = len(labels)
    indices = list(range(n)) if token_indices is None else [int(v) for v in token_indices]
    valid = [True] * n if score_valid_mask is None else [bool(v) for v in score_valid_mask]
    if len(indices) != n or len(valid) != n:
        raise ValueError("labels, token_indices, and score_valid_mask must align")
    mask = []
    for label, token_index, score_valid in zip(labels, indices, valid):
        label = int(label)
        if label not in {0, 1, 2}:
            raise ValueError("labels must contain only 0/1/2 values")
        mask.append(
            bool(score_valid)
            and label in {0, 1}
            and (not exclude_token_zero or token_index >= 1)
        )
    return mask


def support_entries(
    *,
    row_index: int,
    row_id: str | None,
    token_ids: Sequence[int],
    labels: Sequence[int],
    token_indices: Sequence[int] | None = None,
    comparison_mask: Sequence[bool] | None = None,
) -> list[SupportEntry]:
    """Construct canonical entries for one prediction row."""

    n = len(labels)
    indices = list(range(n)) if token_indices is None else [int(v) for v in token_indices]
    mask = (
        build_comparison_mask(labels, token_indices=indices)
        if comparison_mask is None
        else [bool(v) for v in comparison_mask]
    )
    if len(token_ids) != n or len(indices) != n or len(mask) != n:
        raise ValueError("token ids, labels, indices, and comparison mask must align")
    return [
        SupportEntry(row_index, row_id, token_index, int(token_id), int(label))
        for token_id, label, token_index, keep in zip(token_ids, labels, indices, mask)
        if keep
    ]


def support_hash(entries: Iterable[SupportEntry]) -> str:
    """SHA256 of deterministically serialized, canonically sorted support."""

    payload = [entry.to_json_dict() for entry in sorted(entries)]
    encoded = json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(encoded.encode("utf-8")).hexdigest()


def support_summary(entries: Iterable[SupportEntry]) -> dict:
    values = list(entries)
    labels = [entry.label for entry in values]
    return {
        "n_tokens": len(values),
        "n_grounded": labels.count(0),
        "n_hallucinated": labels.count(1),
        "positive_prevalence": labels.count(1) / len(labels) if labels else None,
        "support_sha256": support_hash(values),
    }


def validate_finite_scores(scores: Sequence[float], mask: Sequence[bool]) -> None:
    if len(scores) != len(mask):
        raise ValueError("scores and mask must align")
    bad = [index for index, (score, keep) in enumerate(zip(scores, mask)) if keep and not isfinite(float(score))]
    if bad:
        raise ValueError(f"comparison support contains non-finite scores at indices {bad[:10]}")