"""
cortex/core/retrieval_search.py — Search backends and processing for the retrieval pipeline.

Extracted from cortex/core/retrieval.py (#947 — SRP split).
Contains BM25/FTS5 search, vector search, cosine fallback, embedding,
working memory retrieval, lightweight retrieve, rerank, status penalty,
budget enforcement, and retrieval logging.
"""

from __future__ import annotations

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

if TYPE_CHECKING:
    from cortex.storage.adapter import StorageAdapter

# Voyage AI SDK — optional, required only when retrieval_rerank_enabled=True.
try:
    import voyageai  # type: ignore
except ImportError:
    voyageai = None  # type: ignore

from cortex.core.retrieval_types import (
    RetrievalConstraints,
    RetrievedItem,
    _FTSResult,
    _VectorResult,
)
from cortex.core.retrieval_scoring import (
    _normalise_bm25,
    _sensitivity_level,
    cosine_similarity,
    rrf_fusion,
)
from cortex.core.retrieval_helpers import (
    _decode_embedding,
    _get_content,
    _get_id,
    _get_item_type,
    _get_rank,
    _get_score,
    _render_claim_content,
    _table_to_item_type,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# BM25 / FTS search
# ---------------------------------------------------------------------------

async def bm25_search(
    storage: "StorageAdapter",
    query: str,
    top_k: int,
    owner_id: str,
    include_episodes: bool = False,
    entity_ids: Optional[List[str]] = None,
) -> List[_FTSResult]:
    """Search fts_claims (and optionally fts_episodes), merge, return top_k."""
    tables: List[str] = ["claims"]
    if include_episodes:
        tables.append("episodes")
    results: List[_FTSResult] = []
    for table in tables:
        try:
            _fts_kwargs: Dict[str, Any] = dict(
                table=table, top_k=top_k, owner_id=owner_id,
            )
            if entity_ids and table == "claims":
                _fts_kwargs["entity_ids"] = entity_ids
            raw = await storage.fts_search(query, **_fts_kwargs)
            for r in raw:
                results.append(
                    _FTSResult(
                        item_type=_table_to_item_type(table),
                        item_id=_get_id(r),
                        content=_get_content(r),
                        rank=_get_rank(r),
                    )
                )
        except Exception as exc:
            logger.warning("BM25 search failed for table=%s: %s", table, exc)

    results.sort(key=lambda r: r.rank)
    return results[:top_k]


# ---------------------------------------------------------------------------
# Vector search
# ---------------------------------------------------------------------------

async def vector_search(
    storage: "StorageAdapter",
    query: str,
    top_k: int,
    owner_id: str,
    embed_fn: Any,
) -> List[_VectorResult]:
    """Embed the query, then ask storage for nearest neighbours."""
    embedding = await embed_fn(query)
    if embedding is None:
        logger.info("No embedding available; skipping vector search")
        return []

    try:
        raw = await storage.vector_search(
            embedding, top_k=top_k, owner_id=owner_id
        )
        return [
            _VectorResult(
                item_type=_get_item_type(r),
                item_id=_get_id(r),
                content=_get_content(r),
                score=_get_score(r),
            )
            for r in raw
        ]
    except NotImplementedError:
        logger.warning(
            "sqlite-vec not available; falling back to brute-force cosine similarity. "
            "Install sqlite-vec for faster ANN vector search."
        )
        return await cosine_fallback(storage, embedding, top_k, owner_id=owner_id)
    except Exception as exc:
        logger.warning("Vector search failed: %s", exc)
        return []


async def cosine_fallback(
    storage: "StorageAdapter",
    query_embedding: List[float],
    top_k: int,
    owner_id: str = "default",
) -> List[_VectorResult]:
    """Brute-force cosine similarity fallback scoped by owner_id.

    Delegates to storage.vector_search() which already applies owner_id
    filtering at the SQL level (JOIN to semantic_claim WHERE owner_id = ?).

    Previously this called a non-existent ``get_all_embeddings()`` and
    only post-filtered claims, letting non-claim items pass through
    unscoped (#1528).  Now all items are owner-scoped via the storage
    layer's brute-force implementation.
    """
    try:
        raw = await storage.vector_search(
            query_embedding, item_type=None, top_k=top_k, owner_id=owner_id
        )
    except Exception as exc:
        logger.warning("Cosine fallback via storage.vector_search failed: %s", exc)
        return []

    results = [
        _VectorResult(
            item_type=_get_item_type(r),
            item_id=_get_id(r),
            content=_get_content(r),
            score=_get_score(r),
        )
        for r in raw
    ]

    # Update access tracking for retrieved claims (fire-and-forget, non-blocking)
    _claim_ids_to_track = [r.item_id for r in results if r.item_type == "claim"]
    if _claim_ids_to_track:
        from cortex.core.retrieval_lightweight import _update_access_tracking
        asyncio.create_task(_update_access_tracking(storage, _claim_ids_to_track, owner_id))

    return results


# ---------------------------------------------------------------------------
# Embedding
# ---------------------------------------------------------------------------

async def embed_text(
    text: str,
    storage: "StorageAdapter",
    llm: Any,
    embedding_cache: Any,
) -> Optional[List[float]]:
    """Produce an embedding for *text*, using the cache when available."""
    if embedding_cache is not None:
        cached = embedding_cache.get(text)
        if cached is not None:
            return cached

    if hasattr(storage, "embed"):
        try:
            result = await storage.embed(text)  # type: ignore[attr-defined]
            if result is not None and embedding_cache is not None:
                embedding_cache.set(text, result)
            return result
        except Exception as exc:
            logger.debug("storage.embed failed: %s", exc)

    if llm is not None and hasattr(llm, "embed"):
        try:
            _embed_resp = await llm.embed(texts=[text])
            result = (
                _embed_resp.embeddings[0]
                if hasattr(_embed_resp, "embeddings")
                else _embed_resp
            )
            if result is not None and embedding_cache is not None:
                embedding_cache.set(text, result)
            return result
        except Exception as exc:
            logger.debug("llm.embed failed: %s", exc)

    return None


# ---------------------------------------------------------------------------
# Working memory
# ---------------------------------------------------------------------------

async def working_memory_retrieve(
    working_memory: Any,
    query: str,
    constraints: RetrievalConstraints,
    owner_id: str,
    session_id: Optional[str] = None,
) -> List[RetrievedItem]:
    """Check working memory cache before any DB queries."""
    if working_memory is None or not working_memory.enabled:
        return []

    if session_id is None:
        logger.warning(
            "working_memory_retrieve: no session_id provided, falling back to owner_id=%s — "
            "WM entries written under real session IDs may not be found",
            owner_id,
        )
        session_id = owner_id

    try:
        matches = await working_memory.query(query, session_id=session_id)
    except Exception as exc:
        logger.warning("Working memory query failed (non-fatal): %s", exc)
        return []

    if not matches:
        return []

    items: List[RetrievedItem] = []
    for entry in matches[: (constraints.top_k or 20)]:
        content = str(entry.value)
        items.append(
            RetrievedItem(
                item_type="working_memory",
                item_id=f"wm:{entry.session_id}:{entry.key}",
                content=content,
                score=1.0,
                provenance=f"working_memory:{entry.category}",
            )
        )

    return items


# ---------------------------------------------------------------------------
# Token counting
# ---------------------------------------------------------------------------

def count_tokens(text: str, token_budget_mgr: Any) -> int:
    """Count tokens — delegate to TokenBudgetManager if available, else word-approx."""
    if token_budget_mgr is not None and hasattr(token_budget_mgr, "count_tokens"):
        try:
            return token_budget_mgr.count_tokens(text)
        except Exception as e:
            logger.warning("cortex_core_retrieval: unhandled exception: %s", e)
    return max(1, int(len(text.split()) * 1.3))


# ---------------------------------------------------------------------------
# Lightweight retrieval (BM25 + vector + RRF)
# ---------------------------------------------------------------------------

# ---------------------------------------------------------------------------
# Rerank
# ---------------------------------------------------------------------------

async def rerank_results(
    config: Any,
    query: str,
    items: List[RetrievedItem],
) -> List[RetrievedItem]:
    """Rerank retrieved items using Voyage rerank-2.5-lite (#502)."""
    rerank_model = getattr(config, "retrieval_rerank_model", "rerank-2.5-lite")
    rerank_top_k = getattr(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]


# ---------------------------------------------------------------------------
# Status penalty
# ---------------------------------------------------------------------------

DEFAULT_STATUS_PENALTIES: Dict[str, float] = {
    "active": 1.0,
    "contradicted": 0.3,
    "superseded": 0.0,
    "deprecated": 0.3,
    "cooling": 0.7,
    "needs_review": 0.0,
    "rejected": 0.0,
}


def apply_status_penalty(
    config: Any,
    items: List[RetrievedItem],
    default_penalties: Dict[str, float],
) -> List[RetrievedItem]:
    """Apply score penalties to claims with non-active status."""
    from collections.abc import Mapping
    _QUARANTINE_STATUSES = {"needs_review", "rejected", "superseded", "archived"}
    penalties = dict(default_penalties)
    custom = getattr(config, "retrieval_status_penalties", None)
    if custom is not None:
        if isinstance(custom, Mapping):
            penalties.update(custom)
        else:
            logger.warning(
                "retrieval_status_penalties is not a Mapping; ignoring custom value"
            )
    for qs in _QUARANTINE_STATUSES:
        penalties[qs] = 0.0

    from collections.abc import Mapping
    if not isinstance(penalties, Mapping):
        logger.warning(
            "retrieval_status_penalties is not a Mapping; using default penalties"
        )
        penalties = default_penalties

    filtered: List[RetrievedItem] = []
    for item in items:
        if item.item_type == "claim" and item.metadata:
            # Exclude deleted claims (#1430)
            if item.metadata.get("is_deleted"):
                continue
            status = item.metadata.get("status", "active")
            if status in _QUARANTINE_STATUSES:
                continue
            penalty = penalties.get(status, 0.5)
            if penalty < 1.0:
                item.score *= penalty
        filtered.append(item)

    return filtered


# ---------------------------------------------------------------------------
# Token budget enforcement
# ---------------------------------------------------------------------------

async def apply_budget(
    items: List[RetrievedItem],
    constraints: RetrievalConstraints,
    config: Any,
    token_budget_mgr: Any = None,
) -> Tuple[int, List[RetrievedItem]]:
    """Trim items list so total token count stays within budget."""
    if constraints.token_budget is not None:
        budget = constraints.token_budget
    else:
        budget = getattr(config, "retrieval_token_budget", 3000)

    trimmed: List[RetrievedItem] = []
    tokens_used = 0

    for item in items:
        item_tokens = count_tokens(item.content, token_budget_mgr)
        if tokens_used + item_tokens > budget:
            logger.debug(
                "Budget exhausted at %d/%d tokens — dropping remaining %d items",
                tokens_used, budget, len(items) - len(trimmed),
            )
            break
        trimmed.append(item)
        tokens_used += item_tokens

    return tokens_used, trimmed


# ---------------------------------------------------------------------------
# Retrieval logging
# ---------------------------------------------------------------------------

async def log_retrieval(
    storage: "StorageAdapter",
    config: Any,
    *,
    retrieval_id: str,
    owner_id: str,
    query: str,
    mode: str,
    constraints: RetrievalConstraints,
    items: List[RetrievedItem],
    tokens_used: int,
    latency_ms: int,
) -> None:
    """Log retrieval metadata to retrieval_log table."""
    import hashlib
    import json

    selected = [
        {"item_id": it.item_id, "item_type": it.item_type, "score": it.score}
        for it in items
    ]
    context_hash = hashlib.sha256(
        json.dumps(selected, sort_keys=True).encode()
    ).hexdigest()[:16]

    try:
        await storage.log_retrieval(
            id=retrieval_id,
            owner_id=owner_id,
            query_text=query,
            mode=mode,
            token_budget=constraints.token_budget or getattr(
                config, "retrieval_token_budget", 3000
            ),
            tokens_used=tokens_used,
            selected_json=json.dumps(selected),
            final_context_hash=context_hash,
            latency_ms=latency_ms,
            created_at=datetime.now(timezone.utc).isoformat(),
        )
    except Exception as exc:
        logger.warning("Failed to log retrieval %s: %s", retrieval_id, exc)
