"""
cortex/core/retrieval_lightweight.py — Lightweight hybrid retrieval (BM25 + vector + RRF).

Extracted from cortex/core/retrieval.py (#947 — SRP split).
Contains the core retrieval workhorse: parallel BM25 + vector search,
RRF fusion, sensitivity filtering, temporal scoring, and item assembly.
"""

from __future__ import annotations

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

if TYPE_CHECKING:
    from cortex.storage.adapter import StorageAdapter

from cortex.core.retrieval_types import (
    RetrievalConstraints,
    RetrievedItem,
)
from cortex.core.retrieval_scoring import (
    _normalise_bm25,
    _sensitivity_level,
    rrf_fusion,
)
from cortex.core.retrieval_helpers import (
    _get_content,
    _render_claim_content,
)
from cortex.core.retrieval_search import (
    bm25_search,
    vector_search,
)

logger = logging.getLogger(__name__)


async def lightweight_retrieve(
    storage: "StorageAdapter",
    config: Any,
    query: str,
    constraints: RetrievalConstraints,
    owner_id: str,
    rrf_k: int,
    candidate_multiplier: int,
    embed_fn: Any,
    token_budget_mgr: Any = None,
) -> Tuple[List[RetrievedItem], Dict[str, Any]]:
    """Parallel BM25 + vector search, fused with RRF.

    This is the core retrieval workhorse extracted from
    RetrievalPipeline._lightweight_retrieve for SRP compliance.
    """
    top_k = constraints.top_k
    candidate_k = top_k * candidate_multiplier

    # Run BM25 and vector search in parallel with timeout (#P8-3)
    bm25_task = asyncio.create_task(
        bm25_search(storage, query, top_k=candidate_k, owner_id=owner_id,
                    entity_ids=constraints.entity_ids)
    )
    vector_task = asyncio.create_task(
        vector_search(storage, query, top_k=candidate_k, owner_id=owner_id,
                      embed_fn=embed_fn)
    )

    retrieval_info: Dict[str, Any] = {
        "bm25_ok": True,
        "vector_ok": True,
        "reranked": False,
        "timed_out": False,
    }

    try:
        bm25_results, vector_results = await asyncio.wait_for(
            asyncio.gather(bm25_task, vector_task, return_exceptions=True),
            timeout=5.0,
        )
    except asyncio.TimeoutError:
        logger.error("Retrieval search timed out after 5.0s")
        retrieval_info["timed_out"] = True
        retrieval_info["bm25_ok"] = False
        retrieval_info["vector_ok"] = False
        bm25_results, vector_results = [], []

    if isinstance(bm25_results, BaseException):
        logger.warning("BM25 search failed: %s", bm25_results)
        retrieval_info["bm25_ok"] = False
        bm25_results = []
    if isinstance(vector_results, BaseException):
        logger.warning("Vector search failed: %s", vector_results)
        retrieval_info["vector_ok"] = False
        vector_results = []

    bm25_ranked = _normalise_bm25(bm25_results)
    vector_ranked = [(r.item_id, r.score) for r in vector_results]

    fused = rrf_fusion([bm25_ranked, vector_ranked], k=rrf_k)

    content_map: Dict[str, Tuple[str, str]] = {}
    for r in bm25_results:
        content_map[r.item_id] = (r.item_type, r.content)
    for r in vector_results:
        new_content = _get_content(r)
        if r.item_id not in content_map or new_content:
            content_map[r.item_id] = (r.item_type, new_content)

    # Hydrate empty content + fetch metadata for claim items (batch)
    claim_cache: Dict[str, Any] = {}
    _hydrate_ids = [
        iid for iid, (itype, content) in content_map.items()
        if not content and itype == "claim"
    ]
    _meta_ids = [
        iid for iid, (itype, _) in content_map.items()
        if itype == "claim" and iid not in _hydrate_ids
    ]
    _all_claim_ids = _hydrate_ids + _meta_ids

    _hydrate_failures = 0
    if _all_claim_ids:
        # Check class for method (avoids MagicMock auto-attributes in tests),
        # but call on instance (bound method) to fix unbound-method bug (#1578).
        _has_batch = hasattr(type(storage), "get_claims_batch")
        if _has_batch:
            try:
                _batch = await storage.get_claims_batch(_all_claim_ids)
                claim_cache.update(_batch)
            except Exception as exc:
                logger.error("Batch claim fetch failed (%d ids): %s", len(_all_claim_ids), exc)
                _hydrate_failures = len(_all_claim_ids)
        else:
            # Fallback: per-ID fetch when batch method unavailable
            for _cid in _all_claim_ids:
                try:
                    _c = await storage.get_claim(_cid)
                    if _c:
                        claim_cache[_cid] = _c
                except Exception as exc:
                    logger.warning("Failed to fetch claim %s: %s", _cid[:8], exc)
                    _hydrate_failures += 1

    # Apply hydrated content
    for item_id in _hydrate_ids:
        claim = claim_cache.get(item_id)
        if claim:
            content_map[item_id] = (
                content_map[item_id][0],
                _render_claim_content(claim),
            )
        else:
            _hydrate_failures += 1

    _hydrated = sum(1 for iid in _hydrate_ids if iid in claim_cache)
    logger.info(
        "Hydrated %d/%d claims, %d failures",
        _hydrated, len(_hydrate_ids), _hydrate_failures,
    )

    # --- Sensitivity filtering (issue #272) ---
    claim_ids_in_map = [
        iid for iid, (itype, _) in content_map.items() if itype == "claim"
    ]
    sensitivity_map: Dict[str, str] = {}
    if claim_ids_in_map:
        try:
            sensitivity_map = await storage.get_claims_sensitivity_batch(claim_ids_in_map)
        except Exception as exc:
            logger.warning(
                "Sensitivity batch failed (%s); excluding all claims (fail-closed).", exc,
            )

    max_level = _sensitivity_level(constraints.sensitivity_max)

    def _passes_sensitivity(item_id: str, item_type: str) -> bool:
        if item_type != "claim":
            return True
        sens = sensitivity_map.get(item_id)
        return sens is not None and _sensitivity_level(sens) <= max_level

    # ------------------------------------------------------------------
    # Scoring weight config (#847 — expose all hardcoded weights)
    # ------------------------------------------------------------------
    try:
        _recency_weight = float(getattr(config, "retrieval_recency_weight", 0.1))
    except (TypeError, ValueError):
        _recency_weight = 0.1
    try:
        _recency_half_life_days = float(getattr(config, "retrieval_recency_half_life_days", 30.0))
    except (TypeError, ValueError):
        _recency_half_life_days = 30.0
    try:
        _explicitness_bonus = float(getattr(config, "retrieval_explicitness_bonus", 0.05))
    except (TypeError, ValueError):
        _explicitness_bonus = 0.05
    try:
        _stability_bonus = float(getattr(config, "retrieval_stability_bonus", 0.03))
    except (TypeError, ValueError):
        _stability_bonus = 0.03

    # ------------------------------------------------------------------
    # Temporal scoring (#513)
    # ------------------------------------------------------------------
    _temporal_boost_weight = getattr(config, "retrieval_temporal_boost_weight", 0.15)
    _has_date_filters = bool(constraints.before_date or constraints.after_date)

    _reference_date: Optional[datetime] = None
    if _temporal_boost_weight > 0:
        try:
            from cortex.core.temporal_scoring import parse_temporal_hints
            _reference_date = parse_temporal_hints(query)
        except Exception as _te:
            logger.debug("temporal hint parse failed (non-fatal): %s", _te)

    _date_filter_fn = None
    if _has_date_filters:
        try:
            from cortex.core.temporal_scoring import filter_by_date
            _before = constraints.before_date
            _after = constraints.after_date
            _date_filter_fn = lambda ed: filter_by_date(ed, _before, _after)  # noqa: E731
        except Exception as _te:
            logger.debug("temporal date filter import failed (non-fatal): %s", _te)

    _proximity_boost_fn = None
    if _reference_date is not None and _temporal_boost_weight > 0:
        try:
            from cortex.core.temporal_scoring import temporal_proximity_boost
            _ref = _reference_date
            _weight = _temporal_boost_weight
            _proximity_boost_fn = lambda ed: _weight * temporal_proximity_boost(ed, _ref)  # noqa: E731
        except Exception as _te:
            logger.debug("temporal proximity import failed (non-fatal): %s", _te)

    # Build result items with filtering
    # Iterate over a wider window than top_k so post-filters don't starve
    # the result set.  Cap at top_k * 3 to avoid scanning the full list.
    _candidate_ceiling = min(len(fused), top_k * 3)
    items: List[RetrievedItem] = []
    for item_id, score in fused[:_candidate_ceiling]:
        if item_id not in content_map:
            continue
        item_type, content = content_map[item_id]
        if constraints.memory_types and item_type not in constraints.memory_types:
            continue
        if constraints.exclude_episodes and item_type == "event":
            continue
        if not _passes_sensitivity(item_id, item_type):
            continue
        # Exclude temporally expired claims unless include_expired=True (#925)
        if item_type == "claim" and not constraints.include_expired:
            _cached_claim = claim_cache.get(item_id)
            if _cached_claim is not None:
                _valid_to = getattr(_cached_claim, "valid_to", None)
                if isinstance(_valid_to, str) and _valid_to:
                    try:
                        _vt = _valid_to.replace("Z", "")[:19]
                        _now_str = datetime.now(timezone.utc).isoformat()[:19]
                        if _vt <= _now_str:
                            continue
                    except (ValueError, TypeError):
                        pass

        # Attach LoCoMo classification metadata for claim items
        item_metadata: Optional[Dict[str, Any]] = None
        _event_date: Optional[str] = None
        if item_type == "claim":
            cached = claim_cache.get(item_id)  # B-CF4: single lookup, reused below
            if cached is not None:
                # B-CF5: Normalize is_deleted — SQLite may store as string "1"/"true"
                _raw_is_deleted = getattr(cached, "is_deleted", 0)
                _is_deleted = 1 if _raw_is_deleted in (1, True, "1", "true") else 0
                # B-CF1: mentions defaults to 0 (int), not None
                _raw_mentions = getattr(cached, "mentions", None)
                _mentions = _raw_mentions if _raw_mentions is not None else 0
                # B-CF2: entity_type defaults to "" when NULL (not "agent" — semantic assumption)
                _raw_entity_type = getattr(cached, "entity_type", None)
                _entity_type = _raw_entity_type if _raw_entity_type is not None else ""
                # B-CF3: scope defaults to "" when NULL (not "user" — semantic assumption)
                _raw_scope = getattr(cached, "scope", None)
                _scope = _raw_scope if _raw_scope is not None else ""
                item_metadata = {
                    "memory_class": getattr(cached, "memory_class", "semantic"),
                    "salience": getattr(cached, "salience", "medium"),
                    "status": getattr(cached, "status", "active"),
                    "explicitness": getattr(cached, "explicitness", "inferred"),
                    "stability": getattr(cached, "stability", "volatile"),
                    "category": getattr(cached, "category", "misc"),
                    "is_deleted": _is_deleted,
                    # Priority fields surfaced for API callers (closes #1344)
                    "temporal": getattr(cached, "temporal", None),
                    "temporal_marker": getattr(cached, "temporal_marker", None),
                    "event_date": getattr(cached, "event_date", None),
                    "cause": getattr(cached, "cause", None),
                    "effect": getattr(cached, "effect", None),
                    "state_before": getattr(cached, "state_before", None),
                    "state_after": getattr(cached, "state_after", None),
                    "entity_id": getattr(cached, "entity_id", None),
                    "entity_type": _entity_type,
                    "scope": _scope,
                    "memory_type": getattr(cached, "memory_type", None),
                    "mentions": _mentions,
                    "last_mentioned_at": getattr(cached, "last_mentioned_at", None),
                    "last_accessed": getattr(cached, "last_accessed", None),
                    "superseded_by": getattr(cached, "superseded_by", None),
                }
                _event_date = getattr(cached, "event_date", None)

        if _date_filter_fn is not None and not _date_filter_fn(_event_date):
            continue

        final_score = score
        if _proximity_boost_fn is not None and _event_date:
            final_score = score + _proximity_boost_fn(_event_date)

        if item_metadata and item_type == "claim":
            if item_metadata.get("explicitness") == "explicit":
                final_score += _explicitness_bonus
            if item_metadata.get("stability") == "stable":
                final_score += _stability_bonus

        # Extract source_session_id + created_at from cached claim (#954)
        # B-CF4: reuse `cached` from metadata block above (avoid redundant dict lookup)
        _item_source_session_id: Optional[str] = None
        _item_created_at: Optional[str] = None
        if item_type == "claim":
            _cached = cached  # already fetched above via claim_cache.get(item_id)
            if _cached is not None:
                _item_created_at = getattr(_cached, "created_at", None)
                _meta_json_raw = getattr(_cached, "metadata_json", None)
                if _meta_json_raw:
                    try:
                        import json as _json_prov
                        _meta_parsed = _json_prov.loads(_meta_json_raw) if isinstance(_meta_json_raw, str) else _meta_json_raw
                        _item_source_session_id = _meta_parsed.get("source_session_id")
                    except (ValueError, TypeError, AttributeError):
                        pass
                if not _item_source_session_id:
                    _item_source_session_id = getattr(_cached, "source_session", None)

        # Recency bonus (#847) — exponential decay on created_at age
        if _recency_weight > 0 and _item_created_at:
            try:
                import math as _math
                _created_dt = datetime.fromisoformat(
                    _item_created_at.replace("Z", "+00:00")
                )
                _now = datetime.now(timezone.utc)
                _age_days = max(0.0, (_now - _created_dt).total_seconds() / 86400.0)
                _recency_bonus = _recency_weight * _math.exp(
                    -_age_days * _math.log(2) / max(_recency_half_life_days, 1.0)
                )
                final_score += _recency_bonus
            except (ValueError, TypeError, OverflowError) as _re:
                logger.debug("recency bonus skipped for %s: %s", item_id[:8], _re)

        items.append(
            RetrievedItem(
                item_type=item_type,
                item_id=item_id,
                content=content,
                score=final_score,
                metadata=item_metadata,
                source_session_id=_item_source_session_id,
                created_at=_item_created_at,
            )
        )
        if len(items) >= top_k:
            break

    # Filter ghost entries: items whose DB row was deleted but whose
    # vec_index / embedding_index entry lingered (#1395).
    items = [it for it in items if it.content]

    # Post-retrieval owner_id safety filter (#1398).
    # Even though BM25 and vector search already scope by owner_id,
    # this acts as a defence-in-depth layer — any item whose cached
    # claim has a mismatched owner_id is removed.
    if owner_id:
        _pre_count = len(items)
        _filtered: List[RetrievedItem] = []
        for it in items:
            if it.item_type == "claim":
                cached = claim_cache.get(it.item_id)
                if cached is not None:
                    _claim_owner = getattr(cached, "owner_id", None)
                    if isinstance(_claim_owner, str) and _claim_owner != owner_id:
                        logger.warning(
                            "Post-retrieval owner filter: dropping %s (owner=%s, expected=%s)",
                            it.item_id[:8], _claim_owner, owner_id,
                        )
                        continue
            _filtered.append(it)
        items = _filtered
        if len(items) < _pre_count:
            logger.info(
                "Post-retrieval owner filter removed %d items",
                _pre_count - len(items),
            )

    # Chronological sort when date filters active
    if _has_date_filters and items:
        def _sort_key(item: RetrievedItem) -> str:
            if item.metadata and item.item_type == "claim":
                _cached_item = claim_cache.get(item.item_id)
                ed = getattr(_cached_item, "event_date", None) if _cached_item else None
                if ed:
                    return str(ed)
            return "9999-99-99"
        items.sort(key=_sort_key)

    # Update access tracking for retrieved claims (fire-and-forget, non-blocking)
    _claim_ids_to_track = [it.item_id for it in items if it.item_type == "claim"]
    if _claim_ids_to_track:
        asyncio.create_task(_update_access_tracking(storage, _claim_ids_to_track, owner_id))

    # Hydration observability (#1583)
    retrieval_info["hydration_attempted"] = len(_all_claim_ids)
    retrieval_info["hydration_failed"] = _hydrate_failures
    retrieval_info["hydration_success_rate"] = 1.0 - (
        _hydrate_failures / max(len(_all_claim_ids), 1)
    )

    # Update health endpoint stats
    try:
        from cortex.api.health import update_retrieval_hydration
        update_retrieval_hydration(len(_all_claim_ids), _hydrate_failures)
    except ImportError:
        pass  # API layer not loaded (e.g., CLI-only usage)

    logger.info(
        "Retrieval completed: bm25=%s vector=%s reranked=%s timed_out=%s hydration=%d/%d (%.0f%%)",
        retrieval_info["bm25_ok"],
        retrieval_info["vector_ok"],
        retrieval_info["reranked"],
        retrieval_info["timed_out"],
        len(_all_claim_ids) - _hydrate_failures,
        len(_all_claim_ids),
        retrieval_info["hydration_success_rate"] * 100,
    )

    return items, retrieval_info


async def _update_access_tracking(
    storage: "StorageAdapter",
    claim_ids: List[str],
    owner_id: str,
) -> None:
    """Batch-update access_count and last_accessed for retrieved claims."""
    from datetime import datetime, timezone
    now = datetime.now(timezone.utc).isoformat()
    for claim_id in claim_ids:
        try:
            claim = await storage.get_claim(claim_id)
            if claim:
                current_count = getattr(claim, 'access_count', 0) or 0
                await storage.update_claim(
                    claim_id,
                    owner_id=owner_id,
                    access_count=current_count + 1,
                    last_accessed=now,
                )
        except Exception as exc:
            import logging
            logging.getLogger(__name__).debug(
                "access tracking update failed for %s: %s", claim_id[:8], exc
            )
