from __future__ import annotations

import re
from collections import Counter
from dataclasses import dataclass, field
from decimal import Decimal, InvalidOperation
from typing import Any

MATCH = "match"
MISSED_FIELD = "missed_field"
HALLUCINATED_VALUE = "hallucinated_value"
INCORRECT_TRANSCRIPTION = "incorrect_transcription"
INCORRECT_NORMALIZATION = "incorrect_normalization"
FORMATTING_MISMATCH = "formatting_mismatch"
WRONG_ROW_OR_COLUMN_MAPPING = "wrong_row_or_column_mapping"
# Informational only -- log_date is the reviewer-confirmed upload date, not a page
# transcription, so a mismatch here isn't necessarily an OCR bug. Never counted
# toward accuracy.
DATE_CHECK = "date_check"

ERROR_CATEGORIES = (
    MISSED_FIELD,
    HALLUCINATED_VALUE,
    INCORRECT_TRANSCRIPTION,
    INCORRECT_NORMALIZATION,
    FORMATTING_MISMATCH,
    WRONG_ROW_OR_COLUMN_MAPPING,
)

_TIME_RE = re.compile(r"^(\d{1,2}):(\d{2})(?::(\d{2}))?$")
_TIME_TOLERANCE_MINUTES = 30
_PUNCTUATION_RE = re.compile(r"[\s./\-_,]")


def _normalize_time(text: str) -> str | None:
    match = _TIME_RE.match(text.strip())
    if not match:
        return None
    hour, minute, second = match.groups()
    return f"{int(hour):02d}:{int(minute):02d}:{int(second or 0):02d}"


def _normalize_number(text: str) -> str | None:
    try:
        value = Decimal(text.strip())
    except (InvalidOperation, ValueError):
        return None
    value = value.normalize()
    if value == value.to_integral_value():
        value = value.to_integral_value()
    return format(value, "f")


def normalize(value: Any) -> str | None:
    """Canonicalize a raw OCR/ground-truth value for equality comparison: times to
    HH:MM:SS, numbers to a canonical decimal form, everything else case-folded."""
    if value is None:
        return None
    text = str(value).strip()
    if not text:
        return None

    as_time = _normalize_time(text)
    if as_time is not None:
        return as_time

    as_number = _normalize_number(text)
    if as_number is not None:
        return as_number

    return text.lower()


def loose_equal(a: str | None, b: str | None) -> bool:
    """Equality ignoring case/punctuation/whitespace -- looser than `normalize()`,
    used only to detect formatting_mismatch after an exact-normalized compare fails."""
    if a is None or b is None:
        return False
    return _PUNCTUATION_RE.sub("", a) == _PUNCTUATION_RE.sub("", b)


_CANONICAL_TIME_RE = re.compile(r"^\d{2}:\d{2}:\d{2}$")


def _as_time(normalized: str | None) -> str | None:
    """Guards against `normalize()` falling through to a non-time canonical form
    (e.g. a malformed "10.15" time cell gets read as the number 10.15) when it's used
    for row alignment specifically -- anything not in HH:MM:SS form is unusable here."""
    if normalized is not None and _CANONICAL_TIME_RE.match(normalized):
        return normalized
    return None


def _time_to_minutes(hhmmss: str) -> int:
    hour, minute, _ = hhmmss.split(":")
    return int(hour) * 60 + int(minute)


def classify_field(predicted: dict[str, Any] | None, truth_value: Any) -> str | None:
    """Classify one field comparison. Returns None when both sides are blank (not a
    meaningful comparison, so it's excluded from scoring entirely)."""
    truth_norm = normalize(truth_value)

    pred_value_norm = normalize(predicted.get("value")) if predicted else None
    pred_ntext_norm = normalize(predicted.get("normalized_text")) if predicted else None
    pred_raw_norm = normalize(predicted.get("raw_text")) if predicted else None

    has_truth = truth_norm is not None
    has_pred = pred_value_norm is not None or pred_ntext_norm is not None or pred_raw_norm is not None

    if not has_truth and not has_pred:
        return None
    if not has_truth:
        return HALLUCINATED_VALUE
    if not has_pred:
        return MISSED_FIELD

    if truth_norm in (pred_value_norm, pred_ntext_norm):
        return MATCH
    if truth_norm == pred_raw_norm:
        # the raw OCR reading was right; something downstream (normalization,
        # ocr_spec default substitution) corrupted it
        return INCORRECT_NORMALIZATION

    best_predicted_norm = pred_ntext_norm or pred_value_norm or pred_raw_norm
    if loose_equal(best_predicted_norm, truth_norm):
        return FORMATTING_MISMATCH

    return INCORRECT_TRANSCRIPTION


@dataclass
class FieldDiff:
    page: str
    table: str
    predicted_row: int | None
    truth_row: int | None
    field: str
    ground_truth_value: Any
    predicted_value: Any
    predicted_raw_text: Any
    predicted_normalized_text: Any
    predicted_confidence: Any
    category: str
    note: str = ""


def align_rows(
    predicted_rows: list[dict[str, Any]],
    truth_rows: list[dict[str, Any]],
) -> list[tuple[int | None, int | None]]:
    """Pair up predicted rows with truth rows. Index-aligns when counts match (the
    common case -- reviewers correct rows in place); otherwise aligns by a normalized
    `time` field where the entry type has one, within a tolerance, greedily matching
    closest times first. Unmatched predicted rows come back as (i, None); unmatched
    truth rows as (None, i)."""
    if len(predicted_rows) == len(truth_rows):
        return [(i, i) for i in range(len(predicted_rows))]

    def pred_time(row: dict[str, Any]) -> str | None:
        time_field = row.get("time")
        if not isinstance(time_field, dict):
            return None
        raw = time_field.get("value") or time_field.get("normalized_text") or time_field.get("raw_text")
        return _as_time(normalize(raw))

    def truth_time(row: dict[str, Any]) -> str | None:
        return _as_time(normalize(row.get("time")))

    pred_times = [pred_time(r) for r in predicted_rows]
    truth_times = [truth_time(r) for r in truth_rows]
    has_time_key = any(t is not None for t in pred_times) or any(t is not None for t in truth_times)

    if not has_time_key:
        overlap = min(len(predicted_rows), len(truth_rows))
        pairs = [(i, i) for i in range(overlap)]
        pairs += [(i, None) for i in range(overlap, len(predicted_rows))]
        pairs += [(None, i) for i in range(overlap, len(truth_rows))]
        return pairs

    used_truth: set[int] = set()
    pairs: list[tuple[int | None, int | None]] = []
    for pi, pt in enumerate(pred_times):
        best_ti, best_diff = None, None
        if pt is not None:
            for ti, tt in enumerate(truth_times):
                if ti in used_truth or tt is None:
                    continue
                diff = abs(_time_to_minutes(pt) - _time_to_minutes(tt))
                if best_diff is None or diff < best_diff:
                    best_diff, best_ti = diff, ti
        if best_ti is not None and best_diff <= _TIME_TOLERANCE_MINUTES:
            used_truth.add(best_ti)
            pairs.append((pi, best_ti))
        else:
            pairs.append((pi, None))

    pairs.extend((None, ti) for ti in range(len(truth_rows)) if ti not in used_truth)
    return pairs


def diff_entry_list(
    page: str,
    table: str,
    predicted_rows: list[dict[str, Any]] | None,
    truth_rows: list[dict[str, Any]] | None,
) -> list[FieldDiff]:
    predicted_rows = predicted_rows or []
    truth_rows = truth_rows or []
    diffs: list[FieldDiff] = []

    for pi, ti in align_rows(predicted_rows, truth_rows):
        predicted_row = predicted_rows[pi] if pi is not None else {}
        truth_row = truth_rows[ti] if ti is not None else {}
        # Predicted rows carry schema-internal metadata alongside real OcrField values
        # (e.g. SourceTaggedEntry.source_table: str | None, tagging which table an entry
        # came from) -- never compare those against ground truth, which only ever has
        # actual data field names.
        predicted_field_keys = {k for k, v in predicted_row.items() if v is None or isinstance(v, dict)}
        for field_name in sorted(set(truth_row) | predicted_field_keys):
            predicted = predicted_row.get(field_name)
            truth_value = truth_row.get(field_name)
            category = classify_field(predicted, truth_value)
            if category is None:
                continue
            diffs.append(
                FieldDiff(
                    page=page,
                    table=table,
                    predicted_row=pi,
                    truth_row=ti,
                    field=field_name,
                    ground_truth_value=truth_value,
                    predicted_value=(predicted or {}).get("value"),
                    predicted_raw_text=(predicted or {}).get("raw_text"),
                    predicted_normalized_text=(predicted or {}).get("normalized_text"),
                    predicted_confidence=(predicted or {}).get("confidence"),
                    category=category,
                )
            )
    return diffs


def diff_record(
    predicted_payload: dict[str, Any],
    ground_truth_page_data: dict[str, dict[str, list[dict[str, Any]]]],
) -> list[FieldDiff]:
    diffs: list[FieldDiff] = []
    for page_key, tables in ground_truth_page_data.items():
        predicted_page = predicted_payload.get(page_key) or {}
        for table, truth_rows in tables.items():
            diffs.extend(diff_entry_list(page_key, table, predicted_page.get(table), truth_rows))
    return diffs


def detect_misplacement(diffs: list[FieldDiff]) -> list[FieldDiff]:
    """Second pass over a record's diffs: if a value went missing from one table/row
    and the same value (same field name) shows up as hallucinated elsewhere on the
    record, it's not two unrelated errors -- it's one value in the wrong place. Collapses
    both into a single wrong_row_or_column_mapping finding. Best-effort: only catches
    displacement where the value survives unchanged; a garbled displaced value would
    still show up as two separate findings."""
    missed = [d for d in diffs if d.category == MISSED_FIELD]
    hallucinated = [d for d in diffs if d.category == HALLUCINATED_VALUE]
    used_hallucinated: set[int] = set()

    for missed_diff in missed:
        truth_norm = normalize(missed_diff.ground_truth_value)
        if truth_norm is None:
            continue
        for idx, hall_diff in enumerate(hallucinated):
            if idx in used_hallucinated or hall_diff.field != missed_diff.field:
                continue
            pred_norm = (
                normalize(hall_diff.predicted_value)
                or normalize(hall_diff.predicted_normalized_text)
                or normalize(hall_diff.predicted_raw_text)
            )
            if pred_norm != truth_norm:
                continue
            used_hallucinated.add(idx)
            missed_diff.category = WRONG_ROW_OR_COLUMN_MAPPING
            missed_diff.note = f"found instead at {hall_diff.page}.{hall_diff.table}[row {hall_diff.predicted_row}]"
            hall_diff.category = WRONG_ROW_OR_COLUMN_MAPPING
            hall_diff.note = f"expected at {missed_diff.page}.{missed_diff.table}[row {missed_diff.truth_row}]"
            break

    return diffs


def check_date(predicted_payload: dict[str, Any], log_date: str | None) -> FieldDiff | None:
    """Informational-only check: does any extracted page date match the reviewer-
    confirmed log_date? Not counted toward accuracy -- see DATE_CHECK."""
    if log_date is None:
        return None

    for page_key in ("left_page", "right_page", "coagulation"):
        page = predicted_payload.get(page_key)
        date_field = page.get("date") if page else None
        if date_field:
            break
    else:
        return None

    predicted_date = (date_field.get("value") or "")[:10] or None
    matches = predicted_date == log_date

    return FieldDiff(
        page=page_key,
        table="date",
        predicted_row=None,
        truth_row=None,
        field="date",
        ground_truth_value=log_date,
        predicted_value=predicted_date,
        predicted_raw_text=date_field.get("raw_text"),
        predicted_normalized_text=date_field.get("normalized_text"),
        predicted_confidence=date_field.get("confidence"),
        category=DATE_CHECK,
        note="matches log_date" if matches else "does not match log_date (informational only -- not counted toward accuracy)",
    )


@dataclass
class RecordScore:
    record_id: str
    record_type: str
    comparable_fields: int
    matched_fields: int
    category_counts: Counter = field(default_factory=Counter)
    diffs: list[FieldDiff] = field(default_factory=list)

    @property
    def accuracy(self) -> float | None:
        if self.comparable_fields == 0:
            return None
        return self.matched_fields / self.comparable_fields


def score_record(record_id: str, record_type: str, diffs: list[FieldDiff]) -> RecordScore:
    category_counts = Counter(d.category for d in diffs)
    return RecordScore(
        record_id=record_id,
        record_type=record_type,
        comparable_fields=len(diffs),
        matched_fields=category_counts.get(MATCH, 0),
        category_counts=category_counts,
        diffs=diffs,
    )


@dataclass
class AggregateScore:
    record_type: str
    records_scored: int = 0
    comparable_fields: int = 0
    matched_fields: int = 0
    category_counts: Counter = field(default_factory=Counter)

    @property
    def accuracy(self) -> float | None:
        if self.comparable_fields == 0:
            return None
        return self.matched_fields / self.comparable_fields


def aggregate_scores(scores: list[RecordScore]) -> dict[str, AggregateScore]:
    by_type: dict[str, AggregateScore] = {}
    for score in scores:
        agg = by_type.setdefault(score.record_type, AggregateScore(record_type=score.record_type))
        agg.records_scored += 1
        agg.comparable_fields += score.comparable_fields
        agg.matched_fields += score.matched_fields
        agg.category_counts.update(score.category_counts)
    return by_type
