"""
cortex/capture/wake_pipeline.py — Extraction pipeline orchestration.

Contains the core logic for the remember() pipeline steps 1-4:
  1. Extraction       — LLM-based claim extraction from conversation
  2. Dedup            — exact/near-duplicate filtering
  3. Entity resolution — resolve entity mentions to canonical IDs
  4. Reconciliation   — merge/supersede existing claims in storage
  4b. Activation      — derive + persist activation policy per stored claim (W3, #358)
  + Embedding         — compute and store vector embeddings for new claims

Also owns the Telegram/OpenClaw transport metadata stripping helpers that
are applied to conversation input before extraction.

Used by: CortexPlugin.remember()
"""
from __future__ import annotations

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

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Transport metadata stripping (Telegram/OpenClaw preamble removal)
# ---------------------------------------------------------------------------

# Telegram/OpenClaw transport metadata preamble patterns.
_TELEGRAM_META_FENCE_RE = re.compile(
    r"(?is)^\s*sender\s*\(untrusted\s+metadata\)\s*:\s*```(?:json)?\s*.*?```\s*"
)
_TELEGRAM_META_JSON_RE = re.compile(
    r"(?is)^\s*```(?:json)?\s*\{(?:(?!```)[\s\S])*?\b(?:chat_id|sender|message_id|inbound_meta)\b(?:(?!```)[\s\S])*?\}\s*```\s*"
)


def _contains_telegram_meta_keys(value: Any) -> bool:
    keys = {"chat_id", "sender", "message_id", "inbound_meta"}
    if isinstance(value, dict):
        if keys.intersection(value.keys()):
            return True
        return any(_contains_telegram_meta_keys(v) for v in value.values())
    if isinstance(value, list):
        return any(_contains_telegram_meta_keys(v) for v in value)
    return False


def strip_transport_metadata_text(text: str) -> str:
    """Strip known Telegram/OpenClaw metadata preambles from message text."""
    if not text:
        return text

    cleaned = text
    for _ in range(3):  # allow stacked metadata blocks
        prev = cleaned
        cleaned = _TELEGRAM_META_FENCE_RE.sub("", cleaned, count=1)
        cleaned = _TELEGRAM_META_JSON_RE.sub("", cleaned, count=1)

        # Handle raw JSON metadata object at the start (non-fenced).
        stripped = cleaned.lstrip()
        leading_ws = len(cleaned) - len(stripped)
        if stripped.startswith("{"):
            try:
                obj, end_idx = json.JSONDecoder().raw_decode(stripped)
                if isinstance(obj, dict) and _contains_telegram_meta_keys(obj):
                    cleaned = cleaned[:leading_ws] + stripped[end_idx:]
            except Exception as exc:
                logger.debug("Wake pipeline metadata parse skipped: %s", exc)

        if cleaned == prev:
            break

    return cleaned.strip()


def strip_transport_metadata_conversation(
    conversation: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
    """Return conversation with metadata preambles stripped from string content."""
    cleaned_messages: List[Dict[str, Any]] = []
    for msg in conversation:
        if not isinstance(msg, dict):
            cleaned_messages.append(msg)
            continue
        msg_copy = dict(msg)
        content = msg_copy.get("content")
        if isinstance(content, str):
            msg_copy["content"] = strip_transport_metadata_text(content)
        cleaned_messages.append(msg_copy)
    return cleaned_messages


def conversation_has_signal(conversation: List[Dict[str, Any]]) -> bool:
    """Heuristic: conversation has at least one non-empty text message."""
    for msg in conversation:
        if not isinstance(msg, dict):
            continue
        content = msg.get("content")
        if isinstance(content, str) and content.strip():
            return True
    return False


# ---------------------------------------------------------------------------
# Config helper (duplicated from plugin.py to avoid circular imports)
# ---------------------------------------------------------------------------

def _config_int(config: Any, key: str, default: int) -> int:
    """Read integer config safely, avoiding MagicMock/non-numeric fallthrough."""
    value = getattr(config, key, default)
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return default
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


# ---------------------------------------------------------------------------
# Outcome dataclass
# ---------------------------------------------------------------------------

@dataclass
class ExtractionOutcome:
    """Result returned by run_extraction() to CortexPlugin.remember()."""

    claims_extracted: int = 0
    claims_stored: int = 0
    entities_resolved: int = 0
    noise_filtered: int = 0
    errors: List[str] = field(default_factory=list)
    warnings: List[str] = field(default_factory=list)
    zero_claim_warning: bool = False
    ok: bool = True
    #: Updated counter to write back to plugin._zero_claim_outcomes
    updated_zero_claim_outcomes: int = 0


# ---------------------------------------------------------------------------
# Main extraction pipeline function
# ---------------------------------------------------------------------------

async def run_extraction(
    *,
    extractor: Any,
    storage: Any,
    llm: Any,
    entity_resolver: Any,
    cornerstone_guardian: Any,
    reconciler: Any,
    event_emitter: Any,
    config: Any,
    extraction_semaphore: Any,
    conversation: List[Dict[str, Any]],
    sid: str,
    db_session_pk: Optional[str],
    effective_owner_id: str,
    episode_event_id: Optional[str],
    budget_check: Callable[[int], Optional[str]],
    budget_consume: Callable[[int], None],
    errors: List[str],
    warnings: List[str],
    zero_claim_outcomes: int,
) -> ExtractionOutcome:
    """Run the extraction → dedup → entity resolution → reconcile → embed pipeline.

    This function encapsulates remember() steps 1-4.  All external mutations
    (storage writes, event emits) happen inside here; the caller only needs to
    handle the session setup (step 0) and result assembly.

    Args:
        extractor:           ExtractionPipeline (may be None → error recorded).
        storage:             StorageAdapter.
        llm:                 LLMProvider for embedding.
        entity_resolver:     EntityResolver (may be None → skipped).
        cornerstone_guardian: CornerstoneGuardian (may be None → skipped).
        reconciler:          ReconciliationEngine (may be None → error recorded).
        event_emitter:       EventEmitter for MEMORY_CREATED/UPDATED events.
        config:              CortexConfig.
        extraction_semaphore: asyncio.Semaphore for concurrency cap.
        conversation:        Cleaned/stripped conversation list.
        sid:                 Human-readable session ID.
        db_session_pk:       Internal DB PK of the session_record row (for FK).
        effective_owner_id:  Memory owner namespace.
        episode_event_id:    ID of the already-logged episode event (for provenance).
        budget_check:        Callable(n) → error_str | None.
        budget_consume:      Callable(n) → None.
        errors:              Mutable error list (pre-populated by caller).
        warnings:            Mutable warning list (pre-populated by caller).
        zero_claim_outcomes: Current zero-claim counter from plugin state.

    Returns:
        ExtractionOutcome with all counts and updated state.
    """
    claims_extracted = 0
    claims_stored = 0
    entities_resolved = 0
    noise_filtered = 0
    zero_claim_warning = False

    # 1. Extraction -----------------------------------------------------------
    extraction_result = None
    if extractor is None:
        err = "ExtractionPipeline not available — skipping extraction"
        logger.warning(err)
        errors.append(err)
    else:
        budget_err = budget_check(1)
        if budget_err:
            logger.error("Skipping extraction: %s", budget_err)
            errors.append(budget_err)
            return ExtractionOutcome(
                claims_extracted=0,
                claims_stored=0,
                entities_resolved=0,
                noise_filtered=noise_filtered,
                errors=errors,
                warnings=warnings,
                zero_claim_warning=False,
                ok=False,
                updated_zero_claim_outcomes=zero_claim_outcomes,
            )
        try:
            budget_consume(1)
            async with extraction_semaphore:
                extraction_result = await extractor.extract(
                    conversation,
                    session_id=sid,
                )
            claims_extracted = len(extraction_result.claims)
            noise_filtered = extraction_result.noise_filtered
            logger.debug(
                "Extracted %d claims (noise filtered: %d)",
                claims_extracted,
                noise_filtered,
            )
        except Exception as exc:
            err = f"Extraction failed: {exc}"
            logger.error(err, exc_info=True)
            errors.append(err)

    if extraction_result is None or not extraction_result.claims:
        if extraction_result is not None and conversation_has_signal(conversation) and not errors:
            zero_claim_outcomes += 1
            zero_claim_warning = True
            msg = (
                "Extraction produced zero claims for a valid non-empty input "
                f"(count={zero_claim_outcomes})"
            )
            warnings.append(msg)
            logger.warning("session=%s %s", sid, msg)
        return ExtractionOutcome(
            claims_extracted=claims_extracted,
            claims_stored=0,
            entities_resolved=0,
            noise_filtered=noise_filtered,
            errors=errors,
            warnings=warnings,
            zero_claim_warning=zero_claim_warning,
            ok=not errors,
            updated_zero_claim_outcomes=zero_claim_outcomes,
        )

    # 1a-post. Stamp episode_event_id on each candidate ----------------------
    if episode_event_id is not None:
        for _cand in extraction_result.claims:
            try:
                object.__setattr__(_cand, "episode_event_id", episode_event_id)
            except (TypeError, AttributeError):
                _cand.episode_event_id = episode_event_id  # type: ignore[attr-defined]

    # 1b. Per-batch extraction cap -------------------------------------------
    batch_limit = _config_int(
        config,
        "max_claims_per_extraction",
        _config_int(config, "max_extraction_calls_per_batch", 25),
    )
    if batch_limit > 0 and len(extraction_result.claims) > batch_limit:
        _SENSITIVITY_PRIORITY = {"sensitive": 0, "personal": 1, "normal": 2}
        extraction_result.claims.sort(
            key=lambda c: _SENSITIVITY_PRIORITY.get(
                getattr(c, "sensitivity", "normal"), 2
            )
        )
        dropped = len(extraction_result.claims) - batch_limit
        logger.warning(
            "Capping extraction batch: %d claims → %d "
            "(max_claims_per_extraction=%d, dropped %d, sensitive-first sort applied)",
            len(extraction_result.claims),
            batch_limit,
            batch_limit,
            dropped,
        )
        extraction_result.claims = extraction_result.claims[:batch_limit]

    # 1c. Extraction-time deduplication --------------------------------------
    if extraction_result.claims and storage is not None:
        try:
            from cortex.extraction.dedup import deduplicate_claims
            dedup_result = await deduplicate_claims(
                storage=storage,
                new_claims=extraction_result.claims,
                owner_id=effective_owner_id,
                threshold=0.85,
            )
            for nc, ec in dedup_result.to_merge:
                try:
                    new_confidence = min(1.0, ec.confidence + nc.confidence * 0.1)
                    await storage.update_claim(
                        ec.id,
                        confidence=new_confidence,
                    )
                    logger.debug(
                        "dedup merge: updated claim id=%s confidence %.2f → %.2f",
                        ec.id,
                        ec.confidence,
                        new_confidence,
                    )
                except Exception as merge_exc:
                    logger.warning(
                        "dedup: failed to update merged claim id=%s: %s",
                        ec.id,
                        merge_exc,
                    )
            if dedup_result.total_deduped > 0:
                logger.info(
                    "Extraction dedup: %d skipped (exact), %d merged (similar), %d new",
                    len(dedup_result.skipped),
                    len(dedup_result.to_merge),
                    len(dedup_result.to_store),
                )
            extraction_result.claims = dedup_result.to_store
        except Exception as dedup_exc:
            logger.warning(
                "Extraction dedup failed (non-fatal) — proceeding with all claims: %s",
                dedup_exc,
            )

    if not extraction_result.claims:
        return ExtractionOutcome(
            claims_extracted=claims_extracted,
            claims_stored=0,
            entities_resolved=0,
            noise_filtered=noise_filtered,
            errors=errors,
            warnings=warnings,
            zero_claim_warning=zero_claim_warning,
            ok=not errors,
            updated_zero_claim_outcomes=zero_claim_outcomes,
        )

    # 2. Entity resolution ---------------------------------------------------
    resolved_claims = extraction_result.claims
    if entity_resolver is None:
        logger.warning("EntityResolver not available — skipping entity resolution")
        errors.append("EntityResolver not available")
    else:
        resolved_count = 0
        for claim in extraction_result.claims:
            for mention in getattr(claim, "entities_mentioned", []):
                try:
                    entity = await entity_resolver.resolve(
                        mention=mention.name,
                        owner_id=effective_owner_id,
                        entity_type=getattr(mention, "entity_type", "unknown"),
                    )
                    resolved_count += 1
                    if event_emitter is not None and entity is not None:
                        try:
                            from cortex.events.emitter import CortexEvent, EventType
                            await event_emitter.emit(CortexEvent(
                                event_type=EventType.ENTITY_CREATED,
                                owner_id=effective_owner_id,
                                payload={
                                    "entity_id": getattr(entity, "id", str(entity)),
                                    "entity_name": mention.name,
                                },
                                session_id=sid,
                            ))
                        except Exception as emit_exc:
                            logger.debug("Event emit failed (ENTITY_CREATED): %s", emit_exc)
                except Exception as exc:
                    err = f"Entity resolution failed for '{mention.name}': {exc}"
                    logger.warning(err)
                    errors.append(err)
        entities_resolved = resolved_count

    # 2b. Peer auto-create ---------------------------------------------------
    if storage is not None and entity_resolver is not None:
        try:
            from cortex.peers import PeerManager
            peer_mgr = PeerManager(
                storage=storage,
                entity_resolver=entity_resolver,
            )
            if hasattr(storage, "list_entities"):
                all_entities = await storage.list_entities(
                    owner_id=effective_owner_id,
                    entity_type="person",
                    limit=500,
                )
                await peer_mgr.auto_create_from_entities(
                    owner_id=effective_owner_id,
                    entities=all_entities,
                )
        except Exception as exc:
            logger.warning("Peer auto-create failed (non-fatal): %s", exc)

    # 2c. Cornerstone guard --------------------------------------------------
    if cornerstone_guardian is not None and resolved_claims:
        guarded_claims = []
        for _cand in resolved_claims:
            _subject = getattr(_cand, "subject", "")
            try:
                _protected = await cornerstone_guardian.is_cornerstone_subject(
                    subject=_subject,
                    owner_id=effective_owner_id,
                )
            except Exception as _cs_exc:
                logger.warning(
                    "Cornerstone subject check failed for %r (allowing claim): %s",
                    _subject,
                    _cs_exc,
                )
                _protected = False
            if _protected:
                logger.info(
                    "Cornerstone guard: dropping candidate (subject=%r) "
                    "— touches a protected cornerstone",
                    _subject,
                )
            else:
                guarded_claims.append(_cand)
        resolved_claims = guarded_claims

    # 2d. Convert ClaimExtraction -> CandidateClaim (formalise type contract) --
    # ClaimExtraction is the M2 output type; CandidateClaim is the M4 input
    # contract.  Converting here (after entity resolution, before reconciliation)
    # ensures entity resolution operates on ClaimExtraction objects which have
    # entities_mentioned, while reconciliation receives properly-typed
    # CandidateClaim objects rather than duck-typed instances.
    from cortex.extraction.extractor import ClaimExtraction as _ClaimExtraction
    _candidate_claims = []
    for _raw_claim in resolved_claims:
        if isinstance(_raw_claim, _ClaimExtraction):
            _cand = _raw_claim.to_candidate()
            # Carry through any dynamically-stamped attributes (e.g.
            # episode_event_id set by step 1a-post).  These don't exist
            # on the ClaimExtraction dataclass definition but are set via
            # object.__setattr__() earlier in this pipeline.
            _episode_id = getattr(_raw_claim, "episode_event_id", None)
            if _episode_id is not None:
                _cand.episode_event_id = _episode_id  # type: ignore[attr-defined]
            _candidate_claims.append(_cand)
        else:
            # Already a CandidateClaim (e.g., from tests or direct API callers)
            _candidate_claims.append(_raw_claim)
    resolved_claims = _candidate_claims  # type: ignore[assignment]

    # 3. Reconciliation (writes to storage) ----------------------------------
    reconciliation_report = None
    if reconciler is None:
        err = "ReconciliationEngine not available — memories NOT stored"
        logger.error(err)
        errors.append(err)
    else:
        reconcile_calls = len(resolved_claims)
        budget_err = budget_check(reconcile_calls)
        if budget_err:
            logger.error("Skipping reconciliation: %s", budget_err)
            errors.append(budget_err)
        else:
            budget_consume(reconcile_calls)
            try:
                reconciliation_report = await reconciler.reconcile(
                    candidates=resolved_claims,
                    owner_id=effective_owner_id,
                    session_id=db_session_pk or sid,
                )
                claims_stored = sum(
                    1 for r in reconciliation_report.results
                    if r.action.value in ("ADD", "UPDATE", "SUPERSEDE")
                )
                logger.debug(
                    "Reconciled: %d stored, %d total results",
                    claims_stored,
                    len(reconciliation_report.results),
                )
                # Emit memory events
                if event_emitter is not None:
                    try:
                        from cortex.events.emitter import CortexEvent, EventType
                        for r in reconciliation_report.results:
                            action_val = r.action.value
                            claim_id = getattr(r, "resulting_claim_id", None)
                            if action_val == "ADD" and claim_id:
                                await event_emitter.emit(CortexEvent(
                                    event_type=EventType.MEMORY_CREATED,
                                    owner_id=effective_owner_id,
                                    payload={
                                        "claim_id": claim_id,
                                        "subject": getattr(r, "subject", ""),
                                        "predicate": getattr(r, "predicate", ""),
                                    },
                                    session_id=sid,
                                ))
                            elif action_val in ("UPDATE", "SUPERSEDE") and claim_id:
                                await event_emitter.emit(CortexEvent(
                                    event_type=EventType.MEMORY_UPDATED,
                                    owner_id=effective_owner_id,
                                    payload={
                                        "claim_id": claim_id,
                                        "action": action_val.lower(),
                                        "subject": getattr(r, "subject", ""),
                                        "predicate": getattr(r, "predicate", ""),
                                    },
                                    session_id=sid,
                                ))
                    except Exception as emit_exc:
                        logger.debug("Event emit failed (memory events): %s", emit_exc)
            except Exception as exc:
                err = f"Reconciliation failed: {exc}"
                logger.error(err, exc_info=True)
                errors.append(err)

    # 4b. Activation routing — derive + persist activation policy per stored claim
    if storage is not None and reconciliation_report is not None:
        try:
            from cortex.capture.activation_router import ActivationRouter
            _activation_router = ActivationRouter()
            # Build a fast lookup: resulting_claim_id → candidate claim
            _claim_to_candidate: dict = {}
            for _idx, _cand in enumerate(resolved_claims):
                # Match by position in the results list (reconciliation preserves order).
                if _idx < len(reconciliation_report.results):
                    _res = reconciliation_report.results[_idx]
                    _cid = getattr(_res, "resulting_claim_id", None)
                    if _cid:
                        _claim_to_candidate[_cid] = _cand

            _activation_ok = 0
            _activation_fail = 0
            for _result in reconciliation_report.results:
                if _result.action.value not in ("ADD", "UPDATE", "SUPERSEDE"):
                    continue
                _claim_id = getattr(_result, "resulting_claim_id", None)
                if not _claim_id:
                    continue
                _cand = _claim_to_candidate.get(_claim_id)
                if _cand is None:
                    # Fallback: derive with minimal envelope (safe defaults)
                    _cand = type("_Empty", (), {})()
                _policy = _activation_router.derive(
                    candidate=_cand,
                    claim_id=_claim_id,
                    owner_id=effective_owner_id,
                )
                _written = await _activation_router.write(storage, _policy)
                if _written:
                    _activation_ok += 1
                else:
                    _activation_fail += 1

            if _activation_ok or _activation_fail:
                logger.debug(
                    "session=%s activation routing: written=%d skipped=%d",
                    sid,
                    _activation_ok,
                    _activation_fail,
                )
        except Exception as _act_exc:
            # Activation routing is non-fatal: never block the remember() path.
            logger.warning(
                "session=%s activation routing failed (non-fatal): %s",
                sid,
                _act_exc,
            )

    # 4c. Embed stored claims for vector retrieval ---------------------------
    if storage is not None and llm is not None and reconciliation_report is not None:
        embed_model = config.embedding_model
        for result in reconciliation_report.results:
            if result.action.value not in ("ADD", "UPDATE", "SUPERSEDE"):
                continue
            claim_id = getattr(result, "resulting_claim_id", None)
            if not claim_id:
                continue
            try:
                stored_claim = await storage.get_claim(claim_id)
                if stored_claim is None:
                    continue
                embed_text = (
                    f"{stored_claim.subject} "
                    f"{stored_claim.predicate} "
                    f"{stored_claim.object_value}"
                )
                embed_response = await llm.embed(
                    texts=[embed_text],
                    input_type="document",
                )
                if embed_response and embed_response.embeddings:
                    vector = embed_response.embeddings[0]
                    await storage.store_embedding(
                        item_type="claim",
                        item_id=claim_id,
                        embedding=vector,
                        model_id=embed_model,
                    )
                    logger.debug("Embedded claim %s (%d dims)", claim_id, len(vector))
            except Exception as exc:
                err = f"Embedding failed for claim {claim_id}: {exc}"
                logger.warning(err)
                errors.append(err)

    return ExtractionOutcome(
        claims_extracted=claims_extracted,
        claims_stored=claims_stored,
        entities_resolved=entities_resolved,
        noise_filtered=noise_filtered,
        errors=errors,
        warnings=warnings,
        zero_claim_warning=zero_claim_warning,
        ok=not any(
            ("not available" in e or "budget exhausted" in e)
            for e in errors
        ),
        updated_zero_claim_outcomes=zero_claim_outcomes,
    )
