# fitz_ai/governance/constraints/feature_extractor.py
"""
Feature Extractor for Governance Classifier.

Extracts a flat feature dict from query + chunks + constraint results for
training/inference with a tabular governance classifier. Features come from
three tiers:

- Tier 1: Signals already computed inside constraints (surfaced via metadata)
- Tier 2: Cheap computation on existing data (no LLM, no I/O)
- Tier 3: DetectionSummary from retrieval (threaded via StageContext)

The classifier uses these features to predict one of 3 governance modes:
abstain, disputed, trustworthy.
"""

from __future__ import annotations

import re
import statistics
from typing import TYPE_CHECKING, Any, Sequence

from fitz_ai.governance.protocol import EvidenceItem

if TYPE_CHECKING:
    from fitz_ai.governance.constraints.base import ConstraintResult

from fitz_ai.governance.constraints.plugins.answer_verification import AnswerVerificationConstraint
from fitz_ai.governance.constraints.plugins.causal_attribution import CausalAttributionConstraint
from fitz_ai.governance.constraints.plugins.conflict_aware import ConflictAwareConstraint
from fitz_ai.governance.constraints.plugins.insufficient_evidence import (
    InsufficientEvidenceConstraint,
)
from fitz_ai.governance.constraints.plugins.specific_info_type import SpecificInfoTypeConstraint

_CONSTRAINT_CLASSES: list[type] = [
    InsufficientEvidenceConstraint,
    ConflictAwareConstraint,
    CausalAttributionConstraint,
    SpecificInfoTypeConstraint,
    AnswerVerificationConstraint,
]

# Constraint name → class mapping for result lookup
_CONSTRAINT_NAME_MAP: dict[str, type] = {
    "insufficient_evidence": InsufficientEvidenceConstraint,
    "conflict_aware": ConflictAwareConstraint,
    "causal_attribution": CausalAttributionConstraint,
    "specific_info_type": SpecificInfoTypeConstraint,
    "answer_verification": AnswerVerificationConstraint,
}


def extract_features(
    query: str,
    chunks: Sequence[EvidenceItem],
    constraint_results: dict[str, "ConstraintResult"],
    detection_summary: Any | None = None,
) -> dict[str, Any]:
    """
    Extract a flat feature dict for the governance classifier.

    Args:
        query: The user's question
        chunks: Retrieved chunks
        constraint_results: Map of constraint_name -> ConstraintResult
        detection_summary: Optional DetectionSummary from retrieval

    Returns:
        Flat dict of feature_name -> value (numeric, bool, or string)
    """
    features: dict[str, Any] = {}

    # === Tier 1: Constraint metadata (surfaced by constraint plugins) ===
    _extract_constraint_features(features, constraint_results)

    # === Tier 2: Cheap computation on query + chunks ===
    _extract_query_features(features, query)
    _extract_chunk_features(features, query, chunks)

    # === Tier 3: DetectionSummary from retrieval ===
    _extract_detection_features(features, detection_summary)

    return features


def _extract_constraint_features(
    features: dict[str, Any],
    constraint_results: dict[str, "ConstraintResult"],
) -> None:
    """Extract features from constraint result metadata (schema-driven)."""
    # Aggregate cross-constraint signals (not owned by any single constraint)
    _extract_aggregate_signals(features, constraint_results)

    # Schema-driven per-constraint extraction
    for name, cls in _CONSTRAINT_NAME_MAP.items():
        schema = cls.feature_schema()
        result = constraint_results.get(name)

        for spec in schema:
            if spec.name.endswith("_fired"):
                # Derived from result, not metadata
                features[spec.name] = (not result.allow_decisive_answer) if result else spec.default
            elif spec.name.endswith("_signal"):
                features[spec.name] = result.signal if result else spec.default
            else:
                features[spec.name] = (
                    result.metadata.get(spec.name, spec.default) if result else spec.default
                )

    # AV-specific derived feature: strong denial (2+ jury votes)
    av_votes = features.get("av_jury_votes_no")
    features["av_strong_denial"] = (av_votes or 0) >= 2


def _extract_aggregate_signals(
    features: dict[str, Any],
    constraint_results: dict[str, "ConstraintResult"],
) -> None:
    """Extract cross-constraint aggregate features."""
    num_denials = sum(1 for r in constraint_results.values() if not r.allow_decisive_answer)
    features["num_constraints_fired"] = num_denials

    signals = [r.signal for r in constraint_results.values() if r.signal]
    features["has_qualified_signal"] = "qualified" in signals
    features["has_abstain_signal"] = "abstain" in signals
    features["has_disputed_signal"] = "disputed" in signals
    features["has_any_denial"] = num_denials > 0
    features["num_strong_denials"] = sum(
        1
        for r in constraint_results.values()
        if not r.allow_decisive_answer and r.signal in ("abstain", "disputed")
    )


_COMPARISON_WORDS = {"vs", "versus", "compare", "compared", "differ", "difference", "between"}
_QUESTION_TYPE_RE = re.compile(
    r"^(what|how|why|when|where|who|which|is|are|was|were|do|does|did|can|could|should|will|would)\b",
    re.IGNORECASE,
)


def _extract_query_features(features: dict[str, Any], query: str) -> None:
    """Extract cheap features from the query text."""
    features["query_word_count"] = len(query.split())

    # Comparison words
    q_lower_words = set(query.lower().split())
    features["query_has_comparison_words"] = bool(q_lower_words & _COMPARISON_WORDS)

    # Question type classification
    m = _QUESTION_TYPE_RE.match(query.strip())
    if m:
        word = m.group(1).lower()
        if word in ("what", "which"):
            features["query_question_type"] = "what"
        elif word == "how":
            features["query_question_type"] = "how"
        elif word == "why":
            features["query_question_type"] = "why"
        elif word in ("when", "where"):
            features["query_question_type"] = "when_where"
        elif word == "who":
            features["query_question_type"] = "who"
        else:
            features["query_question_type"] = "yes_no"
    else:
        features["query_question_type"] = "other"


def _extract_chunk_features(
    features: dict[str, Any], query: str, chunks: Sequence[EvidenceItem]
) -> None:
    """Extract cheap features from chunks and their relationship to query."""
    features["num_chunks"] = len(chunks)

    if not chunks:
        features["num_unique_sources"] = 0
        features["mean_vector_score"] = None
        features["max_vector_score"] = None
        features["min_vector_score"] = None
        features["std_vector_score"] = None
        features["score_spread"] = None
        features["vocab_overlap_ratio"] = 0.0
        # Inter-chunk defaults (including ctx_* features)
        for k in (
            "max_pairwise_overlap",
            "min_pairwise_overlap",
            "chunk_length_cv",
            "assertion_density",
            "hedge_density",
            "number_density",
            "ctx_length_mean",
            "ctx_length_std",
            "ctx_total_chars",
            "ctx_contradiction_count",
            "ctx_negation_count",
            "ctx_number_count",
            "ctx_number_variance",
            "ctx_max_pairwise_sim",
            "ctx_mean_pairwise_sim",
            "ctx_min_pairwise_sim",
            "cross_chunk_num_conflicts",
            "cross_chunk_max_divergence",
            "within_chunk_num_conflicts",
            "within_chunk_max_divergence",
        ):
            features[k] = 0.0
        features["has_cross_chunk_divergence"] = False
        features["has_within_chunk_divergence"] = False
        features["year_count"] = 0
        features["has_distinct_years"] = False
        # Text answer feature defaults
        features["query_subject_partial"] = 0.0
        features["best_span_length"] = 0.0
        features["best_sentence_coverage"] = 0.0
        features["entity_substantive_score"] = 0.0
        features["answer_span_coverage"] = 0.0
        return

    # Source diversity
    doc_ids = set()
    for c in chunks:
        doc_id = (
            getattr(c, "doc_id", None) or c.metadata.get("doc_id") or c.metadata.get("source_file")
        )
        if doc_id:
            doc_ids.add(doc_id)
    features["num_unique_sources"] = len(doc_ids) if doc_ids else len(chunks)

    # Vector score distribution
    scores = []
    for c in chunks:
        score = c.metadata.get("vector_score") or c.metadata.get("score")
        if score is not None:
            try:
                scores.append(float(score))
            except (TypeError, ValueError):
                pass

    if scores:
        features["mean_vector_score"] = statistics.mean(scores)
        features["max_vector_score"] = max(scores)
        features["min_vector_score"] = min(scores)
        features["std_vector_score"] = statistics.pstdev(scores) if len(scores) > 1 else 0.0
        features["score_spread"] = max(scores) - min(scores)
    else:
        features["mean_vector_score"] = None
        features["max_vector_score"] = None
        features["min_vector_score"] = None
        features["std_vector_score"] = None
        features["score_spread"] = None

    # Vocabulary overlap between query and chunks
    query_words = {w.strip("?.,!;:()[]\"'") for w in query.lower().split()}
    query_content_words = query_words - _STOP_WORDS - {""}
    if query_content_words:
        chunk_text = " ".join(c.content.lower() for c in chunks)
        chunk_words = {w.strip("?.,!;:()[]\"'") for w in chunk_text.split()}
        overlap = query_content_words & chunk_words
        features["vocab_overlap_ratio"] = len(overlap) / len(query_content_words)
    else:
        features["vocab_overlap_ratio"] = 0.0

    # Inter-chunk text features (Tier 2b — deterministic, no LLM)
    _extract_interchunk_features(features, chunks)

    # Text-level answer features (Tier 2c — query+context text analysis)
    _extract_text_answer_features(features, query, chunks)


_STOP_WORDS = {
    "the",
    "a",
    "an",
    "is",
    "are",
    "was",
    "were",
    "be",
    "been",
    "being",
    "have",
    "has",
    "had",
    "do",
    "does",
    "did",
    "will",
    "would",
    "could",
    "should",
    "may",
    "might",
    "can",
    "shall",
    "in",
    "on",
    "at",
    "to",
    "for",
    "of",
    "with",
    "by",
    "from",
    "and",
    "or",
    "but",
    "not",
    "no",
    "if",
    "then",
    "than",
    "that",
    "this",
    "these",
    "those",
    "what",
    "which",
    "who",
    "how",
    "when",
    "where",
    "why",
    "it",
    "its",
}

_HEDGE_WORDS = {
    "may",
    "might",
    "could",
    "possibly",
    "perhaps",
    "likely",
    "unlikely",
    "sometimes",
    "often",
    "typically",
    "generally",
    "usually",
    "probably",
    "approximately",
    "roughly",
    "about",
    "around",
    "estimated",
    "suggests",
    "appears",
    "seems",
    "potentially",
    "tends",
}

_ASSERTION_WORDS = {
    "always",
    "never",
    "must",
    "certainly",
    "definitely",
    "clearly",
    "obviously",
    "undoubtedly",
    "absolutely",
    "exactly",
    "precisely",
    "proven",
    "confirmed",
    "established",
    "demonstrates",
    "proves",
    "invariably",
    "unquestionably",
}

_NUMBER_RE = re.compile(r"\b\d+(?:\.\d+)?(?:%|st|nd|rd|th)?\b")
_PLAIN_NUMBER_RE = re.compile(r"\b\d+\.?\d*\b")
_YEAR_RE = re.compile(r"\b(19\d{2}|20\d{2})\b")
# Matches dollar amounts, percentages, and plain numbers (for cross-chunk divergence)
_SIGNIFICANT_NUMBER_RE = re.compile(r"[\$]?\d[\d,]*\.?\d*[%]?")

_CONTRADICTION_MARKERS = [
    "however",
    "but",
    "although",
    "contrary",
    "disagree",
    "whereas",
    "nevertheless",
    "conversely",
    "despite",
    "in contrast",
    "on the other hand",
    "contradicts",
    "inconsistent",
    "conflicts with",
    "differs from",
]
# Stronger markers that signal opposing conclusions (not just hedging)
_OPPOSING_MARKERS = [
    "contradicts",
    "inconsistent",
    "conflicts with",
    "disagree",
    "disputed",
    "refuted",
    "disproven",
    "challenged by",
    "rejected by",
    "opposes",
    "incompatible",
    "at odds with",
    "mutually exclusive",
    "cannot both be true",
    "conflicting evidence",
    "conflicting data",
    "conflicting reports",
]
_NEGATION_WORDS = {
    "not",
    "no",
    "never",
    "neither",
    "nor",
    "none",
    "nothing",
    "hardly",
    "barely",
    "scarcely",
    "doesn't",
    "don't",
    "isn't",
    "wasn't",
    "weren't",
    "won't",
    "can't",
    "couldn't",
    "shouldn't",
}


def _extract_interchunk_features(features: dict[str, Any], chunks: Sequence[EvidenceItem]) -> None:
    """Extract inter-chunk text relationship features (deterministic, no LLM).

    Includes features previously only available at training time via
    train_classifier.py compute_context_features(). All features here are
    available at both training and inference time.
    """
    _defaults = {
        "max_pairwise_overlap": 0.0,
        "min_pairwise_overlap": 0.0,
        "chunk_length_cv": 0.0,
        "assertion_density": 0.0,
        "hedge_density": 0.0,
        "number_density": 0.0,
        "ctx_length_mean": 0.0,
        "ctx_length_std": 0.0,
        "ctx_total_chars": 0.0,
        "ctx_contradiction_count": 0.0,
        "ctx_negation_count": 0.0,
        "ctx_number_count": 0.0,
        "ctx_number_variance": 0.0,
        "ctx_max_pairwise_sim": 0.0,
        "ctx_mean_pairwise_sim": 0.0,
        "ctx_min_pairwise_sim": 0.0,
        "cross_chunk_num_conflicts": 0.0,
        "cross_chunk_max_divergence": 0.0,
        "within_chunk_num_conflicts": 0.0,
        "within_chunk_max_divergence": 0.0,
        "year_count": 0,
        "has_distinct_years": False,
        "has_cross_chunk_divergence": False,
        "has_within_chunk_divergence": False,
        "conflict_to_number_ratio": 0.0,
        "opposing_conclusion_count": 0.0,
        "negation_per_char": 0.0,
        "short_ctx_with_overlap": 0.0,
    }
    if len(chunks) < 2:
        features.update(_defaults)
        # Still compute single-chunk stats (ctx_length_mean, ctx_total_chars, etc.)
        if len(chunks) == 1:
            text = chunks[0].content
            features["ctx_length_mean"] = float(len(text))
            features["ctx_total_chars"] = float(len(text))
            words = text.lower().split()
            features["number_density"] = float(len(_NUMBER_RE.findall(text)))
            word_set = set(words)
            total_hedge = len(word_set & _HEDGE_WORDS)
            total_assert = len(word_set & _ASSERTION_WORDS)
            epistemic_total = total_assert + total_hedge
            if epistemic_total > 0:
                features["assertion_density"] = (total_assert - total_hedge) / epistemic_total
            years = set(_YEAR_RE.findall(text))
            features["year_count"] = len(years)
            features["has_distinct_years"] = len(years) > 1
            # Within-chunk divergence works on single chunks too
            _compute_within_chunk_numerical_divergence(features, [text])
            # Single-chunk versions of new features
            text_lower = text.lower()
            features["ctx_negation_count"] = float(sum(1 for w in words if w in _NEGATION_WORDS))
            features["ctx_contradiction_count"] = float(
                sum(1 for m in _CONTRADICTION_MARKERS if m in text_lower)
            )
            features["ctx_number_count"] = float(len(_PLAIN_NUMBER_RE.findall(text)))
            features["opposing_conclusion_count"] = float(
                sum(1 for m in _OPPOSING_MARKERS if m in text_lower)
            )
            features["negation_per_char"] = features["ctx_negation_count"] / max(len(text), 1)
            features["conflict_to_number_ratio"] = 0.0  # no cross-chunk conflicts with 1 chunk
            features["short_ctx_with_overlap"] = float(
                len(text) < 500 and features.get("vocab_overlap_ratio", 0) > 0.3
            )
        return

    # Precompute per-chunk data
    texts = [c.content for c in chunks]
    chunk_word_sets = []
    chunk_lengths = []
    char_lengths = []
    total_hedge = 0
    total_assert = 0
    total_numbers = 0
    all_text_lower = ""
    all_numbers: list[float] = []
    all_years: set[str] = set()

    for c in chunks:
        text = c.content
        text_lower = text.lower()
        all_text_lower += " " + text_lower

        words = text_lower.split()
        content_words = set(words) - _STOP_WORDS
        chunk_word_sets.append(content_words)
        chunk_lengths.append(len(words))
        char_lengths.append(len(text))

        word_set = set(words)
        total_hedge += len(word_set & _HEDGE_WORDS)
        total_assert += len(word_set & _ASSERTION_WORDS)
        total_numbers += len(_NUMBER_RE.findall(text))

        for m in _PLAIN_NUMBER_RE.finditer(text):
            try:
                all_numbers.append(float(m.group(0)))
            except ValueError:
                pass
        for ym in _YEAR_RE.finditer(text):
            all_years.add(ym.group(0))

    # --- Pairwise Jaccard overlap ---
    overlaps = []
    for i in range(len(chunk_word_sets)):
        for j in range(i + 1, len(chunk_word_sets)):
            a, b = chunk_word_sets[i], chunk_word_sets[j]
            union = a | b
            overlaps.append(len(a & b) / len(union) if union else 0.0)

    features["max_pairwise_overlap"] = max(overlaps) if overlaps else 0.0
    features["min_pairwise_overlap"] = min(overlaps) if overlaps else 0.0

    # --- Chunk length coefficient of variation ---
    mean_len = statistics.mean(chunk_lengths)
    if mean_len > 0 and len(chunk_lengths) > 1:
        features["chunk_length_cv"] = statistics.stdev(chunk_lengths) / mean_len
    else:
        features["chunk_length_cv"] = 0.0

    # --- Assertion vs hedge density ---
    epistemic_total = total_assert + total_hedge
    if epistemic_total > 0:
        features["assertion_density"] = (total_assert - total_hedge) / epistemic_total
    else:
        features["assertion_density"] = 0.0

    # --- Hedge density (raw hedge count / total words) ---
    total_words = sum(chunk_lengths)
    features["hedge_density"] = total_hedge / total_words if total_words > 0 else 0.0

    # --- Number density (numbers per chunk) ---
    features["number_density"] = total_numbers / len(chunks)

    # --- Context length stats (ported from train_classifier.py) ---
    features["ctx_length_mean"] = statistics.mean(char_lengths)
    features["ctx_length_std"] = statistics.pstdev(char_lengths)
    features["ctx_total_chars"] = sum(char_lengths)

    # --- Contradiction and negation markers ---
    features["ctx_contradiction_count"] = float(
        sum(1 for m in _CONTRADICTION_MARKERS if m in all_text_lower)
    )
    all_words = all_text_lower.split()
    features["ctx_negation_count"] = float(sum(1 for w in all_words if w in _NEGATION_WORDS))

    # --- Numerical content ---
    features["ctx_number_count"] = float(len(all_numbers))
    if len(all_numbers) > 1:
        mean_n = sum(all_numbers) / len(all_numbers)
        features["ctx_number_variance"] = float(
            sum((x - mean_n) ** 2 for x in all_numbers) / len(all_numbers)
        )
    else:
        features["ctx_number_variance"] = 0.0

    # --- TF-IDF pairwise similarity ---
    try:
        from sklearn.feature_extraction.text import TfidfVectorizer
        from sklearn.metrics.pairwise import cosine_similarity as _cos_sim

        tfidf = TfidfVectorizer(max_features=500, stop_words="english")
        matrix = tfidf.fit_transform(texts)
        sim_matrix = _cos_sim(matrix)
        n = sim_matrix.shape[0]
        pairwise_sims = [float(sim_matrix[i, j]) for i in range(n) for j in range(i + 1, n)]
        features["ctx_max_pairwise_sim"] = max(pairwise_sims)
        features["ctx_mean_pairwise_sim"] = sum(pairwise_sims) / len(pairwise_sims)
        features["ctx_min_pairwise_sim"] = min(pairwise_sims)
    except Exception:
        features["ctx_max_pairwise_sim"] = 0.0
        features["ctx_mean_pairwise_sim"] = 0.0
        features["ctx_min_pairwise_sim"] = 0.0

    # --- Temporal features ---
    features["year_count"] = len(all_years)
    features["has_distinct_years"] = len(all_years) > 1

    # --- Cross-chunk numerical divergence ---
    # Numerical divergence detection (both across and within chunks)
    _compute_cross_chunk_numerical_divergence(features, texts)
    _compute_within_chunk_numerical_divergence(features, texts)

    # --- Q3-specific features: distinguish data-rich documents from real conflicts ---
    features["numerical_richness_per_chunk"] = features["ctx_number_count"] / max(len(chunks), 1)
    features["conflict_internality_ratio"] = features["within_chunk_num_conflicts"] / max(
        features["cross_chunk_num_conflicts"] + features["within_chunk_num_conflicts"],
        1,
    )
    features["chars_per_chunk"] = features["ctx_total_chars"] / max(len(chunks), 1)
    features["contradiction_per_char"] = features["ctx_contradiction_count"] / max(
        features["ctx_total_chars"], 1
    )

    # --- Conflict quality features (Q3/FT-disputed) ---
    # Real disputes have conflicts as large fraction of total numbers;
    # data-rich docs have many numbers with few conflicts proportionally
    features["conflict_to_number_ratio"] = features["cross_chunk_num_conflicts"] / max(
        features["ctx_number_count"], 1
    )
    # Opposing language: stronger contradiction markers beyond simple "however/but"
    features["opposing_conclusion_count"] = float(
        sum(1 for m in _OPPOSING_MARKERS if m in all_text_lower)
    )
    # Negation density: disputes tend to have more negations per unit of text
    features["negation_per_char"] = features["ctx_negation_count"] / max(
        features["ctx_total_chars"], 1
    )

    # --- Q1 recovery features (partial-answer trustworthy) ---
    # Short context but some content overlap = partial answer, not absence
    features["short_ctx_with_overlap"] = float(
        features["ctx_total_chars"] < 500 and features.get("vocab_overlap_ratio", 0) > 0.3
    )


def _extract_numbers_from_text(text: str) -> list[float]:
    """Extract significant numbers from text, excluding years and trivial values."""
    nums = []
    for m in _SIGNIFICANT_NUMBER_RE.finditer(text):
        clean = m.group(0).replace(",", "").replace("$", "").replace("%", "")
        try:
            n = float(clean)
            if n < 2 or (1900 <= n <= 2099):
                continue
            nums.append(n)
        except ValueError:
            pass
    return nums


def _count_divergent_pairs(nums_a: list[float], nums_b: list[float]) -> tuple[int, float]:
    """Count conflicting number pairs and max divergence ratio between two sets."""
    conflict_count = 0
    max_ratio = 0.0
    for ni in nums_a:
        for nj in nums_b:
            if min(ni, nj) <= 0:
                continue
            ratio = max(ni, nj) / min(ni, nj)
            if 1.02 < ratio < 2.0:
                conflict_count += 1
                if ratio > max_ratio:
                    max_ratio = ratio
    return conflict_count, max_ratio


def _compute_cross_chunk_numerical_divergence(features: dict[str, Any], texts: list[str]) -> None:
    """Detect numerical disagreement across chunks.

    Extracts significant numbers from each chunk, then finds cross-chunk pairs
    where two numbers are close in magnitude (ratio < 2.0) but not equal
    (ratio > 1.02). This catches conflicts like $485k vs $520k, 27.78 vs 27.44 mph.
    Years (1900-2099) are excluded.
    """
    per_chunk_nums = [_extract_numbers_from_text(text) for text in texts]

    conflict_count = 0
    max_divergence_ratio = 0.0

    for i in range(len(per_chunk_nums)):
        for j in range(i + 1, len(per_chunk_nums)):
            c, r = _count_divergent_pairs(per_chunk_nums[i], per_chunk_nums[j])
            conflict_count += c
            max_divergence_ratio = max(max_divergence_ratio, r)

    features["cross_chunk_num_conflicts"] = conflict_count
    features["cross_chunk_max_divergence"] = max_divergence_ratio
    features["has_cross_chunk_divergence"] = conflict_count > 0


def _compute_within_chunk_numerical_divergence(features: dict[str, Any], texts: list[str]) -> None:
    """Detect numerical disagreement within individual chunks.

    A single chunk can contain contradictory numbers for the same quantity,
    e.g. "speed limit reduced from 55 to 45 mph" or "revenue was $5M...
    adjusted to $3.2M". This is common in production when chunks contain
    updates, revisions, or multi-source summaries.
    """
    total_conflicts = 0
    max_ratio = 0.0

    for text in texts:
        nums = _extract_numbers_from_text(text)
        if len(nums) < 2:
            continue
        # Compare all pairs within this single chunk
        for i in range(len(nums)):
            for j in range(i + 1, len(nums)):
                ni, nj = nums[i], nums[j]
                if min(ni, nj) <= 0:
                    continue
                ratio = max(ni, nj) / min(ni, nj)
                if 1.02 < ratio < 2.0:
                    total_conflicts += 1
                    max_ratio = max(max_ratio, ratio)

    features["within_chunk_num_conflicts"] = total_conflicts
    features["within_chunk_max_divergence"] = max_ratio
    features["has_within_chunk_divergence"] = total_conflicts > 0


_SENTENCE_SPLIT_RE = re.compile(r"[.!?]+")


def _extract_text_answer_features(
    features: dict[str, Any], query: str, chunks: Sequence[EvidenceItem]
) -> None:
    """Extract text-level features measuring how well context answers the query.

    These features capture whether the retrieved context actually contains an answer
    to the question, beyond simple word overlap. They are cheap (regex/set ops, no LLM).
    """
    query_words = {w.strip("?.,!;:()[]\"'") for w in query.lower().split()}
    q_content = query_words - _STOP_WORDS - {""}
    # Strip question words from subject extraction
    _question_words = {
        "what",
        "how",
        "why",
        "when",
        "where",
        "who",
        "which",
        "is",
        "are",
        "was",
        "were",
        "does",
        "do",
        "did",
        "can",
        "could",
        "should",
        "will",
        "would",
    }
    q_subject_words = q_content - _question_words

    if not q_subject_words or not chunks:
        features["query_subject_partial"] = 0.0
        features["best_span_length"] = 0.0
        features["best_sentence_coverage"] = 0.0
        features["entity_substantive_score"] = 0.0
        features["answer_span_coverage"] = 0.0
        return

    ctx_texts = [c.content for c in chunks]
    ctx_joined_lower = " ".join(t.lower() for t in ctx_texts)

    # 1. query_subject_partial: fraction of query subject words found in context
    found = sum(1 for w in q_subject_words if w in ctx_joined_lower)
    features["query_subject_partial"] = found / len(q_subject_words)

    # 2. entity_substantive_score: query words mentioned 2+ times (discussed, not just mentioned)
    substantive = sum(1 for w in q_subject_words if ctx_joined_lower.count(w) >= 2)
    features["entity_substantive_score"] = substantive / len(q_subject_words)

    # 3-5. Sentence-level features
    best_coverage = 0.0
    best_span_len = 0
    best_span_coverage = 0.0

    for text in ctx_texts:
        for sent in _SENTENCE_SPLIT_RE.split(text):
            sent = sent.strip()
            if not sent:
                continue
            sent_words = [w.strip("?.,!;:()[]\"'-").lower() for w in sent.split()]
            sent_word_set = set(sent_words) - _STOP_WORDS - {""}

            # best_sentence_coverage
            overlap = q_content & sent_word_set
            coverage = len(overlap) / len(q_content)
            if coverage > best_coverage:
                best_coverage = coverage

            # best_span_length / answer_span_coverage: longest contiguous span with 2+ query words
            for start in range(len(sent_words)):
                span_covered: set[str] = set()
                for end in range(start, min(start + 30, len(sent_words))):
                    w = sent_words[end].strip("?.,!;:()[]\"'-")
                    if w in q_content:
                        span_covered.add(w)
                    if len(span_covered) >= 2:
                        span_cov = len(span_covered) / len(q_content)
                        if span_cov > best_span_coverage:
                            best_span_coverage = span_cov
                            best_span_len = end - start + 1

    features["best_sentence_coverage"] = best_coverage
    features["best_span_length"] = float(best_span_len)
    features["answer_span_coverage"] = best_span_coverage


def _extract_detection_features(features: dict[str, Any], detection_summary: Any | None) -> None:
    """Extract features from DetectionSummary (Tier 3)."""
    if detection_summary is None:
        features["detection_temporal"] = None
        features["detection_comparison"] = None
        features["detection_aggregation"] = None
        features["detection_boost_recency"] = None
        features["detection_boost_authority"] = None
        features["detection_needs_rewriting"] = None
        return

    features["detection_temporal"] = getattr(detection_summary, "has_temporal_intent", None)
    features["detection_comparison"] = getattr(detection_summary, "has_comparison_intent", None)
    features["detection_aggregation"] = getattr(detection_summary, "has_aggregation_intent", None)
    features["detection_boost_recency"] = getattr(detection_summary, "boost_recency", None)
    features["detection_boost_authority"] = getattr(detection_summary, "boost_authority", None)
    features["detection_needs_rewriting"] = getattr(detection_summary, "needs_rewriting", None)


# ── Feature type sets (derived from constraint schemas + static Tier 2/3) ──

# Tier 2/3 bool features not owned by constraints
_TIER2_BOOL_FEATURES = {
    "query_has_comparison_words",
    "has_distinct_years",
    "has_cross_chunk_divergence",
    "has_within_chunk_divergence",
}
_TIER3_BOOL_FEATURES = {
    "detection_temporal",
    "detection_comparison",
    "detection_aggregation",
    "detection_boost_recency",
    "detection_boost_authority",
    "detection_needs_rewriting",
}
# Aggregate cross-constraint bool features
_AGGREGATE_BOOL_FEATURES = {
    "has_abstain_signal",
    "has_disputed_signal",
    "has_qualified_signal",
    "has_any_denial",
    "av_strong_denial",
}
_TIER2_CATEGORICAL_FEATURES = {"query_question_type"}


def get_feature_type_sets() -> tuple[set[str], set[str]]:
    """Return (categorical_features, bool_features) derived from constraint schemas + static Tier 2/3."""
    categorical: set[str] = set(_TIER2_CATEGORICAL_FEATURES)
    bools: set[str] = set(_TIER2_BOOL_FEATURES | _TIER3_BOOL_FEATURES | _AGGREGATE_BOOL_FEATURES)
    for cls in _CONSTRAINT_CLASSES:
        for spec in cls.feature_schema():
            if spec.type == "categorical":
                categorical.add(spec.name)
            elif spec.type == "bool":
                bools.add(spec.name)
    return categorical, bools


__all__ = ["extract_features", "get_feature_type_sets"]
