"""Mechanical resolver for self-describing citation ref_ids (see src/ref_id.py).

A citation is a ref_id string like ``2021_..._Brake_Fluid_Check::clip::0002`` or
``manual_2021-ford-f-150::p214::0308``. Resolving one means: split on "::",
look the result up in the evidence set actually retrieved for this query, and
classify it into the three states the UI renders:

  green  = resolves to a real retrieved unit AND that unit is the one that
           actually supports the claim (matches gold, when gold is known).
  yellow = resolves to a real unit that exists in the corpus/retrieved set,
           just not the one supporting this specific claim.
  red    = does not resolve to anything retrieved at all -- fabricated.

Two ways to get yellow-vs-green, because "the one supporting this claim" means
different things with and without ground truth:

- ``classify_citation`` -- eval mode. We have the benchmark's gold_clip_id (and
  optionally gold pages), so the verdict is exact: same source as gold = green,
  a different real source = yellow, unresolvable = red. This is what
  src/eval/faithfulness.py uses against AutoRepair-VideoQA.
- ``live_classify_citation`` -- production mode. There is no gold for an
  arbitrary user query, so "supports this claim" is approximated by whether the
  cited unit's own component/task text meaningfully overlaps the claim text
  it's attached to. This is a heuristic, not a proof -- it catches the common
  failure (citing a real but off-topic unit) without needing another LLM call.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Literal

Verdict = Literal["green", "yellow", "red"]

_WORD_RE = re.compile(r"[a-z0-9一-鿿]{2,}")


@dataclass(frozen=True)
class ParsedRef:
    source: str  # clip id (video/proc) or "manual_{pdf_stem}" (manual page)
    kind: str  # "clip" | "proc" | "p{page}"
    locator: str  # disambiguating index within that source+kind


def parse_ref_id(ref_id: str) -> ParsedRef | None:
    if not ref_id or "::" not in ref_id:
        return None
    parts = ref_id.split("::")
    if len(parts) < 3 or not all(parts[:2]):
        return None
    return ParsedRef(source=parts[0], kind=parts[1], locator="::".join(parts[2:]))


def classify_citation(
    ref_id: str,
    *,
    evidence_by_ref_id: dict[str, dict],
    gold_clip_id: str | None = None,
    gold_pages: set[int] | None = None,
) -> Verdict:
    parsed = parse_ref_id(ref_id)
    if parsed is None:
        return "red"
    evidence = evidence_by_ref_id.get(ref_id)
    if evidence is None:
        return "red"

    if gold_clip_id is None and not gold_pages:
        return "green"  # no gold to compare against; "resolves" is all we can say

    if parsed.kind in ("clip", "proc"):
        return "green" if gold_clip_id is not None and parsed.source == gold_clip_id else "yellow"
    if parsed.kind.startswith("p") and gold_pages:
        page = evidence.get("page_start")
        return "green" if page in gold_pages else "yellow"
    return "yellow"


def _words(text: str) -> set[str]:
    return set(_WORD_RE.findall((text or "").lower()))


def live_classify_citation(
    ref_id: str,
    *,
    evidence_by_ref_id: dict[str, dict],
    claim_text: str,
    overlap_threshold: float = 0.12,
) -> Verdict:
    """Heuristic 3-state verdict for production queries (no gold available)."""
    parsed = parse_ref_id(ref_id)
    if parsed is None:
        return "red"
    evidence = evidence_by_ref_id.get(ref_id)
    if evidence is None:
        return "red"

    topic = " ".join(str(evidence.get(k) or "") for k in ("title", "component", "task"))
    topic_words = _words(topic)
    claim_words = _words(claim_text)
    if not topic_words or not claim_words:
        return "green"
    overlap = len(topic_words & claim_words) / len(topic_words)
    return "green" if overlap >= overlap_threshold else "yellow"