"""
tests/test_graph_visualization.py — Tests for Graph Visualization (#185).

Tests the get_entity_graph, get_entity_neighbors, and find_shortest_path
methods in SQLiteAdapter, plus HTTP endpoint integration.
"""
from __future__ import annotations

import pytest
import pytest_asyncio

from cortex.storage import SQLiteAdapter
from cortex.types import Entity, SemanticClaim


# ---------------------------------------------------------------------------
# Fixture
# ---------------------------------------------------------------------------

@pytest_asyncio.fixture
async def adapter():
    """Fresh in-memory SQLiteAdapter with test data."""
    a = SQLiteAdapter(db_path=":memory:", busy_timeout_ms=1000, enable_wal=False)
    await a.initialize()
    
    # Create test entities
    entity1 = await a.create_entity(
        canonical_name="Andrew",
        entity_type="person",
        owner_id="default",
    )
    entity2 = await a.create_entity(
        canonical_name="Eva",
        entity_type="person",
        owner_id="default",
    )
    entity3 = await a.create_entity(
        canonical_name="Project X",
        entity_type="project",
        owner_id="default",
    )
    entity4 = await a.create_entity(
        canonical_name="Office",
        entity_type="place",
        owner_id="default",
    )
    entity5 = await a.create_entity(
        canonical_name="Acme Corp",
        entity_type="organization",
        owner_id="default",
    )
    
    # Create claims linking entities
    await a.create_claim(
        subject="Andrew",
        predicate="works_with",
        object_value="Eva",
        owner_id="default",
        subject_entity_id=entity1.id,
        object_entity_id=entity2.id,
    )
    await a.create_claim(
        subject="Andrew",
        predicate="works_on",
        object_value="Project X",
        owner_id="default",
        subject_entity_id=entity1.id,
        object_entity_id=entity3.id,
    )
    await a.create_claim(
        subject="Eva",
        predicate="works_on",
        object_value="Project X",
        owner_id="default",
        subject_entity_id=entity2.id,
        object_entity_id=entity3.id,
    )
    await a.create_claim(
        subject="Andrew",
        predicate="located_at",
        object_value="Office",
        owner_id="default",
        subject_entity_id=entity1.id,
        object_entity_id=entity4.id,
    )
    await a.create_claim(
        subject="Acme Corp",
        predicate="employs",
        object_value="Andrew",
        owner_id="default",
        subject_entity_id=entity5.id,
        object_entity_id=entity1.id,
    )
    await a.create_claim(
        subject="Acme Corp",
        predicate="employs",
        object_value="Eva",
        owner_id="default",
        subject_entity_id=entity5.id,
        object_entity_id=entity2.id,
    )
    
    yield a
    await a.close()


# ---------------------------------------------------------------------------
# Tests for get_entity_graph
# ---------------------------------------------------------------------------

class TestEntityGraph:
    async def test_full_graph_returns_nodes_and_edges(self, adapter: SQLiteAdapter):
        """Test that full graph returns entities and their relationships."""
        result = await adapter.get_entity_graph(owner_id="default")
        
        assert "nodes" in result
        assert "edges" in result
        assert "total_entities" in result
        assert "total_relationships" in result
        assert result["total_entities"] >= 3  # At least our test entities
        assert result["total_relationships"] >= 3  # At least our test claims
    
    async def test_full_graph_has_correct_structure(self, adapter: SQLiteAdapter):
        """Test that nodes have correct properties."""
        result = await adapter.get_entity_graph(owner_id="default")
        
        for node in result["nodes"]:
            assert "id" in node
            assert "label" in node
            assert "type" in node
            assert "size" in node
            assert "color" in node
        
        for edge in result["edges"]:
            assert "source" in edge
            assert "target" in edge
            assert "label" in edge
            assert "weight" in edge
    
    async def test_center_entity_filtering(self, adapter: SQLiteAdapter):
        """Test filtering with a center entity."""
        # Get all entities first to find Andrew's ID
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew_node = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        assert andrew_node is not None
        
        # Now query with center
        result = await adapter.get_entity_graph(
            owner_id="default",
            center_entity_id=andrew_node["id"],
            max_depth=1,
        )
        
        # Should include Andrew and immediate neighbors
        assert result["total_entities"] >= 2
        labels = [n["label"] for n in result["nodes"]]
        assert "Andrew" in labels
    
    async def test_max_depth_filtering(self, adapter: SQLiteAdapter):
        """Test that max_depth limits expansion."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew_node = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        assert andrew_node is not None
        
        # depth=1 should return fewer nodes than depth=2
        result_depth1 = await adapter.get_entity_graph(
            owner_id="default",
            center_entity_id=andrew_node["id"],
            max_depth=1,
        )
        result_depth2 = await adapter.get_entity_graph(
            owner_id="default",
            center_entity_id=andrew_node["id"],
            max_depth=2,
        )
        
        # Depth 2 should have at least as many entities as depth 1
        assert result_depth2["total_entities"] >= result_depth1["total_entities"]
    
    async def test_max_nodes_limit(self, adapter: SQLiteAdapter):
        """Test that max_nodes limits the returned nodes."""
        result = await adapter.get_entity_graph(
            owner_id="default",
            max_nodes=2,
        )
        
        assert len(result["nodes"]) <= 2
    
    async def test_empty_graph_no_entities(self, adapter: SQLiteAdapter):
        """Test empty result when no entities exist."""
        # Create a fresh adapter with no data
        empty_adapter = SQLiteAdapter(db_path=":memory:", busy_timeout_ms=1000, enable_wal=False)
        await empty_adapter.initialize()
        
        result = await empty_adapter.get_entity_graph(owner_id="default")
        
        assert result["nodes"] == []
        assert result["edges"] == []
        assert result["total_entities"] == 0
        assert result["total_relationships"] == 0
        
        await empty_adapter.close()


# ---------------------------------------------------------------------------
# Tests for get_entity_neighbors
# ---------------------------------------------------------------------------

class TestEntityNeighbors:
    async def test_neighbors_returns_direct_connections(self, adapter: SQLiteAdapter):
        """Test that neighbors returns only direct connections."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew_node = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        assert andrew_node is not None
        
        result = await adapter.get_entity_neighbors(
            entity_id=andrew_node["id"],
            owner_id="default",
        )
        
        # Should include Andrew and his direct neighbors
        assert result["total_entities"] >= 2
        labels = [n["label"] for n in result["nodes"]]
        assert "Andrew" in labels
    
    async def test_neighbors_edge_labels(self, adapter: SQLiteAdapter):
        """Test that edges have correct labels."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew_node = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        assert andrew_node is not None
        
        result = await adapter.get_entity_neighbors(
            entity_id=andrew_node["id"],
            owner_id="default",
        )
        
        # Edges should have labels (predicates)
        for edge in result["edges"]:
            assert edge["label"] is not None
    
    async def test_neighbors_nonexistent_entity(self, adapter: SQLiteAdapter):
        """Test behavior with nonexistent entity."""
        result = await adapter.get_entity_neighbors(
            entity_id="nonexistent-id",
            owner_id="default",
        )
        
        assert result["nodes"] == []
        assert result["edges"] == []


# ---------------------------------------------------------------------------
# Tests for find_shortest_path
# ---------------------------------------------------------------------------

class TestShortestPath:
    async def test_direct_path(self, adapter: SQLiteAdapter):
        """Test finding direct path between connected entities."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        eva = next((n for n in all_result["nodes"] if n["label"] == "Eva"), None)
        
        assert andrew is not None
        assert eva is not None
        
        path = await adapter.find_shortest_path(
            source_id=andrew["id"],
            target_id=eva["id"],
            owner_id="default",
        )
        
        # Should find a path (they work together)
        assert len(path) > 0
    
    async def test_path_not_found(self, adapter: SQLiteAdapter):
        """Test path not found when entities aren't connected."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        office = next((n for n in all_result["nodes"] if n["label"] == "Office"), None)
        
        # Create disconnected entity
        disconnected = await adapter.create_entity(
            canonical_name="Isolated",
            entity_type="person",
            owner_id="default",
        )
        
        path = await adapter.find_shortest_path(
            source_id=office["id"],
            target_id=disconnected.id,
            owner_id="default",
            max_hops=2,
        )
        
        # Should not find path within 2 hops
        assert len(path) == 0
    
    async def test_same_entity_path(self, adapter: SQLiteAdapter):
        """Test path from entity to itself."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        
        path = await adapter.find_shortest_path(
            source_id=andrew["id"],
            target_id=andrew["id"],
            owner_id="default",
        )
        
        # Should return trivial path
        assert len(path) == 1
        assert path[0]["entity_id"] == andrew["id"]
    
    async def test_max_hops_limit(self, adapter: SQLiteAdapter):
        """Test that max_hops limits the search."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        
        # Create disconnected entity
        disconnected = await adapter.create_entity(
            canonical_name="Isolated",
            entity_type="person",
            owner_id="default",
        )
        
        # With max_hops=1, should not find path to disconnected entity
        path = await adapter.find_shortest_path(
            source_id=andrew["id"],
            target_id=disconnected.id,
            owner_id="default",
            max_hops=1,
        )
        
        # Should not find path within 1 hop (they're not connected)
        assert len(path) == 0


# ---------------------------------------------------------------------------
# Tests for node sizing and coloring
# ---------------------------------------------------------------------------

class TestNodeStyling:
    async def test_node_size_based_on_mention_count(self, adapter: SQLiteAdapter):
        """Test that node size is based on mention_count."""
        result = await adapter.get_entity_graph(owner_id="default")
        
        for node in result["nodes"]:
            assert node["size"] >= 1
            assert node["size"] <= 30  # Capped at 30
    
    async def test_node_color_based_on_type(self, adapter: SQLiteAdapter):
        """Test that node color is based on entity_type."""
        result = await adapter.get_entity_graph(owner_id="default")
        
        type_colors = {
            "person": "#4A90D9",
            "place": "#50C878",
            "organization": "#F5A623",
            "project": "#9013FE",
        }
        
        for node in result["nodes"]:
            expected_color = type_colors.get(node["type"], "#9B9B9B")
            assert node["color"] == expected_color
    
    async def test_mention_count_increases_size(self, adapter: SQLiteAdapter):
        """Test that higher mention_count = larger node."""
        # Note: create_entity doesn't currently support passing mention_count directly.
        # The size defaults to at least 1. This test verifies basic sizing works.
        result = await adapter.get_entity_graph(owner_id="default")
        
        # All nodes should have at least size 1
        for node in result["nodes"]:
            assert node["size"] >= 1
        
        # Verify we have at least some nodes
        assert len(result["nodes"]) >= 3


# ---------------------------------------------------------------------------
# Integration-style tests (HTTP endpoints)
# ---------------------------------------------------------------------------

class TestGraphHTTPEndpoints:
    async def test_graph_endpoint_data_structure(self, adapter: SQLiteAdapter):
        """Test that graph data matches expected HTTP endpoint structure."""
        result = await adapter.get_entity_graph(owner_id="default")
        
        # This is the structure we return from GET /api/v1/graph
        expected_keys = {"nodes", "edges", "total_entities", "total_relationships"}
        assert expected_keys == set(result.keys())
        
        # Validate node structure
        if result["nodes"]:
            node = result["nodes"][0]
            node_keys = {"id", "label", "type", "size", "color"}
            assert node_keys == set(node.keys())
        
        # Validate edge structure
        if result["edges"]:
            edge = result["edges"][0]
            edge_keys = {"source", "target", "label", "weight"}
            assert edge_keys == set(edge.keys())
    
    async def test_neighbors_endpoint_structure(self, adapter: SQLiteAdapter):
        """Test that neighbors data matches expected endpoint structure."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        
        result = await adapter.get_entity_neighbors(
            entity_id=andrew["id"],
            owner_id="default",
        )
        
        # Should have same structure as full graph
        assert "nodes" in result
        assert "edges" in result
        assert "total_entities" in result
        assert "total_relationships" in result
    
    async def test_paths_endpoint_structure(self, adapter: SQLiteAdapter):
        """Test that path data matches expected endpoint structure."""
        all_result = await adapter.get_entity_graph(owner_id="default")
        andrew = next((n for n in all_result["nodes"] if n["label"] == "Andrew"), None)
        eva = next((n for n in all_result["nodes"] if n["label"] == "Eva"), None)
        
        path_result = await adapter.find_shortest_path(
            source_id=andrew["id"],
            target_id=eva["id"],
            owner_id="default",
        )
        
        # Path should be a list of dicts with entity_id, predicate, step
        for step in path_result:
            assert "entity_id" in step
            assert "step" in step