"""cortex/api/routes/memory/retrieve.py — Retrieve and retrieve/explain endpoints."""

from __future__ import annotations

import logging
from typing import Any, Dict, List

from fastapi import Depends, HTTPException, Query, Request, Response, status
from fastapi.responses import JSONResponse

from cortex.api.deps import PREFIX, _resolve_owner_id, _serialise, get_plugin
from cortex.api.deprecation import add_deprecation_headers
from cortex.api.models import (
    EnhancedRetrieveRequest,
    RetrieveExplainRequest,
    RetrieveExplainResponse,
)

from cortex.api.routes.memory.helpers import (
    _MODE_ALIASES,
    _emit_shadow_event,
)

logger = logging.getLogger(__name__)


def _register(router, verify_api_key):
    """Register retrieve and retrieve/explain endpoints."""

    # ---- Canonical: /memories/retrieve/explain  (issue #297) ----
    @router.post(
        f"{PREFIX}/memories/retrieve/explain",
        response_model=RetrieveExplainResponse,
        status_code=status.HTTP_200_OK,
        summary="Retrieval with explanation",
        tags=["memory"],
        operation_id="memories_retrieve_explain",
    )
    async def memories_retrieve_explain(
        request: Request,
        body: RetrieveExplainRequest,
        _: None = Depends(verify_api_key),
        plugin: Any = Depends(get_plugin),
    ) -> RetrieveExplainResponse:
        """Retrieve memories for a query and return explanation of retrieval process.

        Canonical path: ``POST /api/v1/memories/retrieve/explain``.
        """
        owner_id = _resolve_owner_id(request)
        return await _do_retrieve_explain(body, plugin, owner_id=owner_id)

    # ---- DEPRECATED: /retrieve/explain  (issue #297) ----
    @router.post(
        f"{PREFIX}/retrieve/explain",
        response_model=RetrieveExplainResponse,
        status_code=status.HTTP_200_OK,
        summary="Retrieval with explanation — DEPRECATED, use /memories/retrieve/explain",
        tags=["memory (deprecated)"],
        operation_id="retrieve_explain",
        deprecated=True,
    )
    async def retrieve_explain(
        request: Request,
        body: RetrieveExplainRequest,
        response: Response,
        _: None = Depends(verify_api_key),
        plugin: Any = Depends(get_plugin),
    ) -> RetrieveExplainResponse:
        """DEPRECATED — use ``POST /api/v1/memories/retrieve/explain``.

        .. deprecated:: sunset 2026-06-01  (issue #297)
        """
        add_deprecation_headers(response, f"{PREFIX}/memories/retrieve/explain")
        owner_id = _resolve_owner_id(request)
        return await _do_retrieve_explain(body, plugin, owner_id=owner_id)

    async def _do_retrieve_explain(
        body: RetrieveExplainRequest, plugin: Any, *, owner_id: str = "default",
    ) -> RetrieveExplainResponse:
        """Retrieve memories for a query and return explanation of retrieval process."""
        errors: List[str] = []
        results: List[Dict[str, Any]] = []
        explanation: Dict[str, Any] = {}

        storage = getattr(plugin, "storage", None)

        try:
            from cortex.core.retrieval import RetrievalConstraints, RetrievalPipeline

            pipeline = RetrievalPipeline(
                storage=storage,
                config=plugin.config,
                token_budget=getattr(plugin, "token_budget_manager", None),
                cornerstone_guardian=getattr(plugin, "cornerstone_guardian", None),
                llm=getattr(plugin, "llm", None),
            )

            constraints = RetrievalConstraints(top_k=body.limit)
            # owner_id resolved from auth middleware via _resolve_owner_id (#1398)

            import asyncio
            import time

            t0 = time.monotonic()

            candidate_k = body.limit * pipeline.candidate_multiplier

            bm25_task = asyncio.create_task(
                pipeline._bm25_search(body.query, top_k=candidate_k, owner_id=owner_id)
            )
            vector_task = asyncio.create_task(
                pipeline._vector_search(body.query, top_k=candidate_k, owner_id=owner_id)
            )

            from cortex.core.retrieval import _normalise_bm25, rrf_fusion
            bm25_results, vector_results = await asyncio.gather(bm25_task, vector_task)

            fts_matches = len(bm25_results)
            vector_matches = len(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=pipeline.rrf_k)
            rrf_scores = [
                {"item_id": item_id, "score": round(score, 6)}
                for item_id, score in fused[:body.limit]
            ]

            cornerstone_injected = False
            cornerstone_tokens = 0
            guardian = getattr(pipeline, "cornerstone_guardian", None)
            if guardian is not None:
                try:
                    cs_list = await guardian.get_cornerstones(owner_id=owner_id)
                    cornerstone_injected = len(cs_list) > 0
                    for cs in cs_list:
                        cs_text = getattr(cs, "content", "") or ""
                        cornerstone_tokens += pipeline._count_tokens(cs_text)
                except Exception as e:
                    logger.warning("http_server: unhandled exception: %s", e)

            total_budget = getattr(plugin.config, "retrieval_token_budget", 3000)
            remaining = total_budget

            full_result = await pipeline.retrieve(
                query=body.query,
                mode="auto",
                constraints=constraints,
                owner_id=owner_id,
            )

            strategy = getattr(full_result, "mode", "lightweight")
            tokens_used = getattr(full_result, "tokens_used", 0)
            remaining = max(0, total_budget - tokens_used - cornerstone_tokens)

            results = [
                _serialise(item.__dict__) if hasattr(item, "__dict__") else _serialise(item)
                for item in (getattr(full_result, "items", []) or [])
            ]

            explanation = {
                "strategy": strategy,
                "fts_matches": fts_matches,
                "vector_matches": vector_matches,
                "rrf_scores": rrf_scores,
                "cornerstone_injected": cornerstone_injected,
                "token_budget": {
                    "total": total_budget,
                    "used": tokens_used,
                    "remaining": remaining,
                },
            }

        except Exception as exc:
            err = f"Retrieve/explain failed: {exc}"
            logger.error(err, exc_info=True)
            errors.append(err)

        return RetrieveExplainResponse(
            results=results,
            explanation=explanation,
            errors=errors,
        )

    # ---- Canonical: /memories/retrieve  (issue #297) ----
    @router.post(
        f"{PREFIX}/memories/retrieve",
        summary="Enhanced agent retrieval with optional explain",
        tags=["memory"],
        operation_id="memories_retrieve",
    )
    async def memories_retrieve(
        request: Request,
        body: EnhancedRetrieveRequest,
        shadow: bool = Query(False, description="Shadow mode — flag response as shadow"),
        _: None = Depends(verify_api_key),
        plugin: Any = Depends(get_plugin),
    ) -> JSONResponse:
        """Retrieve relevant memories for a query. With ``explain=true``, includes scoring details.

        Canonical path: ``POST /api/v1/memories/retrieve``.
        """
        body.owner_id = _resolve_owner_id(request, body.owner_id)
        return await _do_retrieve_v2(body, shadow, plugin)

    # ---- DEPRECATED: /retrieve/v2  (issue #297) ----
    @router.post(
        f"{PREFIX}/retrieve/v2",
        summary="Enhanced agent retrieval — DEPRECATED, use /memories/retrieve",
        tags=["memory (deprecated)"],
        operation_id="retrieve_v2",
        deprecated=True,
    )
    async def retrieve_v2(
        request: Request,
        body: EnhancedRetrieveRequest,
        response: Response,
        shadow: bool = Query(False, description="Shadow mode — flag response as shadow"),
        _: None = Depends(verify_api_key),
        plugin: Any = Depends(get_plugin),
    ) -> JSONResponse:
        """DEPRECATED — use ``POST /api/v1/memories/retrieve``.

        .. deprecated:: sunset 2026-06-01  (issue #297)

        Security (#818): owner_id resolved from auth middleware, not body.
        """
        add_deprecation_headers(response, f"{PREFIX}/memories/retrieve")
        body.owner_id = _resolve_owner_id(request, body.owner_id)
        return await _do_retrieve_v2(body, shadow, plugin)

    async def _do_retrieve_v2(
        body: EnhancedRetrieveRequest,
        shadow: bool,
        plugin: Any,
    ) -> JSONResponse:
        """Shared implementation for retrieve v2 / memories/retrieve endpoints."""
        import time as _time

        effective_mode = body.strategy or body.mode
        effective_mode = _MODE_ALIASES.get(effective_mode, effective_mode)

        # Build temporal constraints when before_date/after_date are supplied (#513)
        _constraints = None
        _before = getattr(body, "before_date", None)
        _after = getattr(body, "after_date", None)
        _entity_id = getattr(body, "entity_id", None)
        _include_expired = getattr(body, "include_expired", False)
        _exclude_source_session = getattr(body, "exclude_source_session", None)
        if _before or _after or _entity_id or _include_expired or _exclude_source_session:
            try:
                from cortex.core.retrieval import RetrievalConstraints
                _constraints = RetrievalConstraints(
                    token_budget=body.token_budget,
                    before_date=_before,
                    after_date=_after,
                    entity_id=_entity_id,
                    include_expired=_include_expired,
                    exclude_source_session=_exclude_source_session,
                )
            except Exception:
                pass  # fall through; constraints stay None

        result = await plugin.retrieve(
            query=body.query,
            token_budget=body.token_budget,
            mode=effective_mode,
            owner_id=body.owner_id,
            constraints=_constraints,
        )
        if result is None:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="Retrieval pipeline not available",
            )

        data = _serialise(result.__dict__) if hasattr(result, "__dict__") else _serialise(result)

        if shadow:
            data["shadow"] = True
            _emit_shadow_event(endpoint="/api/v1/memories/retrieve", session_id="shadow")

        if body.explain:
            data["_explain"] = {
                "mode":        getattr(result, "mode", body.mode),
                "tokens_used": getattr(result, "tokens_used", None),
                "item_count":  len(getattr(result, "items", []) or []),
                "scores": [
                    {
                        "id":    getattr(i, "item_id", None) or getattr(i, "id", ""),
                        "score": getattr(i, "score", None) or getattr(i, "relevance_score", None),
                    }
                    for i in (getattr(result, "items", []) or [])
                ],
                "query": body.query,
            }

        return JSONResponse(content=data)
