"""R1 GSM8K reconstruction: explicit comparative-derived totals.

This is a narrow serving-safe reader for the first answer-changing GSM8K
reconstruction slice.  It emits a real :class:`MathProblemGraph` and admits only
after the normal solver and verifier replay the graph successfully.

Scope is deliberately small:

* one referenced source quantity;
* one comparative-derived quantity over the same canonical unit;
* a total-style question over the source + derived quantities.

Everything else returns a typed refusal reason so the caller can preserve the
existing refusal path.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Final

from generate.math_candidate_parser import _resolve_value
from generate.math_problem_graph import (
    Comparison,
    InitialPossession,
    MathGraphError,
    MathProblemGraph,
    Operation,
    Quantity,
    Unknown,
)
from generate.math_solver import SolveError, solve
from generate.math_verifier import verify
from language_packs.numerics_loader import parse_compound_cardinal


@dataclass(frozen=True, slots=True)
class R1Reconstruction:
    graph: MathProblemGraph | None
    answer: float | None
    reader_trace: tuple[str, ...]
    refusal_reason: str | None

    @property
    def is_admitted(self) -> bool:
        return self.graph is not None and self.answer is not None


@dataclass(frozen=True, slots=True)
class _Fact:
    entity: str
    value: float
    unit: str
    source_token: str
    sentence_index: int


@dataclass(frozen=True, slots=True)
class _Comparative:
    actor: str
    factor: float
    factor_token: str
    unit: str | None
    reference: str | None
    sentence_index: int
    source_span: str
    implicit_kind: str | None = None


_SENTENCE_SPLIT_RE: Final[re.Pattern[str]] = re.compile(r"(?<=[.?!])\s+")
_WORD_VALUE: Final[str] = (
    r"(?:a\s+|an\s+)?(?:one|two|three|four|five|six|seven|eight|nine|ten|"
    r"eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|eighteen|"
    r"nineteen|twenty|thirty|forty|fifty|sixty|seventy|eighty|ninety|"
    r"hundred|thousand)(?:[-\s]+(?:one|two|three|four|five|six|seven|eight|"
    r"nine|ten|eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|"
    r"eighteen|nineteen|twenty|thirty|forty|fifty|sixty|seventy|eighty|"
    r"ninety|hundred|thousand))*"
)
_VALUE: Final[str] = rf"(?:\$\d+(?:\.\d+)?|\d+(?:\.\d+)?|{_WORD_VALUE})"
_FACTOR: Final[str] = rf"(?:twice|double|triple|quadruple|{_VALUE})"
_VERB: Final[str] = (
    r"(?:has|have|had|having|is|are|was|were|cost|costs|costed|scored|"
    r"taught|teaches|made|makes|owns|owned|caught|collected|received|earned)"
)
_MULTIPLIER_WORDS: Final[frozenset[str]] = frozenset(
    {"twice", "double", "triple", "quadruple"}
)

_THERE_FACT_RE: Final[re.Pattern[str]] = re.compile(
    rf"(?:^|,\s*)there\s+(?:are|were|is|was)\s+(?P<value>{_VALUE})\s+"
    r"(?P<unit>(?:male|female|high school|new|old|large|small)?\s*[A-Za-z]+)",
    re.IGNORECASE,
)
_ENTITY_FACT_RE: Final[re.Pattern[str]] = re.compile(
    rf"(?P<entity>[A-Z][A-Za-z]*(?:\s+[A-Z][A-Za-z]*){{0,3}}|"
    r"(?:the\s+)?[a-z]+(?:\s+[a-z]+){0,2})\s+"
    rf"(?P<verb>{_VERB})\s+"
    rf"(?P<value>{_VALUE})"
    r"(?:\s+(?P<unit>[A-Za-z]+(?:\s+[A-Za-z]+)?))?",
    re.IGNORECASE,
)
_RELATIVE_FACT_RE: Final[re.Pattern[str]] = re.compile(
    rf"(?:as|than)\s+(?P<entity>[^,.?]+?),\s+who\s+{_VERB}\s+"
    rf"(?P<value>{_VALUE})\s+(?P<unit>[A-Za-z]+)",
    re.IGNORECASE,
)

_EXPLICIT_AS_RE: Final[re.Pattern[str]] = re.compile(
    rf"(?P<prefix>.+?)\b(?P<factor>{_FACTOR})\s+"
    r"(?:times\s+)?(?:as\s+many|as\s+much)\s+"
    r"(?P<unit>[A-Za-z]+(?:\s+[A-Za-z]+)?)?\s+as\s+(?P<ref>[^,.?]+)",
    re.IGNORECASE,
)
_THE_NUMBER_AS_RE: Final[re.Pattern[str]] = re.compile(
    rf"(?P<prefix>.+?)\b(?P<factor>{_FACTOR})\s+times\s+the\s+number\s+of\s+"
    r"(?P<unit>[A-Za-z]+(?:\s+[A-Za-z]+)?)\s+as\s+(?P<ref>[^,.?]+)",
    re.IGNORECASE,
)
_GREATER_THAN_RE: Final[re.Pattern[str]] = re.compile(
    rf"(?P<prefix>.+?)\b(?P<factor>{_FACTOR})\s+times\s+"
    r"(?:greater|more)(?:\s+(?P<unit>[A-Za-z]+(?:\s+[A-Za-z]+)?))?\s+(?:than|as)\s+"
    r"(?:the\s+)?(?:cost\s+of\s+)?(?P<ref>[^,.?]+)",
    re.IGNORECASE,
)
_DOUBLE_WHAT_RE: Final[re.Pattern[str]] = re.compile(
    r"(?P<prefix>.+?)\b(?P<factor>twice|double)\s+what\s+(?P<ref>[^,.?]+)",
    re.IGNORECASE,
)
_THAT_MANY_RE: Final[re.Pattern[str]] = re.compile(
    rf"(?P<prefix>.+?)\b(?P<factor>{_FACTOR})\s+times\s+that\s+many\s+"
    r"(?P<unit>[A-Za-z]+)",
    re.IGNORECASE,
)
_IMPLICIT_GENDER_RE: Final[re.Pattern[str]] = re.compile(
    rf"(?P<prefix>.+?)\b(?P<factor>{_FACTOR})\s+"
    r"(?:times\s+)?as\s+many\s+(?P<unit>male\s+students|female\s+students)",
    re.IGNORECASE,
)


def reconstruct_r1_total(problem_text: str) -> R1Reconstruction | None:
    """Attempt R1 reconstruction, returning None when no R1 signal exists."""
    sentences = _sentences(problem_text)
    if not sentences:
        return None
    statements = [s for s in sentences if not s.rstrip().endswith("?")]
    question = next((s for s in reversed(sentences) if s.rstrip().endswith("?")), "")
    if not question:
        return None
    source_clauses = [*statements, *_question_given_clauses(question)]

    comparatives = _comparatives(statements, problem_text)
    if not comparatives:
        return None
    trace: list[str] = []
    if len(comparatives) != 1:
        return _refuse("multiple_comparatives", trace)
    comparative = comparatives[0]
    trace.append(_event("matched_comparative", actor=comparative.actor, sentence_index=comparative.sentence_index))

    if not _question_targets_total(question, comparative):
        return _refuse("question_not_total_target", trace)

    facts = _facts(source_clauses, problem_text)
    ref = _bind_reference(comparative, facts)
    if ref is None:
        return _refuse("reference_not_grounded", trace)
    derived_unit = _canonical_unit(comparative.unit or ref.unit)
    if derived_unit != ref.unit:
        return _refuse("unit_mismatch", trace, derived_unit=derived_unit, reference_unit=ref.unit)

    unused = _unused_source_quantity_tokens(source_clauses, ref, comparative)
    if unused:
        return _refuse("incomplete_source_quantities", trace, unused=unused)

    try:
        graph = _graph(ref, comparative, derived_unit)
        solved = solve(graph)
        verdict = verify(graph, solved)
    except (MathGraphError, SolveError, ValueError) as exc:
        return _refuse("graph_or_solver_refused", trace, error=str(exc))
    if not verdict.passed:
        return _refuse("verifier_refused", trace, reason=verdict.reason)

    trace.append(_event("admitted", answer=solved.answer_value))
    return R1Reconstruction(
        graph=graph,
        answer=solved.answer_value,
        reader_trace=tuple(trace),
        refusal_reason=None,
    )


def _sentences(text: str) -> list[str]:
    return [s.strip() for s in _SENTENCE_SPLIT_RE.split(text.strip()) if s.strip()]


def _question_given_clauses(question: str) -> tuple[str, ...]:
    """Conditional givens inside the question, e.g. ``If X has 22 sharks, ...``."""
    match = re.match(r"\s*if\s+(.+?),\s*(?:how|what|which|who|when|where)\b", question, re.IGNORECASE)
    return (match.group(1).strip(),) if match else ()


def _facts(statements: list[str], problem_text: str) -> list[_Fact]:
    out: list[_Fact] = []
    discourse_subject = _first_proper_noun(problem_text)
    for idx, sentence in enumerate(statements):
        for match in _THERE_FACT_RE.finditer(sentence):
            if _value_is_comparative_factor(sentence, match.end("value")):
                continue
            fact = _fact_from_parts(
                entity=match.group("unit"),
                value_token=match.group("value"),
                unit=match.group("unit"),
                sentence_index=idx,
            )
            if fact is not None:
                out.append(fact)
        for match in _RELATIVE_FACT_RE.finditer(sentence):
            fact = _fact_from_parts(
                entity=match.group("entity"),
                value_token=match.group("value"),
                unit=match.group("unit"),
                sentence_index=idx,
            )
            if fact is not None:
                out.append(fact)
        for match in _ENTITY_FACT_RE.finditer(sentence):
            if _value_is_comparative_factor(sentence, match.end("value")):
                continue
            entity = _clean_entity(match.group("entity"))
            if entity.lower() in {"he", "she"} and discourse_subject is not None:
                entity = discourse_subject
            if entity.lower() in {"there", "who", "what"}:
                continue
            unit = match.group("unit")
            fact = _fact_from_parts(
                entity=entity,
                value_token=match.group("value"),
                unit=unit,
                sentence_index=idx,
            )
            if fact is not None:
                out.append(fact)
    return _dedupe_facts(out)


def _value_is_comparative_factor(sentence: str, value_end: int) -> bool:
    tail = sentence[value_end:value_end + 16].lower()
    return bool(re.match(r"\s+times\b", tail))


def _fact_from_parts(
    *,
    entity: str,
    value_token: str,
    unit: str | None,
    sentence_index: int,
) -> _Fact | None:
    value = _parse_value(value_token)
    if value is None:
        return None
    value_number, value_unit = value
    final_unit = _canonical_unit(value_unit or unit or "")
    if not final_unit:
        return None
    return _Fact(
        entity=_clean_entity(entity),
        value=value_number,
        unit=final_unit,
        source_token=_quantity_surface_token(value_token),
        sentence_index=sentence_index,
    )


def _comparatives(statements: list[str], problem_text: str) -> list[_Comparative]:
    out: list[_Comparative] = []
    discourse_subject = _first_proper_noun(problem_text)
    for idx, sentence in enumerate(statements):
        for pattern in (
            _THE_NUMBER_AS_RE,
            _EXPLICIT_AS_RE,
            _GREATER_THAN_RE,
            _DOUBLE_WHAT_RE,
            _THAT_MANY_RE,
            _IMPLICIT_GENDER_RE,
        ):
            match = pattern.search(sentence)
            if match is None:
                continue
            factor = _parse_factor(match.group("factor"))
            if factor is None:
                continue
            ref = match.groupdict().get("ref")
            unit = match.groupdict().get("unit")
            implicit_kind = None
            if pattern is _THAT_MANY_RE:
                implicit_kind = "that_many"
            elif pattern is _IMPLICIT_GENDER_RE:
                implicit_kind = "gender_pair"
            actor = _actor_from_prefix(match.group("prefix"), discourse_subject)
            if actor is None and implicit_kind in {"that_many", "gender_pair"} and unit:
                actor = _clean_entity(unit)
            if actor is None:
                continue
            out.append(
                _Comparative(
                    actor=actor,
                    factor=factor,
                    factor_token=_quantity_surface_token(match.group("factor")),
                    unit=unit,
                    reference=_clean_reference(ref) if ref else None,
                    sentence_index=idx,
                    source_span=sentence,
                    implicit_kind=implicit_kind,
                )
            )
            break
    return out


def _bind_reference(comparative: _Comparative, facts: list[_Fact]) -> _Fact | None:
    if comparative.reference:
        for fact in reversed(facts):
            if _entity_matches(fact.entity, comparative.reference):
                return fact
        return None

    if comparative.implicit_kind == "gender_pair":
        target_unit = _canonical_unit(comparative.unit or "")
        candidates = [
            f for f in facts
            if f.sentence_index < comparative.sentence_index
            and f.unit == target_unit
            and _is_gender_counterpart(f.entity, comparative.actor)
        ]
        return candidates[0] if len(candidates) == 1 else None

    if comparative.implicit_kind == "that_many":
        target_unit = _canonical_unit(comparative.unit or "")
        candidates = [
            f for f in facts
            if f.sentence_index < comparative.sentence_index and f.unit == target_unit
        ]
        return candidates[0] if len(candidates) == 1 else None
    return None


def _graph(ref: _Fact, comparative: _Comparative, unit: str) -> MathProblemGraph:
    entities = (ref.entity, comparative.actor)
    return MathProblemGraph(
        entities=entities,
        initial_state=(
            InitialPossession(ref.entity, Quantity(ref.value, unit)),
        ),
        operations=(
            Operation(
                actor=comparative.actor,
                kind="compare_multiplicative",
                operand=Comparison(
                    reference_actor=ref.entity,
                    delta=None,
                    factor=comparative.factor,
                    direction="times",
                ),
            ),
        ),
        unknown=Unknown(entity=None, unit=unit),
    )


def _parse_value(token: str) -> tuple[float, str | None] | None:
    raw = token.strip().rstrip(".,?")
    if raw.startswith("$"):
        try:
            return float(raw[1:]), "dollars"
        except ValueError:
            return None
    normalized = _strip_article(raw.lower())
    if re.fullmatch(r"\d+(?:\.\d+)?", normalized):
        return float(normalized), None
    parsed = parse_compound_cardinal(normalized)
    if parsed is not None:
        return float(parsed), None
    resolved = _resolve_value(normalized)
    if resolved is None:
        return None
    return float(resolved.value), resolved.unit_override


def _parse_factor(token: str) -> float | None:
    raw = token.strip().lower().rstrip(".,?")
    if raw in {"twice", "double"}:
        return 2.0
    if raw == "triple":
        return 3.0
    if raw == "quadruple":
        return 4.0
    parsed = _parse_value(raw)
    return parsed[0] if parsed is not None else None


def _question_targets_total(question: str, comparative: _Comparative) -> bool:
    lowered = question.lower()
    total_cue = any(
        cue in lowered
        for cue in ("total", "altogether", "combined", "together", "both", "in all")
    )
    if total_cue:
        return True
    if comparative.unit is not None and "how many" in lowered:
        q_unit = _canonical_unit(_words_after_how_many(question))
        if q_unit and q_unit == _canonical_unit(comparative.unit):
            return True
    if "how much" in lowered and any(word in lowered for word in ("spent", "cost", "accessor")):
        return True
    return False


def _words_after_how_many(question: str) -> str:
    match = re.search(r"how\s+many\s+([A-Za-z]+(?:\s+[A-Za-z]+)?)", question, re.IGNORECASE)
    return match.group(1) if match else ""


def _unused_source_quantity_tokens(
    statements: list[str],
    ref: _Fact,
    comparative: _Comparative,
) -> tuple[str, ...]:
    required: list[str] = []
    for sentence in statements:
        required.extend(_quantity_surfaces(sentence))
    consumed = {ref.source_token.lower(), comparative.factor_token.lower()}
    return tuple(
        token for token in required
        if token.lower() not in consumed
    )


def _quantity_surfaces(text: str) -> tuple[str, ...]:
    surfaces: list[str] = []
    for match in re.finditer(r"\$\d+(?:\.\d+)?|\d+(?:\.\d+)?", text):
        surfaces.append(match.group(0))
    for match in re.finditer(r"\b(twice|double|triple|quadruple)\b", text, re.IGNORECASE):
        surfaces.append(match.group(1).lower())
    word = (
        r"(?:a\s+|an\s+)?(?:one|two|three|four|five|six|seven|eight|nine|ten|"
        r"eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|eighteen|"
        r"nineteen|twenty|thirty|forty|fifty|sixty|seventy|eighty|ninety|"
        r"hundred|thousand)(?:[-\s]+(?:one|two|three|four|five|six|seven|"
        r"eight|nine|ten|eleven|twelve|thirteen|fourteen|fifteen|sixteen|"
        r"seventeen|eighteen|nineteen|twenty|thirty|forty|fifty|sixty|"
        r"seventy|eighty|ninety|hundred|thousand))*"
    )
    for match in re.finditer(rf"\b{word}\b", text, re.IGNORECASE):
        token = _quantity_surface_token(match.group(0))
        if token.lower() not in _MULTIPLIER_WORDS:
            surfaces.append(token)
    return tuple(dict.fromkeys(surfaces))


def _quantity_surface_token(token: str) -> str:
    raw = token.strip().lower().rstrip(".,?")
    raw = _strip_article(raw)
    parts = [p for p in re.split(r"[\s-]+", raw) if p and p != "and"]
    return " ".join(parts) if parts else raw


def _canonical_unit(unit: str) -> str:
    raw = re.sub(r"[^A-Za-z\s]", "", unit).lower().strip()
    if not raw:
        return ""
    tokens = tuple(t for t in raw.split() if t not in {"of", "the", "a", "an", "as"})
    if not tokens:
        return ""
    if "student" in tokens or "students" in tokens:
        return "students"
    if any(t in {"lady", "ladies", "girl", "girls", "boy", "boys", "people", "person"} for t in tokens):
        return "people"
    if "credit" in tokens or "credits" in tokens:
        return "credits"
    head = tokens[-1]
    if head.endswith("ies"):
        return head[:-3] + "ies"
    if head.endswith("s") or head in {"dice"}:
        return head
    return head + "s"


def _actor_from_prefix(prefix: str, discourse_subject: str | None) -> str | None:
    text = re.sub(r"^[,.\s]+", "", prefix.strip())
    text = re.sub(r"^(?:if|and|but|while|meanwhile)\s+", "", text, flags=re.IGNORECASE)
    cost_match = re.search(r"cost\s+of\s+(?:the\s+)?(?P<actor>[A-Za-z][A-Za-z\s]+?)\s+(?:was|is|were|are)?$", text, re.IGNORECASE)
    if cost_match:
        return _clean_entity(cost_match.group("actor"))
    there_match = re.search(r"there\s+(?:are|were|is|was)\s*$", text, re.IGNORECASE)
    if there_match:
        return None
    verb_match = re.search(rf"(?P<actor>.+?)\s+{_VERB}\s*$", text, re.IGNORECASE)
    actor_raw = verb_match.group("actor") if verb_match else text
    actor_raw = re.sub(r"^.*\bcolleague\s+", "", actor_raw, flags=re.IGNORECASE)
    actor_raw = actor_raw.strip()
    if actor_raw.lower() in {"he", "she", "his", "her"}:
        return discourse_subject
    return _clean_entity(actor_raw)


def _clean_reference(raw: str | None) -> str | None:
    if raw is None:
        return None
    text = re.split(r"\b(?:has|have|had|is|are|was|were|teaches|teach|scored)\b", raw, maxsplit=1, flags=re.IGNORECASE)[0]
    return _clean_entity(text)


def _clean_entity(raw: str) -> str:
    text = re.sub(r"[^A-Za-z0-9\s]", " ", raw)
    text = re.sub(r"\s+", " ", text).strip()
    text = re.sub(r"^(?:if|the|a|an|his|her|their|this|that)\s+", "", text, flags=re.IGNORECASE)
    text = re.sub(r"\s+(?:who|which|that)$", "", text, flags=re.IGNORECASE)
    if not text:
        return ""
    lowered = text.lower()
    if lowered in {"dad", "father"}:
        return "dad"
    if lowered in {"mom", "mother"}:
        return "mom"
    if "female student" in lowered:
        return "female students"
    if "male student" in lowered:
        return "male students"
    return " ".join(part.capitalize() if part.islower() else part for part in text.split())


def _first_proper_noun(text: str) -> str | None:
    match = re.search(r"\b([A-Z][a-z]+)\b", text)
    return match.group(1) if match else None


def _entity_matches(entity: str, reference: str) -> bool:
    e = _entity_key(entity)
    r = _entity_key(reference)
    return e == r or e.endswith(r) or r.endswith(e)


def _entity_key(entity: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", entity.lower())


def _is_gender_counterpart(source: str, actor: str) -> bool:
    s = set(source.lower().split())
    a = set(actor.lower().split())
    return (
        ("female" in s and "male" in a)
        or ("male" in s and "female" in a)
        or ("girls" in s and "boys" in a)
        or ("boys" in s and "girls" in a)
    )


def _dedupe_facts(facts: list[_Fact]) -> list[_Fact]:
    out: list[_Fact] = []
    seen: set[tuple[str, str, float, int]] = set()
    for fact in facts:
        key = (_entity_key(fact.entity), fact.unit, fact.value, fact.sentence_index)
        if key in seen:
            continue
        seen.add(key)
        out.append(fact)
    return out


def _strip_article(text: str) -> str:
    return re.sub(r"^(?:a|an)\s+", "", text.strip(), flags=re.IGNORECASE)


def _refuse(reason: str, trace: list[str], **detail: object) -> R1Reconstruction:
    trace.append(_event("refused", reason=reason, **detail))
    return R1Reconstruction(
        graph=None,
        answer=None,
        reader_trace=tuple(trace),
        refusal_reason=reason,
    )


def _event(outcome: str, **payload: object) -> str:
    data = {"layer": "r1_reconstruction", "outcome": outcome}
    data.update(payload)
    return json.dumps(data, sort_keys=True, separators=(",", ":"))
