"""
cortex/extraction/extractor.py — ExtractionPipeline (Module 2 / M2).

The first stage of the remember() pipeline. Converts raw conversation turns
into structured CandidateClaim objects (subject-predicate-object triples)
that downstream modules can reconcile and persist.

Pipeline stages:
    1. CodingNoiseFilter (regex pre-filter — zero LLM cost)
    2. Batch splitting — divide filtered messages into chunks of batch_size
    3. For each batch: format_conversation() + LLM completion (structured JSON)
    4. JSON parse + validation → list of ClaimExtraction dataclasses (per batch)
    5. Aggregate all batch results into ExtractionResult

Does NOT write to storage — output feeds into ReconciliationEngine.reconcile() (M4).

Dependencies:
    - cortex.extraction.coding_noise.CodingNoiseFilter (noise pre-filter)
    - cortex.extraction.prompts (system prompt + conversation formatter)
    - LLMProvider (any implementation — Anthropic, OpenAI, Ollama)

Data flow:
    CortexPlugin.remember()
      → ExtractionPipeline.extract(messages)
          → CodingNoiseFilter.filter_conversation()   (cheap regex)
          → _split_batches()                           (chunking)
          → for each batch:
              → format_conversation()                  (text formatting)
              → LLMProvider.complete()                 (LLM call)
              → _parse_claims()                        (JSON → ClaimExtraction list)
      → ExtractionResult
          → passed to ReconciliationEngine.reconcile() (M4)
"""

from __future__ import annotations

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

from cortex.extraction.coding_noise import CodingNoiseFilter
from cortex.extraction.prompts import EXTRACTION_SYSTEM_PROMPT, format_conversation
from cortex.types import (
    VALID_CATEGORIES,
    CATEGORY_ALIASES,
    CATEGORY_TO_INTRINSIC_MAP,
    HUMAN_FAMILIES,
    VALID_ENTITY_TYPES,
    ENTITY_TYPE_ALIASES,
)

# Backwards-compatible aliases (internal references use underscore-prefixed names)
_VALID_CATEGORIES = VALID_CATEGORIES
_CATEGORY_ALIASES = CATEGORY_ALIASES

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Entity type normalization
# ---------------------------------------------------------------------------

def _normalize_entity_type(ent: Dict[str, Any]) -> str:
    """Extract and normalize entity type from an LLM-returned entity dict.

    Tries ``entity_type`` first (canonical field name), then falls back to
    ``type`` and ``category`` for backward compatibility.  The raw value is
    resolved through ENTITY_TYPE_ALIASES and validated against
    VALID_ENTITY_TYPES.  Returns ``"unknown"`` only when no usable type can
    be determined.
    """
    # Try canonical field first, then common LLM variations
    raw = (
        ent.get("entity_type")
        or ent.get("type")
        or ent.get("category")
        or ""
    )
    normalized = str(raw).strip().lower()
    if not normalized:
        return "unknown"
    # Resolve alias (e.g. "place" → "location", "org" → "organization")
    normalized = ENTITY_TYPE_ALIASES.get(normalized, normalized)
    # Validate against canonical set + human families (used for relationship detection)
    if normalized in VALID_ENTITY_TYPES or normalized in HUMAN_FAMILIES:
        return normalized
    return "unknown"


# ---------------------------------------------------------------------------
# Output types
# ---------------------------------------------------------------------------

@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 | engineering | goals | 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

    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,
        )


@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)


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

class ExtractionPipeline:
    """
    Extracts structured claims from a conversation using LLM + noise filter.

    Supports batch extraction: when more messages than ``batch_size`` are
    provided, the pipeline splits them into chunks and issues one LLM call per
    chunk.  Claims from all batches are aggregated into a single
    ExtractionResult.  This prevents enormous context windows on long
    conversations while keeping per-call latency bounded.

    Usage::

        pipeline = ExtractionPipeline(llm=my_llm_provider, model="gpt-4.1-nano",
                                      batch_size=5)
        result = await pipeline.extract(messages, session_id="abc")
        for claim in result.claims:
            print(claim.action, claim.subject, claim.predicate,
                  claim.object_value, claim.temporal_marker)
    """

    DEFAULT_BATCH_SIZE = 5

    def __init__(
        self,
        llm: Any,
        model: Optional[str] = None,
        noise_filter: Optional[CodingNoiseFilter] = None,
        batch_size: int = DEFAULT_BATCH_SIZE,
        custom_instructions: str = "",
        max_claims_per_extraction: int = 25,
        max_entities_per_extraction: int = 100,
    ) -> None:
        """Initialise the extraction pipeline.

        Args:
            llm: LLMProvider instance used for LLM completions.
            model: Optional model override forwarded to llm.complete().
            noise_filter: CodingNoiseFilter instance; a default is created when
                          None is passed.
            batch_size: Maximum number of messages per LLM call.  Must be ≥ 1.
                        Defaults to 5.  Set to a large number (e.g. 9999) to
                        disable batching and process all messages in one call.
            custom_instructions: Optional deployment-specific rules appended
                                 after EXTRACTION_SYSTEM_PROMPT.  When non-empty,
                                 the string is separated from the base prompt by a
                                 blank line and a section header for readability.
                                 Sourced from CortexConfig.extraction_custom_instructions.
        """
        if batch_size < 1:
            raise ValueError(f"batch_size must be >= 1, got {batch_size}")
        if max_claims_per_extraction < 0:
            raise ValueError(
                f"max_claims_per_extraction must be >= 0, got {max_claims_per_extraction}"
            )
        if max_entities_per_extraction < 0:
            raise ValueError(
                f"max_entities_per_extraction must be >= 0, got {max_entities_per_extraction}"
            )
        self.llm = llm
        self.model = model
        self.noise_filter = noise_filter or CodingNoiseFilter()
        self.batch_size = batch_size
        self.max_claims_per_extraction = max_claims_per_extraction
        self.max_entities_per_extraction = max_entities_per_extraction
        raw_instructions = custom_instructions.strip()

        # Sanitize custom_instructions to prevent prompt injection
        if raw_instructions:
            original_len = len(raw_instructions)
            # Cap length to prevent prompt flooding
            if len(raw_instructions) > 2000:
                raw_instructions = raw_instructions[:2000]
                logger.warning(
                    "ExtractionPipeline: custom_instructions truncated from %d to 2000 chars",
                    original_len,
                )

            # Strip common injection patterns (lines starting with control tokens)
            # Note: "Ignore" is only matched when followed by injection keywords to
            # avoid over-blocking legitimate instructions like "Ignore work topics."
            import re as _re
            _INJECTION_RE = _re.compile(
                r'^(##|System:|You are|Ignore\s+(previous|prior|all|above|the above|everything|instructions))',
                _re.IGNORECASE,
            )
            sanitized_lines = []
            removed_lines = []
            for line in raw_instructions.splitlines():
                stripped = line.lstrip()
                if _INJECTION_RE.match(stripped):
                    removed_lines.append(line)
                else:
                    sanitized_lines.append(line)
            if removed_lines:
                logger.warning(
                    "ExtractionPipeline: custom_instructions stripped %d potentially injected line(s): %r",
                    len(removed_lines),
                    removed_lines[:3],  # log up to 3 examples
                )
                raw_instructions = "\n".join(sanitized_lines)

        self.custom_instructions = raw_instructions

        # Build the effective system prompt once (avoids repeated string ops per batch).
        if self.custom_instructions:
            self._system_prompt = (
                EXTRACTION_SYSTEM_PROMPT
                + "\n\n## Operator Instructions\n"
                + self.custom_instructions
            )
        else:
            self._system_prompt = EXTRACTION_SYSTEM_PROMPT

    async def extract(
        self,
        messages: List[Dict[str, Any]],
        session_id: str = "",
    ) -> ExtractionResult:
        """Run the full extraction pipeline on a batch of conversation messages.

        Pipeline stages:
            1. CodingNoiseFilter — remove build output, diffs, stack traces (free)
            2. Batch splitting — divide filtered messages into chunks of batch_size
            3. For each chunk: format_conversation → LLM completion → JSON parse
            4. Aggregate claims from all chunks into a single ExtractionResult

        Args:
            messages: List of message dicts with 'role' and 'content' keys.
                      Accepts any dict shape; missing keys handled gracefully.
            session_id: Session identifier for logging context.

        Returns:
            ExtractionResult containing extracted claims, filter metrics, and
            the number of batches processed.

        Raises:
            Exception: Re-raised from LLM provider if any completion call fails.
                       The caller (CortexPlugin.remember) catches and logs this.

        Note:
            Output feeds into ReconciliationEngine.reconcile() (M4) via
            CortexPlugin.remember().
        """
        raw_count = len(messages)

        # Step 1: filter noise
        filtered = self.noise_filter.filter_conversation(messages)
        noise_filtered = raw_count - len(filtered)

        if not filtered:
            logger.debug("session=%s all messages filtered as noise", session_id)
            return ExtractionResult(
                claims=[],
                noise_filtered=noise_filtered,
                raw_message_count=raw_count,
                batches_processed=0,
            )

        # Step 2: split into batches
        batches = _split_batches(filtered, self.batch_size)
        all_claims: List[ClaimExtraction] = []

        # Step 3: process each batch
        for batch_idx, batch in enumerate(batches):
            conversation_text = format_conversation(batch)
            if not conversation_text.strip():
                logger.debug(
                    "session=%s batch=%d/%d empty after formatting — skipping",
                    session_id, batch_idx + 1, len(batches),
                )
                continue

            kwargs: Dict[str, Any] = {"response_format": "json"}
            # Model routing (#375): use routed model for extraction op
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self.llm, "extraction", self.model)
            if _routed_model:
                kwargs["model"] = _routed_model

            try:
                _resp = await self.llm.complete(
                    system=self._system_prompt,
                    user=conversation_text,
                    **kwargs,
                )
                raw_json = _resp.text if hasattr(_resp, "text") else str(_resp)
            except Exception as exc:
                logger.error(
                    "session=%s batch=%d/%d LLM extraction failed: %s",
                    session_id, batch_idx + 1, len(batches), exc,
                )
                raise

            logger.debug(
                "session=%s batch=%d/%d raw_llm_response=%r",
                session_id, batch_idx + 1, len(batches), raw_json,
            )
            batch_claims = self._parse_claims(raw_json, session_id)
            all_claims.extend(batch_claims)

            logger.debug(
                "session=%s batch=%d/%d extracted=%d",
                session_id, batch_idx + 1, len(batches), len(batch_claims),
            )

        if self.max_claims_per_extraction > 0 and len(all_claims) > self.max_claims_per_extraction:
            dropped_claims = len(all_claims) - self.max_claims_per_extraction
            logger.warning(
                "session=%s extraction claim cap applied: %d -> %d (dropped=%d)",
                session_id,
                len(all_claims),
                self.max_claims_per_extraction,
                dropped_claims,
            )
            all_claims = all_claims[:self.max_claims_per_extraction]

        if self.max_entities_per_extraction > 0:
            all_claims, trimmed_entities = _cap_entities_per_extraction(
                all_claims,
                self.max_entities_per_extraction,
            )
            if trimmed_entities > 0:
                logger.warning(
                    "session=%s extraction entity cap applied: trimmed=%d (max_entities_per_extraction=%d)",
                    session_id,
                    trimmed_entities,
                    self.max_entities_per_extraction,
                )

        logger.debug(
            "session=%s total_extracted=%d noise_filtered=%d batches=%d",
            session_id, len(all_claims), noise_filtered, len(batches),
        )
        return ExtractionResult(
            claims=all_claims,
            noise_filtered=noise_filtered,
            raw_message_count=raw_count,
            batches_processed=len(batches),
        )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _parse_claims(self, raw_json: str, session_id: str) -> List[ClaimExtraction]:
        """Parse LLM JSON output into ClaimExtraction list.

        Tolerant of common LLM output variations:
        - Bare JSON array: [...]
        - Wrapped in dict: {"claims": [...]}
        - Extra fields in claim objects (ignored)
        - Missing optional fields (defaults applied)
        - NOOP claims (silently skipped)

        Args:
            raw_json: Raw JSON string from LLM completion.
            session_id: Session identifier for log context.

        Returns:
            List of valid ClaimExtraction objects (NOOP filtered out).
        """
        # Strip markdown code fences (common with models served via custom
        # OpenAI-compatible endpoints like Bailian/DashScope)
        _stripped = raw_json.strip()
        if _stripped.startswith("```"):
            # Remove opening fence (```json or ```)
            first_nl = _stripped.find("\n")
            if first_nl != -1:
                _stripped = _stripped[first_nl + 1:]
            # Remove closing fence
            if _stripped.rstrip().endswith("```"):
                _stripped = _stripped.rstrip()[:-3].rstrip()
            raw_json = _stripped

        try:
            data = json.loads(raw_json)
        except json.JSONDecodeError as exc:
            logger.warning("session=%s JSON parse error: %s | raw=%r",
                           session_id, exc, raw_json[:200])
            return []

        if isinstance(data, dict):
            # Accept {"claims": [...]} wrapper — common LLM output shape
            if "claims" in data and isinstance(data["claims"], list):
                data = data["claims"]
            elif "subject" in data and "predicate" in data:
                # Single claim dict (not wrapped in a list) — wrap it
                data = [data]
            else:
                logger.warning("session=%s unexpected JSON dict shape (no 'claims' key): %r",
                               session_id, list(data.keys())[:5])
                return []

        if not isinstance(data, list):
            logger.warning("session=%s unexpected JSON shape (not a list): %r",
                           session_id, type(data))
            return []

        claims: List[ClaimExtraction] = []
        for item in data:
            if not isinstance(item, dict):
                continue
            try:
                claim = self._dict_to_claim(item)
                if claim.action == "NOOP":
                    continue  # skip NOOP — nothing to persist
                claims.append(claim)
            except (KeyError, TypeError, ValueError) as exc:
                logger.debug("session=%s skipping malformed claim %r: %s",
                             session_id, item, exc)
                continue

        return claims

    @staticmethod
    def _dict_to_claim(d: Dict[str, Any]) -> ClaimExtraction:
        action = str(d.get("action", "ADD")).upper()
        if action not in ("ADD", "UPDATE", "DELETE", "NOOP"):
            action = "ADD"

        entities: List[EntityMention] = []
        for ent in d.get("entities_mentioned", []):
            if isinstance(ent, dict):
                entities.append(EntityMention(
                    name=str(ent.get("name", "")),
                    entity_type=_normalize_entity_type(ent),
                    aliases=list(ent.get("aliases", [])),
                ))

        # temporal_marker: accept null / missing → None; non-empty string kept as-is
        raw_marker = d.get("temporal_marker", None)
        temporal_marker: Optional[str] = (
            str(raw_marker).strip() if raw_marker is not None and str(raw_marker).strip()
            else None
        )

        # category: validate against known set, resolving aliases first.
        raw_category = str(d.get("category", "misc")).strip().lower()
        # Resolve alias before validity check (e.g. "technical" → "engineering")
        resolved_category = _CATEGORY_ALIASES.get(raw_category, raw_category)
        category = resolved_category if resolved_category in _VALID_CATEGORIES else "misc"
        if category == "misc" and raw_category != "misc":
            logger.warning(
                "Category %r not in valid set %s, fell through to 'misc' (claim: %s %s)",
                raw_category, _VALID_CATEGORIES, d.get("subject", "?"), d.get("predicate", "?"),
            )

        confidence = float(d.get("confidence", 0.5))
        temporal = str(d.get("temporal", "current"))
        intrinsic_type = CATEGORY_TO_INTRINSIC_MAP.get(category, "fact")

        entity_types = {
            str(getattr(ent, "entity_type", "")).strip().lower()
            for ent in entities
            if str(getattr(ent, "entity_type", "")).strip()
        }
        human_family = next(
            (family for family in HUMAN_FAMILIES if family in entity_types),
            None,
        )

        if confidence >= 0.9:
            explicitness = "explicit"
        elif confidence >= 0.7:
            explicitness = "implied"
        else:
            explicitness = "inferred"

        temporal_key = temporal.strip().lower()
        if temporal_key == "permanent":
            stability = "stable"
        elif temporal_key in {"current", "future"}:
            stability = "mutable"
        else:
            stability = "volatile"

        # --- New LoCoMo-aligned fields (v1.3) ---
        memory_class = str(d.get("memory_class", "semantic")).strip().lower()
        if memory_class not in {"semantic", "episodic", "procedural", "open_loop", "transient"}:
            memory_class = "semantic"

        status = str(d.get("status", "active")).strip().lower()
        if status not in {"active", "completed", "superseded", "blocked"}:
            status = "active"

        raw_event_date = d.get("event_date", None)
        event_date: Optional[str] = None
        if raw_event_date is not None:
            date_str = str(raw_event_date).strip()
            # Accept YYYY-MM-DD format only
            if re.match(r'^\d{4}-\d{2}-\d{2}$', date_str):
                event_date = date_str

        raw_seq = d.get("sequence_order", None)
        sequence_order: Optional[int] = None
        if raw_seq is not None:
            try:
                sequence_order = int(raw_seq)
            except (ValueError, TypeError):
                pass

        salience = str(d.get("salience", "medium")).strip().lower()
        if salience not in {"high", "medium", "low"}:
            salience = "medium"

        # Causal linking (#518)
        cause = d.get("cause", None)
        if cause is not None:
            cause = str(cause).strip() or None

        effect = d.get("effect", None)
        if effect is not None:
            effect = str(effect).strip() or None

        # State transitions (#519)
        state_before = d.get("state_before", None)
        if state_before is not None:
            state_before = str(state_before).strip() or None

        state_after = d.get("state_after", None)
        if state_after is not None:
            state_after = str(state_after).strip() or None

        return ClaimExtraction(
            action=action,
            subject=str(d["subject"]),
            predicate=str(d["predicate"]),
            object_value=str(d["object_value"]),
            confidence=confidence,
            temporal=temporal,
            temporal_marker=temporal_marker,
            entities_mentioned=entities,
            sensitivity=str(d.get("sensitivity", "normal")),
            category=category,
            intrinsic_type=intrinsic_type,
            human_family=human_family,
            explicitness=explicitness,
            stability=stability,
            memory_class=memory_class,
            status=status,
            event_date=event_date,
            sequence_order=sequence_order,
            salience=salience,
            cause=cause,
            effect=effect,
            state_before=state_before,
            state_after=state_after,
        )


# ---------------------------------------------------------------------------
# Module-level helpers
# ---------------------------------------------------------------------------

def _cap_entities_per_extraction(
    claims: List[ClaimExtraction],
    max_entities: int,
) -> tuple[List[ClaimExtraction], int]:
    """Cap total entities across all claims and return (claims, trimmed_count)."""
    if max_entities <= 0:
        return claims, 0

    remaining = max_entities
    trimmed = 0
    for claim in claims:
        entities = list(getattr(claim, "entities_mentioned", []) or [])
        if not entities:
            continue
        if remaining <= 0:
            trimmed += len(entities)
            claim.entities_mentioned = []
            continue
        if len(entities) > remaining:
            trimmed += len(entities) - remaining
            claim.entities_mentioned = entities[:remaining]
            remaining = 0
        else:
            remaining -= len(entities)
    return claims, trimmed


def _split_batches(
    messages: List[Dict[str, Any]],
    batch_size: int,
) -> List[List[Dict[str, Any]]]:
    """Split *messages* into consecutive chunks of at most *batch_size* items.

    Args:
        messages: Flat list of message dicts (already noise-filtered).
        batch_size: Maximum messages per chunk.  Must be ≥ 1.

    Returns:
        List of chunks.  Each chunk is a non-empty sublist of *messages*.
        Returns a list with a single empty list only when *messages* is empty.

    Example:
        >>> _split_batches([m1, m2, m3, m4, m5], batch_size=2)
        [[m1, m2], [m3, m4], [m5]]
    """
    if not messages:
        return [[]]
    n = math.ceil(len(messages) / batch_size)
    return [messages[i * batch_size:(i + 1) * batch_size] for i in range(n)]
