"""Trace diffing and classification."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Iterable


@dataclass(frozen=True)
class DiffResult:
    status: str
    category: str
    first_divergence: dict[str, Any] | None


def first_token_divergence(a: list[int], b: list[int]) -> int | None:
    n = min(len(a), len(b))
    for i in range(n):
        if a[i] != b[i]:
            return i
    if len(a) != len(b):
        return n
    return None


def diff_traces(
    base: list[dict[str, Any]],
    other: list[dict[str, Any]],
) -> DiffResult:
    if len(base) != len(other):
        return DiffResult(
            status="FAIL",
            category="GEN_CONTEXT_MISMATCH",
            first_divergence={
                "reason": "trace lengths differ",
                "left_len": len(base),
                "right_len": len(other),
            },
        )

    for idx, (left, right) in enumerate(zip(base, other)):
        if left.get("prompt_id") != right.get("prompt_id"):
            return DiffResult(
                status="FAIL",
                category="GEN_CONTEXT_MISMATCH",
                first_divergence={
                    "index": idx,
                    "reason": "prompt_id mismatch",
                    "left_prompt_id": left.get("prompt_id"),
                    "right_prompt_id": right.get("prompt_id"),
                },
            )

        if left.get("input_token_ids_hash") != right.get("input_token_ids_hash"):
            return DiffResult(
                status="FAIL",
                category="TOKENIZATION_MISMATCH",
                first_divergence={
                    "index": idx,
                    "reason": "input_token_ids_hash mismatch",
                    "left_hash": left.get("input_token_ids_hash"),
                    "right_hash": right.get("input_token_ids_hash"),
                },
            )

        if _field_mismatch(left, right, "tokenizer_id"):
            return DiffResult(
                status="FAIL",
                category="TOKENIZATION_MISMATCH",
                first_divergence={
                    "index": idx,
                    "reason": "tokenizer_id mismatch",
                    "left_tokenizer_id": left.get("tokenizer_id"),
                    "right_tokenizer_id": right.get("tokenizer_id"),
                },
            )

        decoding_fields = (
            "decoding_max_new_tokens",
            "decoding_do_sample",
            "decoding_temperature",
            "decoding_top_p",
            "decoding_top_k",
        )
        for field_name in decoding_fields:
            if _field_mismatch(left, right, field_name):
                return DiffResult(
                    status="FAIL",
                    category="GEN_CONTEXT_MISMATCH",
                    first_divergence={
                        "index": idx,
                        "reason": f"{field_name} mismatch",
                        "left_value": left.get(field_name),
                        "right_value": right.get(field_name),
                    },
                )

        divergence = first_token_divergence(
            left.get("generated_token_ids", []),
            right.get("generated_token_ids", []),
        )
        if divergence is not None:
            return DiffResult(
                status="FAIL",
                category="RUN_VARIANCE_FIXED_BATCH",
                first_divergence={
                    "index": idx,
                    "token_index": divergence,
                    "left_token": _safe_token(left, divergence),
                    "right_token": _safe_token(right, divergence),
                },
            )

        score_divergence = _first_score_divergence(left.get("scores"), right.get("scores"))
        if score_divergence is not None:
            return DiffResult(
                status="FAIL",
                category="SCORE_VARIANCE",
                first_divergence={
                    "index": idx,
                    "score_index": score_divergence,
                },
            )

    return DiffResult(status="PASS", category="PASS", first_divergence=None)


def _safe_token(row: dict[str, Any], index: int) -> int | None:
    tokens = row.get("generated_token_ids", [])
    if index < len(tokens):
        return tokens[index]
    return None


def _first_score_divergence(
    left_scores: list[float] | None,
    right_scores: list[float] | None,
) -> int | None:
    if left_scores is None or right_scores is None:
        return None
    n = min(len(left_scores), len(right_scores))
    for i in range(n):
        if left_scores[i] != right_scores[i]:
            return i
    if len(left_scores) != len(right_scores):
        return n
    return None


def _field_mismatch(left: dict[str, Any], right: dict[str, Any], field: str) -> bool:
    left_val = left.get(field)
    right_val = right.get(field)
    if left_val is None and right_val is None:
        return False
    return left_val != right_val


def aggregate_diffs(diffs: Iterable[DiffResult]) -> DiffResult:
    for diff in diffs:
        if diff.status != "PASS":
            return diff
    return DiffResult(status="PASS", category="PASS", first_divergence=None)
