"""
tests/test_brain_graph_query.py — Tests for M10b: BrainGraphQueryEngine.

Coverage:
- BrainGraphResult dataclass fields
- BrainGraphQueryEngine.query():
  - Single-graph query → returns relevant node in result
  - Multi-graph query (graph_id=None) → searches all active graphs
  - Empty query raises ValueError
  - No active graphs returns empty result
  - max_depth=0 → only entry points in context
  - token_limit respected
- traverse_bfs() — BFS order, depth limit
- traverse_dfs() — DFS order, depth limit
- get_subgraph() — N-hop neighbourhood, edges included
- find_path() — shortest path found, None when no path
- Result caching (cache_ttl > 0)
- _estimate_tokens() fallback (no tiktoken)
- Performance: query completes within 2s for <1000 nodes (mocked)
"""
from __future__ import annotations

import time
from typing import Any, Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
import pytest_asyncio

from cortex.brain_graph.query import (
    BrainGraphQueryEngine,
    BrainGraphResult,
    _estimate_tokens,
)
from cortex.config import CortexConfig
from cortex.storage import SQLiteAdapter
from cortex.types import BrainGraphEdge, BrainGraphIndex, BrainGraphNode


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest_asyncio.fixture
async def adapter():
    """Fresh in-memory SQLiteAdapter for each test."""
    a = SQLiteAdapter(db_path=":memory:", busy_timeout_ms=1000, enable_wal=False)
    await a.initialize()
    yield a
    await a.close()


@pytest.fixture
def config():
    return CortexConfig(
        embedding_model="text-embedding-3-small",
        embedding_dim=1536,
    )


@pytest.fixture
def mock_llm():
    """LLMProvider mock that returns canned embeddings."""
    llm = MagicMock()
    llm.DEFAULT_EMBED_MODEL = "text-embedding-3-small"

    async def mock_embed(texts, model=None, input_type="query"):
        resp = MagicMock()
        # Return a constant unit vector for all queries
        dim = 1536
        # Slightly vary embedding per text to allow similarity differences
        resp.embeddings = [
            [float(i % 10) / 10.0 for i in range(dim)]
            for _ in texts
        ]
        return resp

    llm.embed = mock_embed
    return llm


@pytest_asyncio.fixture
async def engine(adapter, config, mock_llm):
    """BrainGraphQueryEngine backed by in-memory SQLite."""
    return BrainGraphQueryEngine(storage=adapter, config=config, llm=mock_llm)


# ---------------------------------------------------------------------------
# Helper: build a fixture graph
# ---------------------------------------------------------------------------


async def _build_fixture_graph(
    adapter: SQLiteAdapter,
    name: str = "test-graph",
    owner_id: str = "default",
    n_nodes: int = 5,
) -> tuple[BrainGraphIndex, list[BrainGraphNode], list[BrainGraphEdge]]:
    """
    Create a small graph with *n_nodes* nodes and a chain of edges:
        node_0 → node_1 → node_2 → ...
    Each node has a compact_repr like "Node {i}: some useful content."
    """
    graph = await adapter.create_brain_graph(
        name=name, source_type="test", owner_id=owner_id, status="building"
    )

    nodes: list[BrainGraphNode] = []
    for i in range(n_nodes):
        node = await adapter.add_brain_graph_node(
            graph_id=graph.id,
            node_type="concept",
            label=f"Node{i}",
            content=f"This is the content for node {i}. It contains relevant information.",
            compact_repr=f"Node{i}: compact representation. Concept {i}. Detail {i}.",
        )
        nodes.append(node)

    edges: list[BrainGraphEdge] = []
    for i in range(n_nodes - 1):
        e = await adapter.add_brain_graph_edge(
            graph_id=graph.id,
            source_id=nodes[i].id,
            target_id=nodes[i + 1].id,
            edge_type="RELATED_TO",
            weight=1.0,
        )
        edges.append(e)

    # Store embeddings: node i gets a slightly different vector so
    # similarity is meaningful
    dim = 1536
    for i, node in enumerate(nodes):
        emb = [float((i + 1) % 10) / 10.0 for _ in range(dim)]
        await adapter.store_embedding(
            item_type="brain_graph_node",
            item_id=node.id,
            embedding=emb,
            model_id="text-embedding-3-small",
        )

    return graph, nodes, edges


# ---------------------------------------------------------------------------
# Unit tests — _estimate_tokens
# ---------------------------------------------------------------------------


class TestEstimateTokens:
    def test_empty_string(self):
        assert _estimate_tokens("") == 0  # empty string → 0 tokens

    def test_short_string(self):
        count = _estimate_tokens("hello world")  # 11 chars → ≥1
        assert count >= 1

    def test_longer_string(self):
        text = "word " * 100  # 500 chars → ~125 tokens
        count = _estimate_tokens(text)
        assert count > 50

    def test_tiktoken_fallback(self):
        """Verify fallback when tiktoken is unavailable."""
        with patch("cortex.brain_graph.query.tiktoken", None, create=True):
            # The function should still return an integer
            count = _estimate_tokens("hello world")
            assert isinstance(count, int)
            assert count >= 1


# ---------------------------------------------------------------------------
# Unit tests — BrainGraphResult
# ---------------------------------------------------------------------------


class TestBrainGraphResult:
    def test_result_fields(self):
        """BrainGraphResult must have the required fields."""
        r = BrainGraphResult(
            nodes=[],
            context="some context",
            traversal_path=["id1", "id2"],
            token_count=10,
            graph_id="g1",
        )
        assert r.nodes == []
        assert r.context == "some context"
        assert r.traversal_path == ["id1", "id2"]
        assert r.token_count == 10
        assert r.graph_id == "g1"

    def test_result_optional_graph_id(self):
        r = BrainGraphResult(
            nodes=[],
            context="",
            traversal_path=[],
            token_count=0,
            graph_id=None,
        )
        assert r.graph_id is None


# ---------------------------------------------------------------------------
# Integration tests — BrainGraphQueryEngine
# ---------------------------------------------------------------------------


class TestBrainGraphQueryEngine:

    # --- Constructor -------------------------------------------------------

    async def test_engine_init(self, adapter, config, mock_llm):
        engine = BrainGraphQueryEngine(adapter, config, mock_llm)
        assert engine is not None

    # --- query() -----------------------------------------------------------

    async def test_query_raises_on_empty_string(self, engine):
        with pytest.raises(ValueError, match="non-empty"):
            await engine.query("")

    async def test_query_raises_on_whitespace_only(self, engine):
        with pytest.raises(ValueError, match="non-empty"):
            await engine.query("   ")

    async def test_query_no_active_graphs_returns_empty(self, engine):
        """No graphs in DB → empty result, no exception."""
        result = await engine.query("anything", owner_id="nobody")
        assert isinstance(result, BrainGraphResult)
        assert result.nodes == []
        assert result.context == ""
        assert result.token_count == 0
        assert result.traversal_path == []

    async def test_query_single_graph_returns_result(self, adapter, engine):
        """Basic smoke test: build a fixture graph and assert result is populated."""
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        # Activate the graph
        await adapter.update_brain_graph(graph.id, status="active")

        result = await engine.query(
            "some relevant concept",
            graph_id=graph.id,
            max_depth=3,
        )
        assert isinstance(result, BrainGraphResult)
        assert len(result.nodes) > 0
        assert result.context != ""
        assert result.token_count > 0
        assert len(result.traversal_path) > 0
        assert result.graph_id == graph.id

    async def test_query_entry_point_in_result(self, adapter, engine):
        """The most-similar entry point node should appear in the result."""
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        await adapter.update_brain_graph(graph.id, status="active")

        result = await engine.query("relevant", graph_id=graph.id)
        result_ids = {n.id for n in result.nodes}
        # At least one of the graph's nodes must be in the result
        assert len(result_ids.intersection({n.id for n in nodes})) > 0

    async def test_query_traversal_path_contains_visited_nodes(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=4)
        await adapter.update_brain_graph(graph.id, status="active")

        result = await engine.query("concept", graph_id=graph.id, max_depth=3)
        # traversal_path must be a superset of result node IDs
        result_ids = {n.id for n in result.nodes}
        path_ids = set(result.traversal_path)
        assert result_ids.issubset(path_ids)

    async def test_query_token_limit_respected(self, adapter, engine):
        """Context must not exceed the specified token_limit."""
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=10)
        await adapter.update_brain_graph(graph.id, status="active")

        limit = 50  # very tight
        result = await engine.query(
            "concept", graph_id=graph.id, token_limit=limit
        )
        # token_count must be ≤ limit
        assert result.token_count <= limit

    async def test_query_max_depth_zero_only_entry_points(self, adapter, engine):
        """max_depth=0 means no traversal beyond entry points."""
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        await adapter.update_brain_graph(graph.id, status="active")

        result = await engine.query(
            "concept", graph_id=graph.id, max_depth=0
        )
        # With depth=0, we should only get the entry points themselves
        assert len(result.nodes) <= BrainGraphQueryEngine.TOP_K_ENTRY_POINTS

    async def test_query_multi_graph_no_graph_id(self, adapter, engine):
        """When graph_id=None, all active graphs for the owner are searched."""
        g1, n1, _ = await _build_fixture_graph(adapter, name="g1", n_nodes=3)
        g2, n2, _ = await _build_fixture_graph(adapter, name="g2", n_nodes=3)
        await adapter.update_brain_graph(g1.id, status="active")
        await adapter.update_brain_graph(g2.id, status="active")

        result = await engine.query("concept", graph_id=None, owner_id="default")
        assert len(result.nodes) > 0
        # Result should contain nodes from at least one graph
        result_graph_ids = {n.graph_id for n in result.nodes}
        assert g1.id in result_graph_ids or g2.id in result_graph_ids

    async def test_query_inactive_graphs_excluded(self, adapter, engine):
        """Graphs with status != 'active' must not be searched."""
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        # Leave it in 'building' status (default)

        result = await engine.query("concept", graph_id=None, owner_id="default")
        assert result.nodes == []

    async def test_query_result_graph_id_populated(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        await adapter.update_brain_graph(graph.id, status="active")

        result = await engine.query("concept", graph_id=graph.id)
        assert result.graph_id == graph.id

    async def test_query_caching(self, adapter, engine):
        """Same query twice should return cached result (same object)."""
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        await adapter.update_brain_graph(graph.id, status="active")

        engine._cache_ttl = 60.0
        r1 = await engine.query("concept", graph_id=graph.id)
        r2 = await engine.query("concept", graph_id=graph.id)
        # Both results should be equal in content
        assert r1.context == r2.context
        assert r1.traversal_path == r2.traversal_path

    async def test_query_cache_disabled(self, adapter, engine):
        """cache_ttl=0 should bypass cache."""
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        await adapter.update_brain_graph(graph.id, status="active")

        r1 = await engine.query("concept", graph_id=graph.id, cache_ttl=0)
        r2 = await engine.query("concept", graph_id=graph.id, cache_ttl=0)
        # Should still work, just not use cache
        assert r1.context == r2.context

    async def test_query_performance_under_2s(self, adapter, config, mock_llm):
        """Query must complete in <2s for a graph with <1000 nodes."""
        engine = BrainGraphQueryEngine(adapter, config, mock_llm)

        # Build a moderately sized graph (50 nodes is sufficient for timing test)
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=50)
        await adapter.update_brain_graph(graph.id, status="active")

        t0 = time.monotonic()
        result = await engine.query("concept", graph_id=graph.id, max_depth=3)
        elapsed = time.monotonic() - t0

        assert elapsed < 2.0, f"Query took {elapsed:.2f}s, expected <2s"
        assert isinstance(result, BrainGraphResult)


# ---------------------------------------------------------------------------
# BFS traversal
# ---------------------------------------------------------------------------


class TestBFSTraversal:

    async def test_bfs_from_start_node(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        visited, path = await engine.traverse_bfs(
            nodes[0].id, graph.id, max_depth=4
        )
        # All nodes in a chain should be reachable
        assert len(visited) == 5
        assert len(path) == 5

    async def test_bfs_respects_max_depth(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        visited, path = await engine.traverse_bfs(
            nodes[0].id, graph.id, max_depth=2
        )
        # Depth 0 = node[0], depth 1 = node[1], depth 2 = node[2]
        assert len(visited) == 3

    async def test_bfs_nonexistent_start_node(self, adapter, engine):
        visited, path = await engine.traverse_bfs(
            "nonexistent-id", "nonexistent-graph", max_depth=3
        )
        assert visited == []
        assert path == []

    async def test_bfs_single_node_graph(self, adapter, engine):
        graph = await adapter.create_brain_graph("solo", "test", status="active")
        node = await adapter.add_brain_graph_node(graph.id, "concept", "Only")
        visited, path = await engine.traverse_bfs(node.id, graph.id, max_depth=3)
        assert len(visited) == 1
        assert path == [node.id]

    async def test_bfs_start_node_first_in_path(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        _, path = await engine.traverse_bfs(nodes[0].id, graph.id, max_depth=3)
        assert path[0] == nodes[0].id

    async def test_bfs_no_cycles(self, adapter, engine):
        """Even with cycles, BFS must not visit the same node twice."""
        graph = await adapter.create_brain_graph("cyclic", "test", status="active")
        n1 = await adapter.add_brain_graph_node(graph.id, "concept", "N1")
        n2 = await adapter.add_brain_graph_node(graph.id, "concept", "N2")
        # Create a cycle: N1 → N2 → N1
        await adapter.add_brain_graph_edge(graph.id, n1.id, n2.id, "RELATED_TO")
        await adapter.add_brain_graph_edge(graph.id, n2.id, n1.id, "RELATED_TO")

        _, path = await engine.traverse_bfs(n1.id, graph.id, max_depth=10)
        assert len(path) == len(set(path)), "Duplicate nodes in BFS path"


# ---------------------------------------------------------------------------
# DFS traversal
# ---------------------------------------------------------------------------


class TestDFSTraversal:

    async def test_dfs_visits_all_reachable(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=4)
        visited, path = await engine.traverse_dfs(
            nodes[0].id, graph.id, max_depth=10
        )
        assert len(visited) == 4

    async def test_dfs_respects_max_depth(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        visited, path = await engine.traverse_dfs(
            nodes[0].id, graph.id, max_depth=1
        )
        # Only node[0] (depth 0) and node[1] (depth 1)
        assert len(visited) == 2

    async def test_dfs_nonexistent_start(self, adapter, engine):
        visited, path = await engine.traverse_dfs(
            "no-such-node", "no-such-graph", max_depth=3
        )
        assert visited == []
        assert path == []

    async def test_dfs_no_cycles(self, adapter, engine):
        graph = await adapter.create_brain_graph("dfs-cyclic", "test", status="active")
        n1 = await adapter.add_brain_graph_node(graph.id, "concept", "A")
        n2 = await adapter.add_brain_graph_node(graph.id, "concept", "B")
        await adapter.add_brain_graph_edge(graph.id, n1.id, n2.id, "RELATED_TO")
        await adapter.add_brain_graph_edge(graph.id, n2.id, n1.id, "RELATED_TO")

        _, path = await engine.traverse_dfs(n1.id, graph.id, max_depth=10)
        assert len(path) == len(set(path)), "Duplicate nodes in DFS path"


# ---------------------------------------------------------------------------
# Subgraph extraction
# ---------------------------------------------------------------------------


class TestGetSubgraph:

    async def test_subgraph_1_hop(self, adapter, engine):
        graph, nodes, edges = await _build_fixture_graph(adapter, n_nodes=4)
        sub_nodes, sub_edges = await engine.get_subgraph(
            nodes[0].id, graph.id, n_hops=1
        )
        # 1-hop from nodes[0]: nodes[0] and nodes[1]
        node_ids = {n.id for n in sub_nodes}
        assert nodes[0].id in node_ids
        assert nodes[1].id in node_ids
        assert len(sub_nodes) == 2

    async def test_subgraph_2_hops(self, adapter, engine):
        graph, nodes, edges = await _build_fixture_graph(adapter, n_nodes=5)
        sub_nodes, sub_edges = await engine.get_subgraph(
            nodes[0].id, graph.id, n_hops=2
        )
        # 2-hops from node[0]: node[0], node[1], node[2]
        assert len(sub_nodes) == 3

    async def test_subgraph_edges_within_node_set(self, adapter, engine):
        graph, nodes, edges = await _build_fixture_graph(adapter, n_nodes=4)
        sub_nodes, sub_edges = await engine.get_subgraph(
            nodes[0].id, graph.id, n_hops=3
        )
        node_ids = {n.id for n in sub_nodes}
        for edge in sub_edges:
            assert edge.source_id in node_ids
            assert edge.target_id in node_ids

    async def test_subgraph_full_graph_large_hops(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        sub_nodes, sub_edges = await engine.get_subgraph(
            nodes[0].id, graph.id, n_hops=100
        )
        assert len(sub_nodes) == 5

    async def test_subgraph_isolated_node(self, adapter, engine):
        graph = await adapter.create_brain_graph("iso", "test", status="active")
        lone = await adapter.add_brain_graph_node(graph.id, "concept", "Lone")
        sub_nodes, sub_edges = await engine.get_subgraph(lone.id, graph.id, n_hops=3)
        assert len(sub_nodes) == 1
        assert sub_edges == []


# ---------------------------------------------------------------------------
# Path finding
# ---------------------------------------------------------------------------


class TestFindPath:

    async def test_path_source_to_target(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        path = await engine.find_path(nodes[0].id, nodes[4].id, graph.id)
        assert path is not None
        assert len(path) == 5
        assert path[0].id == nodes[0].id
        assert path[-1].id == nodes[4].id

    async def test_path_adjacent_nodes(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        path = await engine.find_path(nodes[0].id, nodes[1].id, graph.id)
        assert path is not None
        assert len(path) == 2

    async def test_path_same_source_and_target(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        path = await engine.find_path(nodes[0].id, nodes[0].id, graph.id)
        assert path is not None
        assert len(path) == 1
        assert path[0].id == nodes[0].id

    async def test_path_no_path_returns_none(self, adapter, engine):
        """Disconnected nodes: no path should return None."""
        graph = await adapter.create_brain_graph("discon", "test", status="active")
        n1 = await adapter.add_brain_graph_node(graph.id, "concept", "Island1")
        n2 = await adapter.add_brain_graph_node(graph.id, "concept", "Island2")
        # No edges between them
        path = await engine.find_path(n1.id, n2.id, graph.id)
        assert path is None

    async def test_path_max_depth_exceeded_returns_none(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=5)
        # Path length is 5 nodes (4 hops); limit to 2 hops
        path = await engine.find_path(
            nodes[0].id, nodes[4].id, graph.id, max_depth=2
        )
        assert path is None

    async def test_path_preserves_order(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=4)
        path = await engine.find_path(nodes[0].id, nodes[3].id, graph.id)
        assert path is not None
        assert [n.id for n in path] == [n.id for n in nodes]

    async def test_path_nonexistent_source(self, adapter, engine):
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        # BFS won't find anything starting from a nonexistent node
        path = await engine.find_path("ghost", nodes[2].id, graph.id)
        assert path is None


# ---------------------------------------------------------------------------
# Integration: full pipeline (fixture graph → query → assert)
# ---------------------------------------------------------------------------


class TestIntegration:

    async def test_full_pipeline_relevant_node_in_result(self, adapter, config):
        """
        Acceptance criterion: build fixture graph → query → relevant node in result.
        """
        # Custom LLM whose embedding for "auth" is closest to node 0
        class TunedLLM(MagicMock):
            async def embed(self, texts, model=None, input_type="query"):
                resp = MagicMock()
                dim = 1536
                # "query" embedding → all 0.9 (high similarity to node 0)
                resp.embeddings = [[0.9] * dim for _ in texts]
                return resp

        llm = TunedLLM()
        a = SQLiteAdapter(db_path=":memory:", busy_timeout_ms=1000, enable_wal=False)
        await a.initialize()
        try:
            # Build graph where node 0 has embedding close to query
            graph = await a.create_brain_graph(
                "auth-graph", "codebase", owner_id="default", status="active"
            )
            n0 = await a.add_brain_graph_node(
                graph.id, "concept", "Authentication",
                compact_repr="Auth: JWT Bearer token. POST /auth/login. Returns {access_token}.",
            )
            n1 = await a.add_brain_graph_node(
                graph.id, "concept", "User",
                compact_repr="User: entity with id, email, role.",
            )
            await a.add_brain_graph_edge(graph.id, n0.id, n1.id, "RELATED_TO")

            # Embedding: query ≈ n0 (all 0.9), n0 = 0.9, n1 = 0.1
            await a.store_embedding(
                item_type="brain_graph_node",
                item_id=n0.id,
                embedding=[0.9] * 1536,
                model_id="text-embedding-3-small",
            )
            await a.store_embedding(
                item_type="brain_graph_node",
                item_id=n1.id,
                embedding=[0.1] * 1536,
                model_id="text-embedding-3-small",
            )

            engine = BrainGraphQueryEngine(a, config, llm)
            result = await engine.query("how does authentication work?", graph_id=graph.id)

            assert isinstance(result, BrainGraphResult)
            assert len(result.nodes) > 0
            # The auth node (n0) must be in the result
            result_ids = {n.id for n in result.nodes}
            assert n0.id in result_ids, "Entry point node missing from result"
            # Context must contain the compact repr of n0
            assert "Auth:" in result.context or "authentication" in result.context.lower()
            # Traversal path must include n0
            assert n0.id in result.traversal_path
        finally:
            await a.close()

    async def test_context_ready_for_injection(self, adapter, config, mock_llm):
        """BrainGraphResult.context must be a non-empty string with no post-processing."""
        graph, nodes, _ = await _build_fixture_graph(adapter, n_nodes=3)
        await adapter.update_brain_graph(graph.id, status="active")

        engine = BrainGraphQueryEngine(adapter, config, mock_llm)
        result = await engine.query("concept", graph_id=graph.id)

        if result.nodes:
            assert isinstance(result.context, str)
            assert len(result.context) > 0
            # Must not contain raw Python repr artifacts
            assert "BrainGraphNode(" not in result.context
