"""
cortex/extraction/utils.py — Utility functions for the extraction pipeline.

Contains helpers that support ExtractionPipeline but are not methods of the
class itself:

    - _is_llm_api_error: LLM error classification
    - _normalize_entity_type: Entity type normalization
    - dict_to_claim: Convert a raw LLM dict to a ClaimExtraction
    - _cap_entities_per_extraction: Post-extraction entity capping
    - _split_batches: Batch splitting utility
"""

from __future__ import annotations

import asyncio
import logging
import math
import re
from typing import Any, Dict, List, Optional

from cortex.extraction.models import ClaimExtraction, EntityMention
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

# Use the original module logger name for backwards compatibility — tests and
# log filters may capture "cortex.extraction.extractor" specifically.
logger = logging.getLogger("cortex.extraction.extractor")


def _is_llm_api_error(exc: Exception) -> bool:
    """Check if an exception originates from an LLM SDK (not a programming bug).

    Catches API-level errors from OpenAI, Anthropic, httpx, etc. that aren't
    classified as retryable by ``_is_retryable_exception`` but still shouldn't
    kill the entire extraction pipeline (e.g., content filter, bad request to
    the model endpoint, model overloaded).

    Programming errors (KeyError, TypeError, ValueError from our code, etc.)
    return False so they propagate normally.
    """
    # Catch Python/asyncio timeout errors — always LLM-related in this context (#588)
    if isinstance(exc, (TimeoutError, asyncio.TimeoutError)):
        return True

    exc_module = getattr(type(exc), "__module__", "") or ""
    _LLM_MODULES = ("openai", "anthropic", "httpx", "aiohttp", "litellm", "google")
    if any(exc_module.startswith(m) for m in _LLM_MODULES):
        return True

    exc_name = type(exc).__name__
    _LLM_ERROR_NAMES = (
        "APIError", "APIConnectionError", "APIStatusError", "APITimeoutError",
        "RateLimitError", "AuthenticationError", "BadRequestError",
        "InternalServerError", "PermissionDeniedError",
        "ContentFilterError", "ServiceUnavailableError",
    )
    if exc_name in _LLM_ERROR_NAMES:
        return True

    # RuntimeError with LLM-ish messages (e.g., from our own providers)
    if isinstance(exc, RuntimeError):
        msg = str(exc).lower()
        if any(kw in msg for kw in ("quota", "rate limit", "timeout", "overloaded",
                                     "model", "api", "provider")):
            return True

    return False


# ---------------------------------------------------------------------------
# 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"


# ---------------------------------------------------------------------------
# dict_to_claim — Convert a raw LLM dict to a ClaimExtraction
# ---------------------------------------------------------------------------

def dict_to_claim(d: Dict[str, Any]) -> ClaimExtraction:
    """Convert a single LLM-returned dict into a ClaimExtraction instance.

    This is the module-level function called by ExtractionPipeline._dict_to_claim
    static method. It handles all field validation, normalization, and default
    application.
    """
    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" → "decisions")
    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

    # Source role (#321 — split extraction mode); LLM may include it in output
    _raw_source_role = d.get("source_role", None)
    source_role: Optional[str] = None
    if _raw_source_role is not None:
        _sr = str(_raw_source_role).strip().lower()
        if _sr in ("user", "agent", "system"):
            source_role = _sr

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


# ---------------------------------------------------------------------------
# Post-extraction 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)]
