"""
cortex/extraction/models.py — Data models for the extraction pipeline.

Pure dataclasses with no pipeline logic. These are data contracts consumed
by 15+ files across cortex/ and tests/.

Classes:
    EntityMention: An entity reference extracted from conversation text.
    ClaimExtraction: A single extracted memory claim with to_candidate() bridge.
    ExtractionResult: Full result from the extraction pipeline.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

_logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Valid-value sets for ClaimExtraction enum-constrained string fields.
# Used by _validate_field() and __post_init__; exported so tests can import.
# ---------------------------------------------------------------------------
VALID_ACTIONS: frozenset = frozenset({"ADD", "UPDATE", "DELETE", "NOOP"})
VALID_TEMPORAL: frozenset = frozenset({"current", "past", "future", "permanent"})
VALID_MEMORY_CLASS: frozenset = frozenset({"semantic", "episodic", "procedural", "open_loop", "transient"})
VALID_SALIENCE: frozenset = frozenset({"high", "medium", "low"})
VALID_STABILITY: frozenset = frozenset({"stable", "mutable", "volatile"})
VALID_SENSITIVITY: frozenset = frozenset({"normal", "personal", "sensitive"})
from cortex.types.constants import VALID_CATEGORIES as VALID_CATEGORY  # single source of truth
VALID_STATUS: frozenset = frozenset({"active", "completed", "superseded", "blocked"})
VALID_INTRINSIC_TYPE: frozenset = frozenset({"fact", "preference", "goal"})
VALID_EXPLICITNESS: frozenset = frozenset({"explicit", "implied", "inferred"})


def _validate_field(value: str, valid_set: frozenset, field_name: str, default: str) -> str:
    """Return *value* unchanged if valid; otherwise clamp to *default* with a warning.

    Never raises — extraction must remain robust even when the LLM returns
    unexpected values.  Handles non-string types (e.g. dict/list from
    malformed LLM JSON) and performs case-insensitive matching.
    """
    # Type guard: LLM may return dict/list instead of str — unhashable for `in`
    if not isinstance(value, str):
        _logger.warning(
            "ClaimExtraction: non-string %s=%r (type %s), clamping to %r",
            field_name,
            value,
            type(value).__name__,
            default,
        )
        return default

    # Case-insensitive lookup: find canonical form from valid_set
    lower = value.lower()
    for canonical in valid_set:
        if canonical.lower() == lower:
            return canonical

    _logger.warning(
        "ClaimExtraction: invalid %s=%r (expected one of %s), clamping to %r",
        field_name,
        value,
        sorted(valid_set),
        default,
    )
    return default


@dataclass
class EntityMention:
    """An entity reference extracted from conversation text.

    These are passed to EntityResolver (M3) for canonicalization — the
    same real-world entity might be mentioned as "Andrew", "Andy", "@andrew".

    Attributes:
        name: Primary name as it appeared in the conversation.
        entity_type: Classification — 'person', 'technology', 'project', etc.
        aliases: Other names/spellings observed in the same context.
    """
    name: str
    entity_type: str
    aliases: List[str] = field(default_factory=list)


@dataclass
class ClaimExtraction:
    """A single extracted memory claim from the conversation.

    This is the intermediate format between LLM extraction and storage.
    The ReconciliationEngine (M4) decides whether each claim should be
    added, merged with an existing claim, or discarded.

    Attributes:
        action: LLM-suggested action — 'ADD' (new fact), 'UPDATE' (revises
                existing), 'DELETE' (negates prior claim), 'NOOP' (no info).
        subject: The entity or concept this claim is about.
        predicate: The relationship or property (verb/adjective).
        object_value: The value or target of the relationship.
        confidence: LLM-assigned confidence (0.0-1.0). Lower for inferred items.
        temporal: Temporal scope — 'current', 'past', 'future', 'permanent'.
        temporal_marker: Explicit time reference extracted verbatim from the
                         source text, or None if no specific time phrase was
                         mentioned.  Examples: "yesterday", "last week",
                         "March 2026", "two years ago", "recently".
        entities_mentioned: All entities referenced in this claim (for M3).
        sensitivity: Privacy classification — 'normal', 'personal', 'sensitive'.

    Note:
        Output feeds into ReconciliationEngine.reconcile() as candidates.
    """
    action: str                           # ADD | UPDATE | DELETE | NOOP
    subject: str
    predicate: str
    object_value: str
    confidence: float
    temporal: str = "current"             # current | past | future | permanent
    temporal_marker: Optional[str] = None # e.g. "yesterday", "March 2026", None
    entities_mentioned: List[EntityMention] = field(default_factory=list)
    sensitivity: str = "normal"           # normal | personal | sensitive
    category: str = "misc"               # identity | personal | preferences | decisions | goals | behavioral | relational | episodic | misc
    intrinsic_type: str = "fact"         # fact | preference | goal
    human_family: Optional[str] = None    # family | friend | colleague | partner | mentor | mentee
    explicitness: str = "inferred"       # explicit | implied | inferred
    stability: str = "volatile"          # stable | mutable | volatile
    # --- LoCoMo-aligned fields (v1.3) ---
    memory_class: str = "semantic"        # semantic|episodic|procedural|open_loop|transient
    status: str = "active"               # active|completed|superseded|blocked
    event_date: Optional[str] = None     # ISO 8601 YYYY-MM-DD when parseable
    sequence_order: Optional[int] = None  # within-conversation ordering (1,2,3...)
    salience: str = "medium"             # high|medium|low
    # Causal linking (#518)
    cause: Optional[str] = None           # reason/trigger for this fact
    effect: Optional[str] = None          # outcome/result of the cause
    # State transitions (#519)
    state_before: Optional[str] = None    # previous state being replaced
    state_after: Optional[str] = None     # new state after the change
    # Source role (#321 — split extraction mode)
    source_role: Optional[str] = None     # 'user' | 'agent' | 'system' — who said the source message

    def __post_init__(self) -> None:
        """Validate and clamp all enum-constrained string fields coming from LLM output.

        Invalid values are replaced with a safe default and a warning is logged.
        This never raises so that extraction remains robust under unexpected LLM output.
        """
        self.action = _validate_field(self.action, VALID_ACTIONS, "action", "NOOP")
        self.temporal = _validate_field(self.temporal, VALID_TEMPORAL, "temporal", "current")
        self.sensitivity = _validate_field(self.sensitivity, VALID_SENSITIVITY, "sensitivity", "normal")
        self.category = _validate_field(self.category, VALID_CATEGORY, "category", "misc")
        self.intrinsic_type = _validate_field(self.intrinsic_type, VALID_INTRINSIC_TYPE, "intrinsic_type", "fact")
        self.explicitness = _validate_field(self.explicitness, VALID_EXPLICITNESS, "explicitness", "inferred")
        self.stability = _validate_field(self.stability, VALID_STABILITY, "stability", "volatile")
        self.memory_class = _validate_field(self.memory_class, VALID_MEMORY_CLASS, "memory_class", "semantic")
        self.status = _validate_field(self.status, VALID_STATUS, "status", "active")
        self.salience = _validate_field(self.salience, VALID_SALIENCE, "salience", "medium")

    def to_candidate(self) -> "CandidateClaim":
        """Convert this extraction result to the reconciliation contract type.

        ``ClaimExtraction`` is the ExtractionPipeline (M2) output format.
        ``CandidateClaim`` is the ReconciliationEngine (M4) input contract.
        Calling this method is the *only* supported way to bridge between them;
        manual field-by-field construction should not be used.

        Mapping decisions
        -----------------
        * Shared fields (subject, predicate, object_value, confidence,
          sensitivity, intrinsic_type, human_family, explicitness, stability)
          are copied directly.
        * ``action`` is mapped to ``is_user_explicit``:
          - "ADD" + explicitness=="explicit"  -> True (direct user statement)
          - Everything else                  -> False
          The reconciliation engine W4 recency-precedence logic uses this flag
          to let explicit current statements supersede older inferences.
        * ``temporal_marker`` is preserved on CandidateClaim (field added in
          types.py) because reconciliation already reads it via getattr().
          No information is lost.
        * ``entity_type`` is derived from the *first* entity mention's type,
          falling back to None when no entities were mentioned.  The
          reconciliation engine falls back to "unknown" when this is None.
        * ``source_text`` is not available at extraction output - set to None.
        * ``claim_type`` defaults to "semantic" (the only type currently used).
        * ``metadata`` is set to None; callers may populate it after conversion.

        Returns:
            A ``CandidateClaim`` ready for consumption by ReconciliationEngine.
        """
        from cortex.types import CandidateClaim  # local import avoids any circular risk

        # Derive entity_type from the first mentioned entity (if any).
        entity_type: Optional[str] = None
        if self.entities_mentioned:
            raw_et = getattr(self.entities_mentioned[0], "entity_type", None)
            if raw_et:
                entity_type = str(raw_et).strip() or None

        # is_user_explicit: True only when the LLM judged the claim as an
        # explicit (high-confidence) direct statement ("ADD" is the normal
        # action for new explicit information).
        is_user_explicit = (
            self.action == "ADD" and self.explicitness == "explicit"
        )

        return CandidateClaim(
            subject=self.subject,
            predicate=self.predicate,
            object_value=self.object_value,
            confidence=self.confidence,
            source_text=None,
            entity_type=entity_type,
            claim_type="semantic",
            metadata=None,
            intrinsic_type=self.intrinsic_type,
            human_family=self.human_family,
            explicitness=self.explicitness,
            stability=self.stability,
            sensitivity=self.sensitivity,
            category=self.category,
            is_user_explicit=is_user_explicit,
            temporal=self.temporal,
            temporal_marker=self.temporal_marker,
            memory_class=self.memory_class,
            status=self.status,
            event_date=self.event_date,
            sequence_order=self.sequence_order,
            salience=self.salience,
            cause=self.cause,
            effect=self.effect,
            state_before=self.state_before,
            state_after=self.state_after,
            action=self.action,
            source_role=self.source_role,
        )


@dataclass
class ExtractionResult:
    """Full result from the extraction pipeline for one conversation batch.

    Attributes:
        claims: List of extracted ClaimExtraction objects (NOOP claims already filtered).
        noise_filtered: How many input messages were dropped by CodingNoiseFilter.
        raw_message_count: Total input messages before any filtering.
        batches_processed: Number of LLM batch calls made during extraction.
                           A value > 1 means the input was split into multiple
                           batches of batch_size messages each.

    Note:
        Passed to CortexPlugin.remember() which forwards claims to
        ReconciliationEngine.reconcile() (M4).
    """
    claims: List[ClaimExtraction]
    noise_filtered: int      # number of messages dropped by noise filter
    raw_message_count: int
    batches_processed: int = 1  # number of LLM calls made (≥1)
    batches_failed: int = 0     # batches lost to LLM errors (partial result)
