from __future__ import annotations

import re
from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Optional, Sequence, Set, Tuple


class FailureType(str, Enum):
    """Taxonomy of verifiable hallucination/grounding failures."""

    # The claim is not supported by any provided evidence item.
    UNGROUNDED = "ungrounded"
    # Claim cites evidence ids that do not exist.
    CITATION_INVENTED = "citation_invented"
    # Claim cites evidence ids that exist but do not support the claim.
    CITATION_MISATTRIBUTED = "citation_misattributed"
    # Claim contains numbers that do not appear in the supporting evidence.
    NUMERIC_MISMATCH = "numeric_mismatch"
    # The system had no evidence set, so it cannot be verified.
    UNVERIFIABLE = "unverifiable"
    # The response contains internal contradictions (detected heuristically).
    CONTRADICTION = "contradiction"


@dataclass(frozen=True)
class EvidenceItem:
    evidence_id: str
    text: str


@dataclass(frozen=True)
class Claim:
    text: str
    # Evidence ids explicitly cited by the claim, like [E1] or (E2).
    cited_evidence_ids: Tuple[str, ...] = ()
    claim_type: str = "statement"  # statement|numeric|date|citation


@dataclass(frozen=True)
class ClaimAssessment:
    claim: Claim
    supported: bool
    support_score: float
    evidence_ids: Tuple[str, ...]
    failure_types: Tuple[FailureType, ...] = ()
    explain: str = ""


@dataclass(frozen=True)
class GroundednessSummary:
    total_claims: int
    supported_claims: int
    unsupported_claims: int
    citation_invented: int
    citation_misattributed: int
    numeric_mismatch: int
    contradiction: int
    unverifiable: bool


@dataclass(frozen=True)
class GroundednessReport:
    summary: GroundednessSummary
    assessments: Tuple[ClaimAssessment, ...]


@dataclass(frozen=True)
class GroundednessConfig:
    """Tunable, deterministic configuration.

    Note: This is not intended to be "perfect NLP". It's designed to be:
    - offline
    - deterministic
    - testable
    - good enough to drive measurable improvements
    """

    # Keep low enough to capture short factual statements like "Alpha is 3."
    min_claim_chars: int = 10
    max_claims: int = 32

    # Similarity thresholds for claim->evidence mapping.
    min_support_score: float = 0.14
    top_k_evidence: int = 3

    # Tokenization.
    min_token_len: int = 3
    stop_words: Tuple[str, ...] = (
        "the",
        "a",
        "an",
        "is",
        "are",
        "was",
        "were",
        "be",
        "been",
        "being",
        "to",
        "of",
        "in",
        "for",
        "on",
        "with",
        "at",
        "by",
        "from",
        "as",
        "into",
        "through",
        "during",
        "before",
        "after",
        "above",
        "below",
        "between",
        "under",
        "again",
        "further",
        "then",
        "once",
        "here",
        "there",
        "when",
        "where",
        "why",
        "how",
        "all",
        "each",
        "few",
        "more",
        "most",
        "other",
        "some",
        "such",
        "no",
        "nor",
        "not",
        "only",
        "own",
        "same",
        "so",
        "than",
        "too",
        "very",
        "just",
        "and",
        "but",
        "or",
        "if",
        "because",
        "until",
        "while",
        "it",
        "this",
        "that",
        "these",
        "those",
        "i",
        "you",
        "he",
        "she",
        "we",
        "they",
    )


_CITE_RE = re.compile(r"(?:\[|\()\s*(E\d{1,4})\s*(?:\]|\))")
_SENTENCE_SPLIT_RE = re.compile(r"[.!?]+\s+")
_CODEBLOCK_RE = re.compile(r"```.*?```", re.DOTALL)
_INLINE_CODE_RE = re.compile(r"`[^`]*`")
_NUM_RE = re.compile(r"\b\d+(?:\.\d+)?\b")


def analyze_groundedness(
    response_text: str,
    evidence: Sequence[EvidenceItem] | Sequence[Dict[str, str]],
    *,
    config: Optional[GroundednessConfig] = None,
) -> GroundednessReport:
    cfg = config or GroundednessConfig()
    ev_items = _normalize_evidence(evidence)
    if not ev_items:
        claims = extract_claims(response_text, config=cfg)
        assessments = tuple(
            ClaimAssessment(
                claim=c,
                supported=False,
                support_score=0.0,
                evidence_ids=(),
                failure_types=(FailureType.UNVERIFIABLE,),
                explain="No evidence set was provided; claims cannot be verified.",
            )
            for c in claims
        )
        return GroundednessReport(
            summary=_summarize(assessments, unverifiable=True),
            assessments=assessments,
        )

    claims = extract_claims(response_text, config=cfg)
    assessments = tuple(_assess_claim(c, ev_items, cfg) for c in claims)

    # Response-level contradiction heuristic: if contradictory, tag all claims (coarse but verifiable).
    if _detect_contradiction(response_text):
        updated: List[ClaimAssessment] = []
        for a in assessments:
            if FailureType.CONTRADICTION in a.failure_types:
                updated.append(a)
            else:
                updated.append(
                    ClaimAssessment(
                        claim=a.claim,
                        supported=a.supported,
                        support_score=a.support_score,
                        evidence_ids=a.evidence_ids,
                        failure_types=tuple(list(a.failure_types) + [FailureType.CONTRADICTION]),
                        explain=(a.explain + " Contradiction detected in response.").strip(),
                    )
                )
        assessments = tuple(updated)

    return GroundednessReport(
        summary=_summarize(assessments, unverifiable=False), assessments=assessments
    )


def extract_claims(text: str, *, config: Optional[GroundednessConfig] = None) -> List[Claim]:
    cfg = config or GroundednessConfig()
    cleaned = _preclean(text)

    parts = _SENTENCE_SPLIT_RE.split(cleaned)
    claims: List[Claim] = []
    for p in parts:
        s = p.strip()
        if not s:
            continue
        if len(s) < cfg.min_claim_chars:
            continue

        cited = tuple(dict.fromkeys(_CITE_RE.findall(s)).keys())
        ctype = "statement"
        if cited:
            ctype = "citation"
        elif _NUM_RE.search(s):
            ctype = "numeric"

        # Basic heuristic: avoid pure instructions/questions.
        if s.endswith("?"):
            continue

        claims.append(Claim(text=s, cited_evidence_ids=cited, claim_type=ctype))
        if len(claims) >= cfg.max_claims:
            break

    return claims


def _normalize_evidence(
    evidence: Sequence[EvidenceItem] | Sequence[Dict[str, str]],
) -> List[EvidenceItem]:
    out: List[EvidenceItem] = []
    for e in evidence:
        if isinstance(e, EvidenceItem):
            eid = str(e.evidence_id)
            txt = str(e.text)
        elif isinstance(e, dict):
            # Support eval datasets: {"id": "E1", "text": "..."}.
            eid = str(e.get("id") or e.get("evidence_id") or "")
            txt = str(e.get("text") or "")
        else:
            continue

        eid = eid.strip()
        txt = txt.strip()
        if not eid or not txt:
            continue
        out.append(EvidenceItem(evidence_id=eid, text=txt))
    return out


def _preclean(text: str) -> str:
    s = text or ""
    s = _CODEBLOCK_RE.sub(" ", s)
    s = _INLINE_CODE_RE.sub(" ", s)
    # Normalize whitespace.
    s = re.sub(r"\s+", " ", s).strip()
    return s


def _tokenize(s: str, cfg: GroundednessConfig) -> Set[str]:
    words = re.findall(r"[a-zA-Z0-9]+", (s or "").lower())
    stop = set(cfg.stop_words)
    toks = {w for w in words if len(w) >= cfg.min_token_len and w not in stop}
    return toks


def _jaccard(a: Set[str], b: Set[str]) -> float:
    if not a or not b:
        return 0.0
    inter = len(a & b)
    union = len(a | b)
    return inter / union if union else 0.0


def _rank_evidence_for_claim(
    claim: Claim, evidence: Sequence[EvidenceItem], cfg: GroundednessConfig
) -> List[Tuple[str, float]]:
    ctoks = _tokenize(claim.text, cfg)
    ranked: List[Tuple[str, float]] = []
    for e in evidence:
        etoks = _tokenize(e.text, cfg)
        score = _jaccard(ctoks, etoks)
        ranked.append((e.evidence_id, float(score)))
    ranked.sort(key=lambda x: x[1], reverse=True)
    return ranked


def _assess_claim(
    claim: Claim, evidence: Sequence[EvidenceItem], cfg: GroundednessConfig
) -> ClaimAssessment:
    ranked = _rank_evidence_for_claim(claim, evidence, cfg)
    top = ranked[: max(1, int(cfg.top_k_evidence))]
    top_ids = tuple(eid for eid, _ in top)

    evidence_by_id = {e.evidence_id: e for e in evidence}
    failures: List[FailureType] = []

    # If the claim cites evidence ids, enforce existence and support.
    if claim.cited_evidence_ids:
        missing = [eid for eid in claim.cited_evidence_ids if eid not in evidence_by_id]
        if missing:
            failures.append(FailureType.CITATION_INVENTED)
            return ClaimAssessment(
                claim=claim,
                supported=False,
                support_score=0.0,
                evidence_ids=(),
                failure_types=tuple(failures),
                explain=f"Cited evidence ids not found: {missing}.",
            )

        cited_scores = []
        for eid in claim.cited_evidence_ids:
            e = evidence_by_id[eid]
            cited_scores.append(_jaccard(_tokenize(claim.text, cfg), _tokenize(e.text, cfg)))
        support = float(max(cited_scores) if cited_scores else 0.0)

        if support < float(cfg.min_support_score):
            failures.append(FailureType.CITATION_MISATTRIBUTED)
            return ClaimAssessment(
                claim=claim,
                supported=False,
                support_score=support,
                evidence_ids=tuple(claim.cited_evidence_ids),
                failure_types=tuple(failures),
                explain="Cited evidence exists but does not support the claim by similarity threshold.",
            )

        # Numeric mismatch check: if claim has numbers, they should appear in cited evidence text.
        nums = _NUM_RE.findall(claim.text)
        if nums:
            ev_text = " ".join(evidence_by_id[eid].text for eid in claim.cited_evidence_ids)
            if any(n not in ev_text for n in nums):
                failures.append(FailureType.NUMERIC_MISMATCH)
                return ClaimAssessment(
                    claim=claim,
                    supported=False,
                    support_score=support,
                    evidence_ids=tuple(claim.cited_evidence_ids),
                    failure_types=tuple(failures),
                    explain="Claim contains numbers not present in the cited evidence.",
                )

        return ClaimAssessment(
            claim=claim,
            supported=True,
            support_score=support,
            evidence_ids=tuple(claim.cited_evidence_ids),
            failure_types=(),
            explain="Supported by cited evidence.",
        )

    # No explicit citations: pick best matching evidence.
    best_id, best_score = ranked[0] if ranked else ("", 0.0)
    best_score = float(best_score)
    supported = best_score >= float(cfg.min_support_score)
    if not supported:
        failures.append(FailureType.UNGROUNDED)
        return ClaimAssessment(
            claim=claim,
            supported=False,
            support_score=best_score,
            evidence_ids=top_ids,
            failure_types=tuple(failures),
            explain="No evidence item exceeds the support threshold.",
        )

    # Numeric mismatch check for best evidence if it looks "supported".
    nums = _NUM_RE.findall(claim.text)
    if nums:
        ev_text = evidence_by_id.get(best_id).text if best_id in evidence_by_id else ""
        if any(n not in ev_text for n in nums):
            failures.append(FailureType.NUMERIC_MISMATCH)
            return ClaimAssessment(
                claim=claim,
                supported=False,
                support_score=best_score,
                evidence_ids=top_ids,
                failure_types=tuple(failures),
                explain="Claim appears supported by similarity but contains numbers not present in the top evidence.",
            )

    return ClaimAssessment(
        claim=claim,
        supported=True,
        support_score=best_score,
        evidence_ids=top_ids,
        failure_types=(),
        explain="Supported by top evidence match.",
    )


def _detect_contradiction(text: str) -> bool:
    # Simple, deterministic contradiction check (same pattern used in core detector).
    sents = [s.strip() for s in re.split(r"[.!?]+", text or "") if s.strip()]
    if len(sents) < 2:
        return False

    negation_pairs = [
        ("is", "is not"),
        ("are", "are not"),
        ("was", "was not"),
        ("can", "cannot"),
        ("will", "will not"),
        ("do", "do not"),
        ("should", "should not"),
        ("must", "must not"),
        ("true", "false"),
        ("yes", "no"),
        ("always", "never"),
    ]

    def _tok(s: str) -> Set[str]:
        return set(re.findall(r"\b[a-z]+\b", (s or "").lower()))

    for i, a in enumerate(sents):
        for b in sents[i + 1 :]:
            al = a.lower()
            bl = b.lower()
            at = _tok(al)
            bt = _tok(bl)
            if len(at & bt) < 4:
                continue
            for pos, neg in negation_pairs:
                if (pos in al and neg in bl) or (neg in al and pos in bl):
                    return True
    return False


def _summarize(
    assessments: Sequence[ClaimAssessment], *, unverifiable: bool
) -> GroundednessSummary:
    total = len(assessments)
    supported = sum(1 for a in assessments if a.supported)
    unsupported = total - supported

    def _cnt(ft: FailureType) -> int:
        return sum(1 for a in assessments if ft in a.failure_types)

    return GroundednessSummary(
        total_claims=int(total),
        supported_claims=int(supported),
        unsupported_claims=int(unsupported),
        citation_invented=int(_cnt(FailureType.CITATION_INVENTED)),
        citation_misattributed=int(_cnt(FailureType.CITATION_MISATTRIBUTED)),
        numeric_mismatch=int(_cnt(FailureType.NUMERIC_MISMATCH)),
        contradiction=int(_cnt(FailureType.CONTRADICTION)),
        unverifiable=bool(unverifiable),
    )
