"""PEF updater — apply extraction results to PEF state.

Backend-agnostic: any ExtractionResult (from spaCy, LLM, or anything else)
feeds through the same updater. Zero spaCy dependency.

Extraction proposes deltas; this module applies them to the existing world.
"""

from __future__ import annotations

from aurora_lens.pef.span import Span
from aurora_lens.pef.state import PEFState, Relationship
from aurora_lens.interpret.schema import ExtractedClaim, ExtractionResult


_PRONOUNS = {"he", "she", "it", "they", "him", "her", "them", "his", "its", "their"}

# Generic concept labels that must not become PEF entity placeholders.
# These are category-level abbreviations (e.g. in "AI governance", "LLM safety").
# Registering them as unresolved placeholders causes false UNSUPPORTED_ATTRIBUTE
# flags when the LLM makes ordinary claims about the concept.
_GENERIC_CONCEPT_LABELS: frozenset[str] = frozenset({
    "ai", "ml", "llm", "nlp", "api", "os", "iot", "ar", "vr",
    "ui", "ux", "ci", "cd", "qa", "hr", "pr", "llms",
})

# Definite description prefixes — subjects starting with these require a prior
# antecedent; they must NOT be auto-resolved as new concept literals.
_DEFINITE_PREFIXES: tuple[str, ...] = ("the ", "that ", "this ", "those ", "these ")


def _is_definite_description(s: str) -> bool:
    """Return True if s is a multi-word definite description (starts with 'the/that/this/...')."""
    lower = s.strip().lower()
    return any(lower.startswith(p) for p in _DEFINITE_PREFIXES)


def update_pef(result: ExtractionResult, pef: PEFState) -> None:
    """Apply extraction results to PEF state.

    Creates/updates entities, adds relationships, records bindings.
    """
    # Apply pronoun bindings
    for pronoun_key, entity_name in result.pronoun_candidates.items():
        entity = pef.find_entity_by_name(entity_name)
        if entity:
            pef.resolve_pronoun(pronoun_key, entity.id)

    # Process claims
    for claim in result.claims:
        subject_name = claim.subject

        # Resolve pronoun subjects
        if subject_name.lower() in _PRONOUNS:
            resolved_name = _resolve_single_pronoun(subject_name, pef)
            if resolved_name:
                subject_name = resolved_name

        # Get or create subject entity
        subj_entity, _ = pef.get_or_create_entity(subject_name)

        # Admit bare nouns and proper nouns as resolved concept literals.
        # Entities previously created as unresolved placeholders (from
        # entity_mentions) are promoted here when a claim is asserted about them.
        # Pronouns and definite descriptions are excluded — they require explicit
        # PEF grounding and must not be silently resolved.
        if subject_name.lower() not in _PRONOUNS and not _is_definite_description(subject_name):
            subj_entity.resolved = True

        # Determine if object is an entity reference or literal
        obj_entity = pef.find_entity_by_name(claim.obj)

        if obj_entity:
            rel = Relationship(
                subject_id=subj_entity.id,
                relation=claim.relation,
                object_entity_id=obj_entity.id,
                object_literal=None,
                span=claim.span,
                source_turn=pef.current_turn,
                evidence=claim.evidence,
                negated=claim.negated,
                provenance=claim.provenance,
                extractor_backend=claim.extractor_backend,
                document_id=claim.document_id,
                document_locator=claim.document_locator,
            )
        else:
            rel = Relationship(
                subject_id=subj_entity.id,
                relation=claim.relation,
                object_entity_id=None,
                object_literal=claim.obj,
                span=claim.span,
                source_turn=pef.current_turn,
                evidence=claim.evidence,
                negated=claim.negated,
                provenance=claim.provenance,
                extractor_backend=claim.extractor_backend,
                document_id=claim.document_id,
                document_locator=claim.document_locator,
            )

        pef.add_relationship(rel)

        # Record echo traces
        pef.record_echo_trace(subj_entity.id, claim.evidence)
        if subj_entity.resolved is False and claim.relation != "IS":
            # Asserting a fact about the entity resolves it
            subj_entity.resolved = True

    # Create unresolved placeholders for mentioned-but-not-asserted entities.
    # Skip generic concept labels (AI, ML, LLM, API, etc.) - they are category-level
    # abbreviations, not specific tracked world entities.
    for name in result.entity_mentions:
        if name.lower() not in _PRONOUNS and name.lower() not in _GENERIC_CONCEPT_LABELS:
            pef.get_or_create_entity(name, resolved=False)


def apply_prepopulated_claims(
    claims: list[ExtractedClaim],
    pef: PEFState,
    document_id: str,
    document_locator: str | None = None,
) -> None:
    """Apply claims from a pre-populated corpus/document to PEF.

    Forces provenance="pre_populated", keeps extractor_backend from each claim,
    sets document_id and document_locator. Uses turn 0 for source_turn.
    """
    # Override provenance for all claims
    prepop_claims = [
        ExtractedClaim(
            subject=c.subject,
            relation=c.relation,
            obj=c.obj,
            span=c.span,
            negated=c.negated,
            evidence=c.evidence,
            provenance="pre_populated",
            extractor_backend=c.extractor_backend,
            document_id=document_id,
            document_locator=document_locator,
        )
        for c in claims
    ]
    # Use turn 0 for pre-populated facts (before conversation)
    saved_turn = pef.current_turn
    pef.current_turn = 0
    try:
        result = ExtractionResult(
            claims=prepop_claims,
            entity_mentions=[],
            span=prepop_claims[0].span if prepop_claims else Span.PRESENT,
        )
        update_pef(result, pef)
    finally:
        pef.current_turn = saved_turn


def _resolve_single_pronoun(pronoun: str, pef: PEFState) -> str | None:
    """Resolve a pronoun to the most recently active matching entity.

    Uses recency heuristic. Returns None if no resolved entities exist.
    """
    resolved_entities = [
        e for e in pef.entities.values() if e.resolved
    ]
    if not resolved_entities:
        return None

    by_recency = sorted(
        resolved_entities,
        key=lambda e: e.turn_last_active,
        reverse=True,
    )

    if by_recency:
        return by_recency[0].name

    return None
