"""RLM-inspired intelligent recall flow for memory retrieval.

Implements adaptive-depth retrieval with:
- LLM query distillation into targeted sub-queries
- Time-based filtering from temporal hints
- Parallel multi-query, multi-scope search
- Confidence-based routing with iterative deepening (budget loop)
- Evidence gap tracking propagated to results
"""

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor, as_completed
import contextvars
from datetime import datetime
import logging
from typing import Any
from uuid import uuid4

from pydantic import BaseModel, Field

from crewai.flow.flow import Flow, listen, router, start
from crewai.memory.analyze import QueryAnalysis, analyze_query
from crewai.memory.types import (
    _RECALL_OVERSAMPLE_FACTOR,
    MemoryConfig,
    MemoryMatch,
    MemoryRecord,
    compute_composite_score,
    embed_texts,
)


logger = logging.getLogger(__name__)


class RecallState(BaseModel):
    """State for the recall flow."""

    id: str = Field(default_factory=lambda: str(uuid4()))
    query: str = ""
    scope: str | None = None
    categories: list[str] | None = None
    time_cutoff: datetime | None = None
    source: str | None = None
    include_private: bool = False
    limit: int = 10
    query_embeddings: list[tuple[str, list[float]]] = Field(default_factory=list)
    query_analysis: QueryAnalysis | None = None
    candidate_scopes: list[str] = Field(default_factory=list)
    chunk_findings: list[Any] = Field(default_factory=list)
    evidence_gaps: list[str] = Field(default_factory=list)
    confidence: float = 0.0
    final_results: list[MemoryMatch] = Field(default_factory=list)
    exploration_budget: int = 1


class RecallFlow(Flow[RecallState]):
    """RLM-inspired intelligent memory recall flow.

    Analyzes the query via LLM to produce targeted sub-queries and filters,
    embeds each sub-query, searches across candidate scopes in parallel,
    and iteratively deepens exploration when confidence is low.
    """

    _skip_auto_memory: bool = True

    initial_state: type[RecallState] = RecallState

    def __init__(
        self,
        storage: Any,
        llm: Any,
        embedder: Any,
        config: MemoryConfig | None = None,
    ) -> None:
        super().__init__(suppress_flow_events=True)
        self._storage = storage
        self._llm = llm
        self._embedder = embedder
        self._config = config or MemoryConfig()

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _merged_categories(self) -> list[str] | None:
        """Return caller-supplied categories, or None if empty."""
        return self.state.categories or None

    def _do_search(self) -> list[dict[str, Any]]:
        """Run parallel search across (embeddings x scopes) with filters.

        Populates ``state.chunk_findings`` and ``state.confidence``.
        Returns the findings list.
        """
        search_categories = self._merged_categories()

        def _search_one(
            embedding: list[float], scope: str
        ) -> tuple[str, list[tuple[MemoryRecord, float]]]:
            raw = self._storage.search(
                embedding,
                scope_prefix=scope,
                categories=search_categories,
                limit=self.state.limit * _RECALL_OVERSAMPLE_FACTOR,
                min_score=0.0,
            )
            # Post-filter by time cutoff
            if self.state.time_cutoff and raw:
                raw = [(r, s) for r, s in raw if r.created_at >= self.state.time_cutoff]
            # Privacy filter
            if not self.state.include_private and raw:
                raw = [
                    (r, s)
                    for r, s in raw
                    if not r.private or r.source == self.state.source
                ]
            return scope, raw

        # Build (embedding, scope) task list
        tasks: list[tuple[list[float], str]] = [
            (embedding, scope)
            for _query_text, embedding in self.state.query_embeddings
            for scope in self.state.candidate_scopes
        ]

        findings: list[dict[str, Any]] = []

        if len(tasks) <= 1:
            for emb, sc in tasks:
                try:
                    scope, results = _search_one(emb, sc)
                except Exception:
                    logger.warning(
                        "Storage search failed in recall flow, skipping scope",
                        exc_info=True,
                    )
                    continue
                if results:
                    top_composite, _ = compute_composite_score(
                        results[0][0], results[0][1], self._config
                    )
                    findings.append(
                        {
                            "scope": scope,
                            "results": results,
                            "top_score": top_composite,
                        }
                    )
        else:
            with ThreadPoolExecutor(max_workers=min(len(tasks), 4)) as pool:
                futures = {
                    pool.submit(contextvars.copy_context().run, _search_one, emb, sc): (
                        emb,
                        sc,
                    )
                    for emb, sc in tasks
                }
                for future in as_completed(futures):
                    try:
                        scope, results = future.result()
                    except Exception:
                        logger.warning(
                            "Storage search failed in recall flow, skipping scope",
                            exc_info=True,
                        )
                        continue
                    if results:
                        top_composite, _ = compute_composite_score(
                            results[0][0], results[0][1], self._config
                        )
                        findings.append(
                            {
                                "scope": scope,
                                "results": results,
                                "top_score": top_composite,
                            }
                        )

        self.state.chunk_findings = findings
        self.state.confidence = max((f["top_score"] for f in findings), default=0.0)
        return findings

    # ------------------------------------------------------------------
    # Flow steps
    # ------------------------------------------------------------------

    @start()
    def analyze_query_step(self) -> QueryAnalysis:
        """Analyze the query, embed distilled sub-queries, extract filters.

        Short queries (below ``query_analysis_threshold`` characters) skip
        the LLM call entirely and embed the raw query directly -- saving
        ~1-3s per recall. Longer queries (e.g. full task descriptions)
        benefit from LLM distillation into targeted sub-queries.

        Sub-queries are embedded in a single batch ``embed_texts()`` call
        rather than sequential ``embed_text()`` calls.
        """
        self.state.exploration_budget = self._config.exploration_budget

        query_len = len(self.state.query)
        skip_llm = query_len < self._config.query_analysis_threshold

        if skip_llm:
            # Short query: skip LLM, embed raw query directly
            analysis = QueryAnalysis(
                keywords=[],
                suggested_scopes=[],
                complexity="simple",
                recall_queries=[self.state.query],
            )
            self.state.query_analysis = analysis
        else:
            # Long query: use LLM to distill sub-queries and extract filters
            available = self._storage.list_scopes(self.state.scope or "/")
            if not available:
                available = ["/"]
            scope_info = (
                self._storage.get_scope_info(self.state.scope or "/")
                if self.state.scope
                else None
            )
            analysis = analyze_query(
                self.state.query,
                available,
                scope_info,
                self._llm,
            )
            self.state.query_analysis = analysis

            # Parse time_filter into a datetime cutoff
            if analysis.time_filter:
                try:
                    self.state.time_cutoff = datetime.fromisoformat(
                        analysis.time_filter
                    )
                except ValueError:
                    pass

        # Batch-embed all sub-queries in ONE call
        queries = (
            analysis.recall_queries if analysis.recall_queries else [self.state.query]
        )
        queries = queries[:3]
        embeddings = embed_texts(self._embedder, queries)
        pairs: list[tuple[str, list[float]]] = [
            (q, emb) for q, emb in zip(queries, embeddings, strict=False) if emb
        ]
        if not pairs:
            # Fallback: embed the raw query if distilled queries all failed
            fallback_emb = embed_texts(self._embedder, [self.state.query])
            if fallback_emb and fallback_emb[0]:
                pairs = [(self.state.query, fallback_emb[0])]
        self.state.query_embeddings = pairs
        return analysis

    @listen(analyze_query_step)
    def filter_and_chunk(self) -> list[str]:
        """Select candidate scopes based on LLM analysis."""
        analysis = self.state.query_analysis
        scope_prefix = (self.state.scope or "/").rstrip("/") or "/"
        if analysis and analysis.suggested_scopes:
            candidates = [s for s in analysis.suggested_scopes if s]
        else:
            try:
                candidates = self._storage.list_scopes(scope_prefix)
            except Exception:
                logger.warning(
                    "Storage list_scopes failed in filter_and_chunk, "
                    "falling back to scope prefix",
                    exc_info=True,
                )
                candidates = []
        if not candidates:
            candidates = [scope_prefix]
        self.state.candidate_scopes = candidates[:20]
        return self.state.candidate_scopes

    @listen(filter_and_chunk)
    def search_chunks(self) -> list[Any]:
        """Initial parallel search across (embeddings x scopes) with filters."""
        return self._do_search()

    @router(search_chunks)
    def decide_depth(self) -> str:
        """Route based on confidence, complexity, and remaining budget."""
        analysis = self.state.query_analysis
        if (
            analysis
            and analysis.complexity == "complex"
            and self.state.confidence < self._config.complex_query_threshold
        ):
            if self.state.exploration_budget > 0:
                return "explore_deeper"
        if self.state.confidence >= self._config.confidence_threshold_high:
            return "synthesize"
        if (
            self.state.exploration_budget > 0
            and self.state.confidence < self._config.confidence_threshold_low
        ):
            return "explore_deeper"
        return "synthesize"

    @listen("explore_deeper")
    def recursive_exploration(self) -> list[Any]:
        """Feed top results back to LLM for deeper context extraction.

        Decrements the exploration budget so the loop terminates.
        """
        self.state.exploration_budget -= 1

        enhanced = []
        for finding in self.state.chunk_findings:
            if not finding.get("results"):
                continue
            content_parts = [r[0].content for r in finding["results"][:5]]
            chunk_text = "\n---\n".join(content_parts)
            prompt = (
                f"Query: {self.state.query}\n\n"
                f"Relevant memory excerpts:\n{chunk_text}\n\n"
                "Extract the most relevant information for the query. "
                "If something is missing, say what's missing in one short line."
            )
            try:
                response = self._llm.call([{"role": "user", "content": prompt}])
                if isinstance(response, str) and "missing" in response.lower():
                    self.state.evidence_gaps.append(response[:200])
                enhanced.append(
                    {
                        "scope": finding["scope"],
                        "extraction": response,
                        "results": finding["results"],
                    }
                )
            except Exception:
                enhanced.append(
                    {
                        "scope": finding["scope"],
                        "extraction": "",
                        "results": finding["results"],
                    }
                )
        self.state.chunk_findings = enhanced
        return enhanced

    @listen(recursive_exploration)
    def re_search(self) -> list[Any]:
        """Re-search after exploration to update confidence for the router loop."""
        return self._do_search()

    @router(re_search)
    def re_decide_depth(self) -> str:
        """Re-evaluate depth after re-search. Same logic as decide_depth."""
        return self.decide_depth()  # type: ignore[call-arg]

    @listen("synthesize")
    def synthesize_results(self) -> list[MemoryMatch]:
        """Deduplicate, composite-score, rank, and attach evidence gaps."""
        seen_ids: set[str] = set()
        matches: list[MemoryMatch] = []
        for finding in self.state.chunk_findings:
            if not isinstance(finding, dict):
                continue
            results = finding.get("results", [])
            if not isinstance(results, list):
                continue
            for item in results:
                if isinstance(item, (list, tuple)) and len(item) >= 2:
                    record, score = item[0], item[1]
                else:
                    continue
                if isinstance(record, MemoryRecord) and record.id not in seen_ids:
                    seen_ids.add(record.id)
                    composite, reasons = compute_composite_score(
                        record, float(score), self._config
                    )
                    matches.append(
                        MemoryMatch(
                            record=record,
                            score=composite,
                            match_reasons=reasons,
                        )
                    )
        matches.sort(key=lambda m: m.score, reverse=True)
        self.state.final_results = matches[: self.state.limit]

        # Attach evidence gaps to the first result so callers can inspect them
        if self.state.evidence_gaps and self.state.final_results:
            self.state.final_results[0].evidence_gaps = list(self.state.evidence_gaps)

        return self.state.final_results
