"""
cortex/core/retrieval_pipeline.py — RetrievalPipeline class.

Extracted from cortex/core/retrieval.py (#947 — SRP split).

The orchestrator for retrieval. Delegates search, scoring, reranking,
and budget enforcement to sibling modules.

Dependencies:
    - cortex.core.retrieval_search  — search backends, rerank, budget
    - cortex.core.retrieval_scoring — query classification, fusion
    - cortex.core.retrieval_types   — DTOs and constants
    - cortex.core.retrieval_helpers — content rendering, accessors
"""

from __future__ import annotations

import asyncio
import logging
import os
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

try:
    import sentry_sdk
except ImportError:  # pragma: no cover
    sentry_sdk = None  # type: ignore[assignment]

if TYPE_CHECKING:
    from cortex.core.token_budget import TokenBudgetManager
    from cortex.storage.adapter import StorageAdapter

# Embedding cache — optional, imported lazily to keep module importable
# even if cachetools is not installed.
try:
    from cortex.core.embedding_cache import EmbeddingCache as _EmbeddingCache
except ImportError:  # pragma: no cover
    _EmbeddingCache = None  # type: ignore[assignment,misc]

from cortex.core.retrieval_types import (
    ContextBlock,
    RetrievalConstraints,
    RetrievalResult,
    RetrievalTrace,
    RetrievedItem,
    _FTSResult,
    _VectorResult,
)
from cortex.core.retrieval_scoring import _is_complex_query
from cortex.core.retrieval_helpers import _new_id
try:
    import voyageai  # type: ignore
except ImportError:
    voyageai = None  # type: ignore

from cortex.core.retrieval_search import (
    DEFAULT_STATUS_PENALTIES,
    apply_budget,
    apply_status_penalty,
    bm25_search,
    cosine_fallback,
    count_tokens,
    embed_text,
    log_retrieval,
    vector_search,
    working_memory_retrieve,
)
from cortex.core.retrieval_lightweight import lightweight_retrieve
from cortex.core.retrieval_policy_routing import apply_task_mode_policy

logger = logging.getLogger(__name__)
trace_logger = logging.getLogger("cortex.retrieval.trace")


class RetrievalPipeline:
    """
    Phase 1 hybrid retrieval: BM25 (FTS5) + vector similarity, fused with RRF.

    Phase 2 additions (issue #34): query expansion, multi-round retrieval,
    LLM rerank, brain-graph traversal.
    """

    def __init__(
        self,
        storage: "StorageAdapter",
        config: Any,           # CortexConfig
        token_budget: Optional["TokenBudgetManager"] = None,
        cornerstone_guardian: Any = None,  # Phase 2
        llm: Any = None,                   # Phase 2 — agentic mode
        working_memory: Optional[Any] = None,  # WorkingMemoryCache (#371)
        embedding_cache: Optional[Any] = None,  # EmbeddingCache (v1.2.1)
        *,
        vector_weight: float = 0.7,
        bm25_weight: float = 0.3,
        rrf_k: int = 60,
        candidate_multiplier: int = 3,
        recency_weight: Optional[float] = None,
    ) -> None:
        self.storage = storage
        self.config = config
        self.token_budget = token_budget
        self.cornerstone_guardian = cornerstone_guardian
        self.llm = llm
        self.working_memory = working_memory

        if embedding_cache is not None:
            self.embedding_cache: Optional[Any] = embedding_cache
        elif _EmbeddingCache is not None:
            _ec_enabled = getattr(config, "embedding_cache_enabled", True)
            _ec_maxsize = getattr(config, "embedding_cache_max_size", 1000)
            _ec_ttl = getattr(config, "embedding_cache_ttl_seconds", 3600)
            self.embedding_cache = _EmbeddingCache(
                maxsize=_ec_maxsize, ttl=_ec_ttl, enabled=_ec_enabled,
            )
        else:
            self.embedding_cache = None

        self.vector_weight = vector_weight
        self.bm25_weight = bm25_weight
        self.rrf_k = getattr(config, "retrieval_rrf_k", rrf_k)
        self.candidate_multiplier = candidate_multiplier
        if recency_weight is not None:
            self.recency_weight = recency_weight
        else:
            self.recency_weight = getattr(config, "retrieval_recency_weight", 0.1)
        self.recency_half_life_days = getattr(config, "retrieval_recency_half_life_days", 30.0)
        self._last_retrieval_info: Optional[Dict[str, Any]] = None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def retrieve(
        self,
        query: str,
        mode: str = "auto",
        constraints: Optional[RetrievalConstraints] = None,
        owner_id: str = "default",
        task_mode_context: Optional[Dict[str, str]] = None,
        session_id: Optional[str] = None,
    ) -> RetrievalResult:
        """Main retrieval entry point — find relevant memories for a query."""
        t0 = time.monotonic()
        retrieval_id = _new_id()
        constraints = constraints or RetrievalConstraints()

        _sentry_ret_span = None
        if sentry_sdk is not None:
            try:
                _sentry_ret_span = sentry_sdk.start_span(
                    op="cortex.retrieval", description=f"retrieve({mode})",
                )
                _sentry_ret_span.set_tag("owner_id", owner_id)
                _sentry_ret_span.set_tag("mode", mode)
                _sentry_ret_span.set_data("query_length", len(query))
            except Exception:  # pragma: no cover
                _sentry_ret_span = None

        # Task-mode policy routing (#362 / #608)
        _policy_skip_wm, _policy_set_top_k, constraints = apply_task_mode_policy(
            self.config, query, task_mode_context, constraints,
        )

        # Working Memory pre-fetch (#371)
        if _policy_skip_wm:
            wm_items: List[RetrievedItem] = []
        else:
            wm_items = await self._working_memory_retrieve(query, constraints, owner_id, session_id=session_id)

        if mode == "auto":
            mode = self._route_query(query, constraints)

        # Dynamic top-k + selective rerank (#495)
        _rerank_cfg_enabled = getattr(self.config, "retrieval_rerank_enabled", False)
        _is_complex = _is_complex_query(query)
        _rerank_enabled = _rerank_cfg_enabled
        _dynamic_top_k = 10 if _is_complex else 8
        if constraints.top_k == 20 and not _policy_set_top_k:
            from copy import copy as _copy
            constraints = _copy(constraints)
            constraints.top_k = _dynamic_top_k

        _rerank_initial_candidates = getattr(
            self.config, "retrieval_rerank_initial_candidates", 15,
        )
        _retrieval_constraints = constraints
        if _rerank_enabled and _rerank_initial_candidates > constraints.top_k:
            from copy import copy as _copy
            _retrieval_constraints = _copy(constraints)
            _retrieval_constraints.top_k = _rerank_initial_candidates

        _retrieval_t0 = time.monotonic()
        if mode == "lightweight":
            items = await self._lightweight_retrieve(query, _retrieval_constraints, owner_id)
        elif mode == "chained":
            # Chained multi-hop retrieval (#1279)
            max_hops = getattr(self.config, "retrieval_max_chained_hops", 2)
            if max_hops > 0:
                from cortex.core.agentic_retrieval import chained_retrieve
                return await chained_retrieve(
                    pipeline=self,
                    query=query,
                    owner_id=owner_id,
                    max_hops=max_hops,
                    constraints=constraints,
                )
            else:
                logger.debug("Chained retrieval requested but max_hops=0; falling back to lightweight")
                items = await self._lightweight_retrieve(query, _retrieval_constraints, owner_id)
        elif mode in ("agentic", "brain_graph"):
            if self.llm is not None:
                from cortex.core.agentic_retrieval import AgenticRetrievalPipeline
                agentic = AgenticRetrievalPipeline(
                    storage=self.storage, config=self.config,
                    token_budget=self.token_budget,
                    cornerstone_guardian=self.cornerstone_guardian,
                    llm=self.llm, vector_weight=self.vector_weight,
                    bm25_weight=self.bm25_weight, rrf_k=self.rrf_k,
                    candidate_multiplier=self.candidate_multiplier,
                )
                return await agentic.retrieve(
                    query, mode=mode, constraints=constraints, owner_id=owner_id,
                )
            else:
                logger.warning(
                    "Mode '%s' requested but no LLM provided. Falling back to 'lightweight'.",
                    mode,
                )
                mode = "lightweight"
                items = await self._lightweight_retrieve(query, constraints, owner_id)
        else:
            raise ValueError(
                f"Invalid retrieval mode {mode!r}. "
                f"Supported modes: 'auto', 'lightweight', 'chained', 'agentic', 'brain_graph'."
            )

        # Merge WM items with DB results
        if wm_items:
            seen_ids = {item.item_id for item in wm_items}
            db_deduped = [item for item in items if item.item_id not in seen_ids]
            items = wm_items + db_deduped
            mode = f"wm+{mode}"

        # Rerank step (#502)
        _rerank_input_count = 0
        _rerank_output_count = 0
        _rerank_top_score: Optional[float] = None
        _rerank_latency_ms = 0
        if _rerank_enabled and items:
            _rerank_input_count = len(items)
            _rerank_t0 = time.monotonic()
            try:
                items = await asyncio.wait_for(
                    self._rerank_results(query, items), timeout=0.4,
                )
                _rerank_latency_ms = int((time.monotonic() - _rerank_t0) * 1000)
                _rerank_output_count = len(items)
                if items:
                    _rerank_top_score = items[0].score
            except asyncio.TimeoutError:
                _rerank_latency_ms = int((time.monotonic() - _rerank_t0) * 1000)
                _rerank_top_k = getattr(self.config, "retrieval_rerank_top_k", 5)
                logger.warning(
                    "retrieve: rerank timed out after 400ms, returning top %d unranked",
                    _rerank_top_k,
                )
                items = items[:_rerank_top_k]
                _rerank_output_count = len(items)

        # Status penalty (#574)
        items = self._apply_status_penalty(items)
        items = [it for it in items if it.score > 0]

        # Exclude claims from a specific source session (#1220 — echo loop prevention)
        _excl_session = getattr(constraints, "exclude_source_session", None)
        if _excl_session:
            items = [
                it for it in items
                if getattr(it, "source_session_id", None) != _excl_session
            ]

        _has_date_constraints = bool(constraints.before_date or constraints.after_date)
        if not _has_date_constraints:
            items.sort(key=lambda it: it.score, reverse=True)

        tokens_used, items = await self._apply_budget(items, constraints)
        latency_ms = int((time.monotonic() - t0) * 1000)

        # Structured retrieval trace (#1347)
        if getattr(self.config, "retrieval_trace_enabled", False):
            self._emit_retrieval_trace(
                retrieval_id=retrieval_id,
                owner_id=owner_id,
                query=query,
                mode=mode,
                items=items,
                rerank_input_count=_rerank_input_count,
                rerank_output_count=_rerank_output_count,
                rerank_top_score=_rerank_top_score,
                rerank_latency_ms=_rerank_latency_ms,
                total_latency_ms=latency_ms,
            )

        await self._log_retrieval(
            retrieval_id=retrieval_id, owner_id=owner_id, query=query,
            mode=mode, constraints=constraints, items=items,
            tokens_used=tokens_used, latency_ms=latency_ms,
        )

        if _sentry_ret_span is not None:
            try:
                _sentry_ret_span.set_data("result_count", len(items))
                _sentry_ret_span.set_data("latency_ms", latency_ms)
                _sentry_ret_span.finish()
            except Exception:  # pragma: no cover
                pass

        # Populate retrieval_info from lightweight retrieve (#P8-4)
        _retrieval_info = getattr(self, "_last_retrieval_info", None)
        if _retrieval_info is not None and _rerank_enabled and items:
            # rerank succeeded if we got here without exception
            _retrieval_info["reranked"] = True

        return RetrievalResult(
            items=items, retrieval_id=retrieval_id, mode=mode,
            latency_ms=latency_ms, tokens_used=tokens_used,
            retrieval_info=_retrieval_info,
        )

    async def assemble_context(
        self,
        query: str,
        owner_id: str = "default",
        cornerstone_guardian: Any = None,
        retrieval_id: Optional[str] = None,
    ) -> ContextBlock:
        """Assemble a formatted context block ready for system prompt injection."""
        if retrieval_id is None:
            retrieval_id = _new_id()

        guardian = cornerstone_guardian if cornerstone_guardian is not None else self.cornerstone_guardian
        total_budget = getattr(self.config, "retrieval_token_budget", 3000)

        parts: List[Tuple[str, str, int]] = []
        total_tokens = 0

        cornerstone_tokens = 0
        if guardian is not None:
            try:
                cs_text = await guardian.format_injection_block(owner_id=owner_id)
                if cs_text:
                    cornerstone_tokens = self._count_tokens(cs_text)
                    parts.append(("cornerstone", cs_text, cornerstone_tokens))
                    total_tokens += cornerstone_tokens
                    logger.debug(
                        "assemble_context: injected %d cornerstone tokens", cornerstone_tokens,
                    )
            except Exception as exc:
                logger.warning("assemble_context: cornerstone injection failed: %s", exc)

        available_for_retrieval = max(0, total_budget - cornerstone_tokens)

        if self.token_budget is not None:
            try:
                budget_obj = self.token_budget.create_budget(total=available_for_retrieval)
                available_for_retrieval = min(available_for_retrieval, budget_obj.remaining)
            except Exception as exc:
                logger.warning("Token budget creation failed, proceeding with defaults: %s", exc)

        result = await self.retrieve(
            query,
            constraints=RetrievalConstraints(token_budget=available_for_retrieval),
            owner_id=owner_id,
        )

        for item in result.items:
            item_tokens = self._count_tokens(item.content)
            if total_tokens + item_tokens > total_budget:
                break
            parts.append((item.item_type, item.content, item_tokens))
            total_tokens += item_tokens

        return ContextBlock(
            parts=parts, total_tokens=total_tokens,
            budget_remaining=total_budget - total_tokens,
            retrieval_id=retrieval_id,
        )

    async def explain(self, retrieval_id: str) -> Dict[str, Any]:
        """Return provenance and scoring details for a retrieval (stub for Phase 1)."""
        return {"retrieval_id": retrieval_id, "note": "Full explain available in Phase 2"}

    # ------------------------------------------------------------------
    # Delegation methods — maintain backward compat for tests and
    # consumers that call pipeline._{method} directly or mock
    # cortex.core.retrieval.RetrievalPipeline._{method}.
    # ------------------------------------------------------------------

    _DEFAULT_STATUS_PENALTIES = DEFAULT_STATUS_PENALTIES

    def _apply_status_penalty(self, items: List[RetrievedItem]) -> List[RetrievedItem]:
        return apply_status_penalty(self.config, items, self._DEFAULT_STATUS_PENALTIES)

    async def _working_memory_retrieve(
        self, query: str, constraints: RetrievalConstraints, owner_id: str,
        session_id: Optional[str] = None,
    ) -> List[RetrievedItem]:
        return await working_memory_retrieve(
            self.working_memory, query, constraints, owner_id,
            session_id=session_id,
        )

    def _route_query(self, query: str, constraints: RetrievalConstraints) -> str:
        words = query.lower().split()
        has_semantic_cues = any(
            w in words for w in ("explain", "why", "how", "what", "when", "who", "describe")
        )
        return "lightweight"

    async def _lightweight_retrieve(
        self, query: str, constraints: RetrievalConstraints, owner_id: str,
    ) -> List[RetrievedItem]:
        items, self._last_retrieval_info = await lightweight_retrieve(
            self.storage, self.config, query, constraints, owner_id,
            rrf_k=self.rrf_k, candidate_multiplier=self.candidate_multiplier,
            embed_fn=self._embed, token_budget_mgr=self.token_budget,
        )
        return items

    async def _embed(self, text: str) -> Optional[List[float]]:
        return await embed_text(text, self.storage, self.llm, self.embedding_cache)

    async def _apply_budget(
        self, items: List[RetrievedItem], constraints: RetrievalConstraints,
    ) -> Tuple[int, List[RetrievedItem]]:
        return await apply_budget(items, constraints, self.config, self.token_budget)

    async def _rerank_results(
        self, query: str, items: List[RetrievedItem],
    ) -> List[RetrievedItem]:
        """Rerank retrieved items using Voyage rerank-2.5-lite (#502).

        Note: tests patch ``cortex.core.retrieval.voyageai``, so we resolve
        the voyageai reference through that module to remain mockable.
        """
        import cortex.core.retrieval as _ret_mod
        _voyageai = getattr(_ret_mod, "voyageai", voyageai)

        rerank_model = getattr(self.config, "retrieval_rerank_model", "rerank-2.5-lite")
        rerank_top_k = getattr(self.config, "retrieval_rerank_top_k", 5)
        documents = [item.content for item in items]
        non_empty_mask = [bool(doc.strip()) for doc in documents]
        if not any(non_empty_mask):
            logger.debug("_rerank_results: all items have empty content, skipping rerank")
            return items[:rerank_top_k]
        t0 = time.monotonic()
        try:
            if _voyageai is None:
                logger.warning(
                    "_rerank_results: voyageai package not installed; skipping rerank. "
                    "Install with: pip install voyageai"
                )
                return items[:rerank_top_k]
            api_key = os.environ.get("VOYAGE_API_KEY")
            if not api_key:
                logger.warning("_rerank_results: VOYAGE_API_KEY not set, skipping rerank")
                return items[:rerank_top_k]
            client = _voyageai.AsyncClient(api_key=api_key)
            rerank_result = await client.rerank(
                query=query, documents=documents, model=rerank_model,
                top_k=min(rerank_top_k, len(documents)),
            )
            latency_ms = int((time.monotonic() - t0) * 1000)
            logger.debug(
                "_rerank_results: reranked %d → %d items with %s in %dms",
                len(items), len(rerank_result.results), rerank_model, latency_ms,
            )
            reranked: List[RetrievedItem] = []
            for rr in rerank_result.results[:rerank_top_k]:
                original_item = items[rr.index]
                reranked.append(RetrievedItem(
                    item_type=original_item.item_type,
                    item_id=original_item.item_id,
                    content=original_item.content,
                    score=float(rr.relevance_score),
                    provenance=f"rerank:{rerank_model}",
                    metadata=original_item.metadata,
                    source_session_id=original_item.source_session_id,
                    created_at=original_item.created_at,
                ))
            return reranked
        except Exception as exc:
            latency_ms = int((time.monotonic() - t0) * 1000)
            logger.warning(
                "_rerank_results: rerank failed after %dms (%s: %s); returning unranked results",
                latency_ms, type(exc).__name__, exc,
            )
            return items[:rerank_top_k]

    def _count_tokens(self, text: str) -> int:
        return count_tokens(text, self.token_budget)

    async def _bm25_search(self, query: str, top_k: int, owner_id: str,
                           include_episodes: bool = False,
                           entity_ids: Optional[List[str]] = None) -> List[_FTSResult]:
        return await bm25_search(
            self.storage, query, top_k=top_k, owner_id=owner_id,
            include_episodes=include_episodes, entity_ids=entity_ids,
        )

    async def _vector_search(
        self, query: str, top_k: int, owner_id: str,
    ) -> List[_VectorResult]:
        """Delegate vector search to the module-level helper (mockable seam)."""
        return await vector_search(
            self.storage, query, top_k=top_k, owner_id=owner_id,
            embed_fn=self._embed,
        )

    async def _cosine_fallback(
        self, query_embedding: List[float], top_k: int,
    ) -> List[_VectorResult]:
        """Delegate cosine fallback to the module-level helper (mockable seam)."""
        return await cosine_fallback(self.storage, query_embedding, top_k)

    async def _log_retrieval(
        self, *, retrieval_id: str, owner_id: str, query: str, mode: str,
        constraints: RetrievalConstraints, items: List[RetrievedItem],
        tokens_used: int, latency_ms: int,
    ) -> None:
        await log_retrieval(
            self.storage, self.config, retrieval_id=retrieval_id,
            owner_id=owner_id, query=query, mode=mode,
            constraints=constraints, items=items,
            tokens_used=tokens_used, latency_ms=latency_ms,
        )

    def _emit_retrieval_trace(
        self,
        *,
        retrieval_id: str,
        owner_id: str,
        query: str,
        mode: str,
        items: List[RetrievedItem],
        rerank_input_count: int,
        rerank_output_count: int,
        rerank_top_score: Optional[float],
        rerank_latency_ms: int,
        total_latency_ms: int,
    ) -> None:
        """Emit a structured JSON retrieval trace (#1347).

        Called only when ``retrieval_trace_enabled`` is True. Constructs a
        RetrievalTrace, serialises to JSON, and logs to ``cortex.retrieval.trace``.
        Errors are silently swallowed — tracing must never affect retrieval.
        """
        import json as _json
        from datetime import datetime, timezone

        try:
            # Build category distribution from final items
            category_dist: Dict[str, int] = {}
            for it in items:
                cat = "unknown"
                if it.metadata:
                    cat = it.metadata.get("category", "unknown") or "unknown"
                category_dist[cat] = category_dist.get(cat, 0) + 1

            trace = RetrievalTrace(
                retrieval_id=retrieval_id,
                owner_id=owner_id,
                query=query,
                mode=mode,
                timestamp_utc=datetime.now(timezone.utc).isoformat(),
                rerank_input_count=rerank_input_count,
                rerank_output_count=rerank_output_count,
                rerank_top_score=rerank_top_score,
                rerank_latency_ms=rerank_latency_ms,
                final_count=len(items),
                final_top_score=items[0].score if items else None,
                final_min_score=items[-1].score if items else None,
                final_item_ids=[it.item_id for it in items],
                category_distribution=category_dist,
                total_latency_ms=total_latency_ms,
            )

            log_level_name = getattr(self.config, "retrieval_trace_log_level", "DEBUG").upper()
            log_level = getattr(logging, log_level_name, logging.DEBUG)
            trace_logger.log(log_level, _json.dumps(trace.to_dict()))

        except Exception as _exc:  # pragma: no cover — tracing must never break retrieval
            logger.debug("retrieval trace emit failed (non-fatal): %s", _exc)
