"""
Expandable Recall Engine — Phase 23.5

Two-phase recall:
  Phase 1 (fast): Retrieve summary nodes matching the query
  Phase 2 (expand): If expand_level > 1, fetch raw children of top summaries

This gives agents a summary-first view (low latency, low token cost) while
still enabling drill-down into the full raw memory trail.

Progressive recall stages:
  1. index search  (BM25/keyword on all nodes)
  2. summary filter (return summaries first)
  3. expand detail  (traverse DAG children on demand)
"""

import logging
import time
from typing import Dict, List, Optional

logger = logging.getLogger("openmemo.recall")


class ExpandableRecallEngine:
    """
    Summary-first DAG recall with configurable expand depth.

    Args:
        memory_graph    — MemoryGraph instance
        max_expand      — hard cap on expand_level (default 3)
    """

    def __init__(self, memory_graph, max_expand: int = 3):
        self._graph = memory_graph
        self._max_expand = max_expand

    def recall(self, query: str, scene: str = "",
               agent_id: str = "", expand_level: int = 1,
               top_k: int = 5) -> Dict:
        """
        Recall memories matching query with optional DAG expansion.

        Args:
            query        — search string
            scene        — filter by scene
            agent_id     — filter by agent
            expand_level — 1=summary only, 2+=include raw children
            top_k        — max summary nodes to return

        Returns:
            {
              "summary": "...",           # concatenated summary text
              "nodes": [...],             # summary node dicts
              "details": [...],           # raw children (if expand_level > 1)
              "expand_level": int,
              "total_raw": int,
              "query_time_ms": float,
            }
        """
        t0 = time.time()
        expand_level = max(1, min(expand_level, self._max_expand))

        summaries = self._graph.list_summaries(
            agent_id=agent_id, scene=scene, limit=top_k * 3
        )

        if summaries:
            ranked = self._rank(summaries, query, scene)[:top_k]
        else:
            raw_nodes = self._graph.list_raw(
                agent_id=agent_id, scene=scene, limit=top_k * 5
            )
            ranked = self._rank(raw_nodes, query, scene)[:top_k]

        summary_texts = [n.content for n in ranked]
        combined_summary = "\n".join(summary_texts) if summary_texts else ""

        details = []
        if expand_level > 1 and ranked:
            seen = set()
            for node in ranked:
                if node.type == "summary":
                    children = self._graph.get_children(node.id)
                    for c in children:
                        if c.id not in seen:
                            details.append(c.to_dict())
                            seen.add(c.id)
                else:
                    if node.id not in seen:
                        details.append(node.to_dict())
                        seen.add(node.id)

        elapsed = (time.time() - t0) * 1000
        return {
            "summary": combined_summary,
            "nodes": [n.to_dict() for n in ranked],
            "details": details,
            "expand_level": expand_level,
            "total_raw": len(details),
            "query_time_ms": round(elapsed, 2),
        }

    def expand(self, node_id: str) -> Dict:
        """
        Expand a single node to show its raw children.

        Returns:
            {
              "node": {...},
              "children": [...],
              "child_count": int,
            }
        """
        node = self._graph.get_node(node_id)
        if not node:
            return {"error": f"node {node_id} not found", "children": []}

        children = self._graph.get_children(node_id)
        return {
            "node": node.to_dict(),
            "children": [c.to_dict() for c in children],
            "child_count": len(children),
        }

    def progressive_recall(self, query: str, scene: str = "",
                           agent_id: str = "") -> Dict:
        """
        Three-stage progressive recall (most efficient for agents):
          Stage 1: keyword/BM25 index search
          Stage 2: summary filter
          Stage 3: expand detail on top result

        Returns condensed result ready for prompt injection.
        """
        stage1 = self.recall(query, scene=scene, agent_id=agent_id,
                             expand_level=1, top_k=10)

        top_nodes = stage1["nodes"][:3]
        if not top_nodes:
            return {"context": "", "stages": 1, "nodes_found": 0}

        stage2_text = "\n".join(n["content"] for n in top_nodes)

        details = []
        if top_nodes and top_nodes[0].get("type") == "summary":
            expanded = self.expand(top_nodes[0]["id"])
            details = expanded.get("children", [])[:5]

        context_parts = [stage2_text]
        if details:
            context_parts.append("\nDetails:")
            for d in details:
                context_parts.append(f"  - {d.get('content', '')}")

        return {
            "context": "\n".join(context_parts),
            "summary": stage2_text,
            "details": details,
            "stages": 3 if details else 2,
            "nodes_found": len(top_nodes),
        }

    def _rank(self, nodes: List, query: str, scene: str = "") -> List:
        """Simple keyword + scene + recency ranking."""
        if not nodes:
            return []
        query_terms = set(query.lower().split())
        now = time.time()

        def score(node) -> float:
            text = node.content.lower()
            keyword_score = sum(1.0 for t in query_terms if t in text)
            keyword_score /= max(len(query_terms), 1)

            scene_score = 0.2 if (scene and node.scene == scene) else 0.0

            age_days = (now - node.timestamp) / 86400
            recency = max(0.0, 1.0 - age_days / 30)

            importance = getattr(node, "importance", 0.5)
            is_summary = 0.15 if node.type == "summary" else 0.0

            return (keyword_score * 0.5 + recency * 0.2 +
                    scene_score + importance * 0.15 + is_summary)

        return sorted(nodes, key=score, reverse=True)
