"""
cortex/brain_graph/query.py — BrainGraphQueryEngine (M10b)

Read path for the brain graph index: takes a natural language query,
traverses the graph from embedding-similar entry points, assembles a
compact context block within a token budget, and returns a BrainGraphResult.

All public methods are async.
"""
from __future__ import annotations

import asyncio
import hashlib
import logging
import time
from collections import deque
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple

from cortex.config import CortexConfig
from cortex.llm.provider import LLMProvider
from cortex.storage.adapter import StorageAdapter
from cortex.types import BrainGraphEdge, BrainGraphIndex, BrainGraphNode

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Result type
# ---------------------------------------------------------------------------


@dataclass
class BrainGraphResult:
    """
    Result returned by BrainGraphQueryEngine.query().

    ``context`` is ready for direct injection into an agent context window;
    no post-processing is needed.
    """

    nodes: List[BrainGraphNode]
    """Ordered list of nodes included in the context (most relevant first)."""

    context: str
    """Assembled compact context string, ready for injection."""

    traversal_path: List[str]
    """All node IDs visited during BFS traversal (entry points first)."""

    token_count: int
    """Estimated token count of ``context``."""

    graph_id: Optional[str]
    """Graph that contributed the most nodes, or None for multi-graph results."""


# ---------------------------------------------------------------------------
# Internal scoring helper
# ---------------------------------------------------------------------------


@dataclass(order=True)
class _ScoredNode:
    """A node tagged with relevance score and BFS depth."""

    score: float = field(compare=True)
    depth: int = field(compare=False)
    node: BrainGraphNode = field(compare=False)
    similarity: float = field(compare=False, default=0.0)


# ---------------------------------------------------------------------------
# Token estimation
# ---------------------------------------------------------------------------

_CHARS_PER_TOKEN = 4  # rough estimate; tiktoken is preferred when available


def _estimate_tokens(text: str) -> int:
    """
    Estimate token count for *text*.

    Uses tiktoken (cl100k_base) when available; falls back to a
    character-based approximation (4 chars ≈ 1 token).
    """
    if not text:
        return 0
    try:
        import tiktoken  # type: ignore

        enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))
    except Exception:
        return max(1, len(text) // _CHARS_PER_TOKEN)


# ---------------------------------------------------------------------------
# Main engine
# ---------------------------------------------------------------------------


class BrainGraphQueryEngine:
    """
    Query engine for brain graph indexes.

    Given a natural language query it:
    1. Embeds the query with the configured embedding model.
    2. Finds top-K entry point nodes via cosine similarity.
    3. Expands via BFS up to *max_depth* hops.
    4. Scores each visited node (similarity + graph distance).
    5. Assembles the most relevant nodes into a compact context string
       that fits within *token_limit* tokens.
    6. Returns a :class:`BrainGraphResult`.

    Usage::

        engine = BrainGraphQueryEngine(storage, config, llm)
        result = await engine.query("how does auth work?", graph_id="abc123")
        print(result.context)
    """

    #: Number of entry-point nodes selected by embedding similarity.
    TOP_K_ENTRY_POINTS: int = 5

    #: Relative weight of embedding similarity vs. graph distance in scoring.
    SIMILARITY_WEIGHT: float = 0.7
    DISTANCE_WEIGHT: float = 0.3

    def __init__(
        self,
        storage: StorageAdapter,
        config: CortexConfig,
        llm: LLMProvider,
    ) -> None:
        self._storage = storage
        self._config = config
        self._llm = llm

        # Simple per-instance result cache: (query_hash, graph_id) → result
        self._cache: Dict[Tuple[str, Optional[str]], Tuple[float, BrainGraphResult]] = {}
        self._cache_ttl: float = 300.0  # seconds; configurable via query() kwargs

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def query(
        self,
        query: str,
        graph_id: Optional[str] = None,
        max_depth: int = 3,
        token_limit: int = 2000,
        owner_id: str = "default",
        *,
        cache_ttl: Optional[float] = None,
        similarity_weight: Optional[float] = None,
        distance_weight: Optional[float] = None,
    ) -> BrainGraphResult:
        """
        Query brain graph indexes and return assembled context.

        Args:
            query:             Natural language query string.
            graph_id:          Restrict search to this graph.  If None,
                               searches all active graphs for *owner_id*.
            max_depth:         Maximum BFS hops from entry points (default 3).
            token_limit:       Maximum tokens in the assembled context (default 2000).
            owner_id:          Owner namespace (default "default").
            cache_ttl:         Override default cache TTL in seconds.
                               Pass 0 to disable caching for this call.
            similarity_weight: Override SIMILARITY_WEIGHT for scoring.
            distance_weight:   Override DISTANCE_WEIGHT for scoring.

        Returns:
            :class:`BrainGraphResult` with nodes, context, traversal_path,
            token_count, and graph_id.

        Raises:
            ValueError: If *query* is empty.
            RuntimeError: If no active graphs are found for *owner_id*.
        """
        if not query.strip():
            raise ValueError("query must be non-empty")

        sim_w = similarity_weight if similarity_weight is not None else self.SIMILARITY_WEIGHT
        dist_w = distance_weight if distance_weight is not None else self.DISTANCE_WEIGHT
        ttl = cache_ttl if cache_ttl is not None else self._cache_ttl

        # --- Cache check ---------------------------------------------------
        cache_key = (self._query_hash(query, max_depth, token_limit), graph_id)
        if ttl > 0:
            cached = self._cache.get(cache_key)
            if cached is not None:
                cached_at, result = cached
                if time.monotonic() - cached_at < ttl:
                    logger.debug("brain_graph.query cache hit: %s", cache_key)
                    return result

        t0 = time.monotonic()

        # --- 1. Resolve graphs to search -----------------------------------
        graphs = await self._resolve_graphs(graph_id, owner_id)
        if not graphs:
            logger.warning(
                "brain_graph.query: no active graphs found for owner=%s graph_id=%s",
                owner_id,
                graph_id,
            )
            return BrainGraphResult(
                nodes=[],
                context="",
                traversal_path=[],
                token_count=0,
                graph_id=graph_id,
            )

        # --- 2. Embed query ------------------------------------------------
        query_embedding = await self._embed_query(query)

        # --- 3. Find entry points via similarity ----------------------------
        entry_points = await self._find_entry_points(
            query_embedding, graphs, graph_id
        )
        if not entry_points:
            logger.info("brain_graph.query: no entry points found for query=%r", query)
            return BrainGraphResult(
                nodes=[],
                context="",
                traversal_path=[],
                token_count=0,
                graph_id=graph_id,
            )

        # --- 4. Build adjacency maps for all target graphs -----------------
        adjacency = await self._build_adjacency(graphs)

        # --- 5. BFS traversal from entry points ----------------------------
        scored_nodes, traversal_path = await self._bfs_traverse(
            entry_points=entry_points,
            adjacency=adjacency,
            max_depth=max_depth,
            sim_weight=sim_w,
            dist_weight=dist_w,
        )

        # --- 6. Assemble context within token budget -----------------------
        selected, context, token_count = self._assemble_context(
            scored_nodes, token_limit
        )

        # --- 7. Determine primary graph_id ---------------------------------
        primary_graph = self._dominant_graph(selected) or graph_id

        result = BrainGraphResult(
            nodes=selected,
            context=context,
            traversal_path=traversal_path,
            token_count=token_count,
            graph_id=primary_graph,
        )

        elapsed = time.monotonic() - t0
        logger.info(
            "brain_graph.query completed: nodes=%d tokens=%d graphs=%d elapsed=%.3fs",
            len(selected),
            token_count,
            len(graphs),
            elapsed,
        )

        # --- Cache store ---------------------------------------------------
        if ttl > 0:
            self._cache[cache_key] = (time.monotonic(), result)

        return result

    async def traverse_bfs(
        self,
        start_node_id: str,
        graph_id: str,
        max_depth: int = 3,
    ) -> Tuple[List[BrainGraphNode], List[str]]:
        """
        BFS traversal starting from a specific node.

        Args:
            start_node_id: ID of the starting node.
            graph_id:      Graph to traverse.
            max_depth:     Maximum hops from start_node_id.

        Returns:
            (nodes, traversal_path) — ordered list of nodes visited and
            the ordered list of node IDs (BFS order).
        """
        graphs = await self._resolve_graphs(graph_id, owner_id="default")
        adjacency = await self._build_adjacency(graphs)

        start = await self._storage.get_brain_graph_node(start_node_id)
        if start is None:
            return [], []

        fake_entry = _ScoredNode(score=1.0, depth=0, node=start, similarity=1.0)
        scored, path = await self._bfs_traverse(
            entry_points=[fake_entry],
            adjacency=adjacency,
            max_depth=max_depth,
            sim_weight=1.0,
            dist_weight=0.0,
        )
        return [s.node for s in scored], path

    async def traverse_dfs(
        self,
        start_node_id: str,
        graph_id: str,
        max_depth: int = 3,
    ) -> Tuple[List[BrainGraphNode], List[str]]:
        """
        DFS traversal starting from a specific node.

        Args:
            start_node_id: ID of the starting node.
            graph_id:      Graph to traverse.
            max_depth:     Maximum depth from start_node_id.

        Returns:
            (nodes, traversal_path) — ordered list of nodes visited and
            the ordered list of node IDs (DFS order).
        """
        graphs = await self._resolve_graphs(graph_id, owner_id="default")
        adjacency = await self._build_adjacency(graphs)

        start = await self._storage.get_brain_graph_node(start_node_id)
        if start is None:
            return [], []

        visited: List[BrainGraphNode] = []
        path: List[str] = []
        seen: Set[str] = set()

        def _dfs(node: BrainGraphNode, depth: int) -> None:
            if node.id in seen or depth > max_depth:
                return
            seen.add(node.id)
            visited.append(node)
            path.append(node.id)
            for neighbor in adjacency.get(node.id, []):
                _dfs(neighbor, depth + 1)

        _dfs(start, 0)
        return visited, path

    async def get_subgraph(
        self,
        center_node_id: str,
        graph_id: str,
        n_hops: int = 2,
    ) -> Tuple[List[BrainGraphNode], List[BrainGraphEdge]]:
        """
        Extract the N-hop neighbourhood around *center_node_id*.

        Returns all nodes within N hops AND all edges where both
        source and target are in that node set.

        Args:
            center_node_id: Center node ID.
            graph_id:       Graph to search within.
            n_hops:         Neighbourhood radius.

        Returns:
            (nodes, edges) in the subgraph.
        """
        graphs = await self._resolve_graphs(graph_id, owner_id="default")
        adjacency = await self._build_adjacency(graphs)
        all_edges = await self._storage.get_brain_graph_edges(graph_id)

        nodes, _ = await self.traverse_bfs(center_node_id, graph_id, max_depth=n_hops)
        node_ids: Set[str] = {n.id for n in nodes}

        subgraph_edges = [
            e for e in all_edges
            if e.source_id in node_ids and e.target_id in node_ids
        ]
        return nodes, subgraph_edges

    async def find_path(
        self,
        source_node_id: str,
        target_node_id: str,
        graph_id: str,
        max_depth: int = 6,
    ) -> Optional[List[BrainGraphNode]]:
        """
        Find the shortest path between two nodes using BFS.

        Args:
            source_node_id: Starting node ID.
            target_node_id: Goal node ID.
            graph_id:       Graph to search.
            max_depth:      Maximum path length to search.

        Returns:
            Ordered list of BrainGraphNode from source to target
            (inclusive), or None if no path exists within max_depth.
        """
        graphs = await self._resolve_graphs(graph_id, owner_id="default")
        adjacency = await self._build_adjacency(graphs)

        if source_node_id == target_node_id:
            node = await self._storage.get_brain_graph_node(source_node_id)
            return [node] if node else None

        # BFS with parent tracking
        queue: deque[Tuple[str, List[str]]] = deque([(source_node_id, [source_node_id])])
        visited: Set[str] = {source_node_id}

        while queue:
            current_id, path = queue.popleft()
            if len(path) - 1 >= max_depth:
                continue
            for neighbor in adjacency.get(current_id, []):
                nid = neighbor.id
                if nid == target_node_id:
                    full_path_ids = path + [nid]
                    return await self._resolve_path(full_path_ids)
                if nid not in visited:
                    visited.add(nid)
                    queue.append((nid, path + [nid]))

        return None  # no path found

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    async def _resolve_graphs(
        self,
        graph_id: Optional[str],
        owner_id: str,
    ) -> List[BrainGraphIndex]:
        """Return the list of graphs to search."""
        if graph_id is not None:
            g = await self._storage.get_brain_graph(graph_id)
            return [g] if g is not None else []
        # Multi-graph: return all active graphs for the owner
        all_graphs = await self._storage.list_brain_graphs(owner_id)
        return [g for g in all_graphs if g.status == "active"]

    async def _embed_query(self, query: str) -> List[float]:
        """Embed the query string and return the vector."""
        resp = await self._llm.embed(
            texts=[query],
            model=self._config.embedding_model,
            input_type="query",
        )
        return resp.embeddings[0]

    async def _find_entry_points(
        self,
        query_embedding: List[float],
        graphs: List[BrainGraphIndex],
        graph_id: Optional[str],
    ) -> List[_ScoredNode]:
        """
        Find the top-K most similar nodes across all target graphs.

        Returns a list of _ScoredNode sorted by similarity descending.
        """
        # Gather candidate nodes via vector search (item_type="brain_graph_node")
        results = await self._storage.vector_search(
            query_embedding=query_embedding,
            item_type="brain_graph_node",
            top_k=self.TOP_K_ENTRY_POINTS * len(graphs),
        )

        # Filter to nodes that belong to our target graphs
        target_graph_ids: Set[str] = {g.id for g in graphs}
        scored: List[_ScoredNode] = []

        for vr in results:
            if len(scored) >= self.TOP_K_ENTRY_POINTS:
                break
            node = await self._storage.get_brain_graph_node(vr.item_id)
            if node is None:
                continue
            if node.graph_id not in target_graph_ids:
                continue
            scored.append(
                _ScoredNode(
                    score=vr.score,
                    depth=0,
                    node=node,
                    similarity=vr.score,
                )
            )

        return sorted(scored, key=lambda s: s.score, reverse=True)

    async def _build_adjacency(
        self,
        graphs: List[BrainGraphIndex],
    ) -> Dict[str, List[BrainGraphNode]]:
        """
        Build a node-id → list-of-neighbor-nodes adjacency map.

        Loads all edges (and the nodes at their targets) for all
        supplied graphs.  Edges are directed; traversal follows
        source → target direction.
        """
        # First, load all nodes into a lookup table
        node_lookup: Dict[str, BrainGraphNode] = {}
        for g in graphs:
            for node in await self._storage.get_brain_graph_nodes(g.id):
                node_lookup[node.id] = node

        # Then, build adjacency from edges
        adjacency: Dict[str, List[BrainGraphNode]] = {
            nid: [] for nid in node_lookup
        }
        for g in graphs:
            for edge in await self._storage.get_brain_graph_edges(g.id):
                src_id = edge.source_id
                tgt = node_lookup.get(edge.target_id)
                if tgt is None:
                    continue
                if src_id in adjacency:
                    adjacency[src_id].append(tgt)
                else:
                    adjacency[src_id] = [tgt]

        return adjacency

    async def _bfs_traverse(
        self,
        entry_points: List[_ScoredNode],
        adjacency: Dict[str, List[BrainGraphNode]],
        max_depth: int,
        sim_weight: float,
        dist_weight: float,
    ) -> Tuple[List[_ScoredNode], List[str]]:
        """
        BFS traversal from entry points.

        Scores each node by:
            score = sim_weight * similarity + dist_weight * (1 / (depth + 1))

        Returns (scored_nodes, traversal_path) both in BFS order.
        """
        visited: Set[str] = set()
        queue: deque[_ScoredNode] = deque()
        all_scored: List[_ScoredNode] = []
        traversal_path: List[str] = []

        for ep in entry_points:
            if ep.node.id not in visited:
                visited.add(ep.node.id)
                queue.append(ep)
                traversal_path.append(ep.node.id)
                all_scored.append(ep)

        while queue:
            current = queue.popleft()
            if current.depth >= max_depth:
                continue

            for neighbor in adjacency.get(current.node.id, []):
                if neighbor.id in visited:
                    continue
                visited.add(neighbor.id)
                new_depth = current.depth + 1

                # Distance-decayed score (inherits parent's similarity)
                distance_score = 1.0 / (new_depth + 1)
                combined_score = (
                    sim_weight * current.similarity
                    + dist_weight * distance_score
                )

                scored = _ScoredNode(
                    score=combined_score,
                    depth=new_depth,
                    node=neighbor,
                    similarity=current.similarity,
                )
                queue.append(scored)
                all_scored.append(scored)
                traversal_path.append(neighbor.id)

        # Sort by score descending for context assembly
        all_scored.sort(key=lambda s: s.score, reverse=True)
        return all_scored, traversal_path

    def _assemble_context(
        self,
        scored_nodes: List[_ScoredNode],
        token_limit: int,
    ) -> Tuple[List[BrainGraphNode], str, int]:
        """
        Greedily assemble context from the highest-scored nodes.

        Picks nodes in relevance order, stopping when adding the next
        node would exceed *token_limit*.

        Returns (selected_nodes, context_str, token_count).
        """
        selected: List[BrainGraphNode] = []
        parts: List[str] = []
        total_tokens = 0

        for sn in scored_nodes:
            node = sn.node
            # Prefer compact_repr; fall back to content, then label
            text = node.compact_repr or node.content or node.label
            if not text:
                continue

            node_tokens = _estimate_tokens(text)
            if total_tokens + node_tokens > token_limit:
                # Skip nodes that are too large but keep trying smaller ones
                continue

            selected.append(node)
            parts.append(text)
            total_tokens += node_tokens

        context = "\n\n".join(parts)
        return selected, context, total_tokens

    async def _resolve_path(self, node_ids: List[str]) -> List[BrainGraphNode]:
        """Fetch BrainGraphNode objects for a list of IDs (preserving order)."""
        nodes: List[BrainGraphNode] = []
        for nid in node_ids:
            n = await self._storage.get_brain_graph_node(nid)
            if n is not None:
                nodes.append(n)
        return nodes

    @staticmethod
    def _dominant_graph(nodes: List[BrainGraphNode]) -> Optional[str]:
        """Return the graph_id that appears most often in *nodes*."""
        if not nodes:
            return None
        counts: Dict[str, int] = {}
        for n in nodes:
            counts[n.graph_id] = counts.get(n.graph_id, 0) + 1
        return max(counts, key=lambda k: counts[k])

    @staticmethod
    def _query_hash(query: str, max_depth: int, token_limit: int) -> str:
        """Stable cache key for a query + params combination."""
        raw = f"{query}|{max_depth}|{token_limit}"
        return hashlib.sha256(raw.encode()).hexdigest()[:16]
