"""Unit tests for LlamaIndex PropertyGraphStore — mocked, no DB required."""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from piggie._integrations.llamaindex_property import (
    PiggiePropertyGraphStore,
    _validate_label,
    _vertex_to_labelled_node,
)
from piggie.exceptions import IntegrationError, LabelError
from piggie.result import Result
from piggie.types import Edge, Vertex

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


@pytest.fixture()
def mock_driver():
    driver = MagicMock()
    driver.graph = "test_graph"
    # Default: sql and cypher return empty result
    driver.sql.return_value = Result([], [])
    driver.cypher.return_value = Result([], [])
    return driver


@pytest.fixture()
def store(mock_driver):
    """Store with table creation disabled to avoid setup noise."""
    return PiggiePropertyGraphStore(
        mock_driver, graph="test_graph", create_tables=False,
    )


# ---------------------------------------------------------------------------
# _validate_label
# ---------------------------------------------------------------------------


class TestValidateLabel:
    def test_valid(self):
        assert _validate_label("Person") == "Person"
        assert _validate_label("_node") == "_node"
        assert _validate_label("Label_123") == "Label_123"

    def test_invalid(self):
        with pytest.raises(LabelError, match="Invalid label"):
            _validate_label("123bad")
        with pytest.raises(LabelError, match="Invalid label"):
            _validate_label("has space")
        with pytest.raises(LabelError, match="Invalid label"):
            _validate_label("DROP;--")
        with pytest.raises(LabelError, match="Invalid label"):
            _validate_label("")


# ---------------------------------------------------------------------------
# Construction
# ---------------------------------------------------------------------------


class TestConstruction:
    def test_create_tables(self, mock_driver):
        PiggiePropertyGraphStore(
            mock_driver, graph="g", create_tables=True,
        )
        # Should call sql to create embedding table
        mock_driver.sql.assert_called_once()
        call_args = mock_driver.sql.call_args[0][0]
        assert "CREATE TABLE IF NOT EXISTS g_node_embeddings" in call_args
        assert "node_id TEXT PRIMARY KEY" in call_args
        mock_driver.commit.assert_called_once()

    def test_skip_table_creation(self, mock_driver):
        PiggiePropertyGraphStore(
            mock_driver, graph="g", create_tables=False,
        )
        mock_driver.sql.assert_not_called()

    def test_custom_embedding_table(self, mock_driver):
        PiggiePropertyGraphStore(
            mock_driver, graph="g",
            embedding_table="custom_emb",
            create_tables=True,
        )
        call_args = mock_driver.sql.call_args[0][0]
        assert "custom_emb" in call_args

    def test_embedding_dimensions(self, mock_driver):
        PiggiePropertyGraphStore(
            mock_driver, graph="g",
            embedding_dimensions=1536,
            create_tables=True,
        )
        call_args = mock_driver.sql.call_args[0][0]
        assert "vector(1536)" in call_args

    def test_supports_flags(self, mock_driver):
        store = PiggiePropertyGraphStore(
            mock_driver, graph="g", create_tables=False,
        )
        assert store.supports_structured_queries is True
        assert store.supports_vector_queries is True


# ---------------------------------------------------------------------------
# _vertex_to_labelled_node
# ---------------------------------------------------------------------------


class TestVertexToLabelledNode:
    @patch("piggie._integrations.llamaindex_property._ChunkNode", MagicMock())
    @patch("piggie._integrations.llamaindex_property._EntityNode")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_entity_node(self, mock_load, MockEntityNode):
        """Vertex with 'name' property → EntityNode."""
        MockEntityNode.return_value = SimpleNamespace(
            id_="42", label="Person", name="Alice"
        )
        v = Vertex(id=42, label="Person", properties={"name": "Alice", "age": 30})
        _vertex_to_labelled_node(v)
        MockEntityNode.assert_called_once_with(
            id_="42", label="Person", name="Alice", properties={"age": 30},
        )

    @patch("piggie._integrations.llamaindex_property._EntityNode", MagicMock())
    @patch("piggie._integrations.llamaindex_property._ChunkNode")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_chunk_node(self, mock_load, MockChunkNode):
        """Vertex with 'text' property → ChunkNode."""
        MockChunkNode.return_value = SimpleNamespace(
            id_="99", label="Chunk", text="hello"
        )
        v = Vertex(id=99, label="Chunk", properties={"text": "hello world"})
        _vertex_to_labelled_node(v)
        MockChunkNode.assert_called_once_with(
            id_="99", label="Chunk", text="hello world", properties={},
        )

    @patch("piggie._integrations.llamaindex_property._ChunkNode", MagicMock())
    @patch("piggie._integrations.llamaindex_property._EntityNode")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_id_property_preferred(self, mock_load, MockEntityNode):
        """Vertex with 'id' in properties uses that over vertex.id."""
        MockEntityNode.return_value = SimpleNamespace(id_="custom_id")
        v = Vertex(id=42, label="X", properties={"id": "custom_id"})
        _vertex_to_labelled_node(v)
        assert MockEntityNode.call_args[1]["id_"] == "custom_id"


# ---------------------------------------------------------------------------
# upsert_nodes
# ---------------------------------------------------------------------------


class TestUpsertNodes:
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    @patch("piggie._integrations.llamaindex_property._EntityNode", SimpleNamespace)
    def test_entity_node(self, mock_load, store, mock_driver):
        node = SimpleNamespace(
            id_="alice", label="Person",
            properties={"age": 30}, name="Alice",
            embedding=None,
        )
        # Make isinstance check work
        with patch("piggie._integrations.llamaindex_property._EntityNode", type(node)):
            store.upsert_nodes([node])
        mock_driver.cypher.assert_called_once()
        call_args = mock_driver.cypher.call_args
        assert "MERGE" in call_args[0][0]
        assert "Person" in call_args[0][0]
        mock_driver.commit.assert_called_once()

    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    @patch("piggie._integrations.llamaindex_property._EntityNode", None)
    @patch("piggie._integrations.llamaindex_property._ChunkNode", None)
    def test_node_with_embedding(self, mock_load, store, mock_driver):
        node = SimpleNamespace(
            id_="chunk1", label="Chunk",
            properties={"text": "hello"}, embedding=[0.1, 0.2, 0.3],
        )
        store.upsert_nodes([node])
        # Should call cypher (MERGE vertex) + sql (INSERT embedding)
        mock_driver.cypher.assert_called_once()
        mock_driver.sql.assert_called_once()
        sql_args = mock_driver.sql.call_args
        assert "INSERT INTO" in sql_args[0][0]
        assert "ON CONFLICT" in sql_args[0][0]
        mock_driver.commit.assert_called_once()

    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    @patch("piggie._integrations.llamaindex_property._EntityNode", None)
    @patch("piggie._integrations.llamaindex_property._ChunkNode", None)
    def test_node_without_embedding(self, mock_load, store, mock_driver):
        node = SimpleNamespace(
            id_="n1", label="Thing", properties={}, embedding=None,
        )
        store.upsert_nodes([node])
        mock_driver.cypher.assert_called_once()
        # No SQL call for embedding
        mock_driver.sql.assert_not_called()


# ---------------------------------------------------------------------------
# upsert_relations
# ---------------------------------------------------------------------------


class TestUpsertRelations:
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_basic_relation(self, mock_load, store, mock_driver):
        rel = SimpleNamespace(
            source_id="alice", target_id="bob", label="KNOWS",
            source_label="Person", target_label="Person",
            properties={},
        )
        store.upsert_relations([rel])
        mock_driver.cypher.assert_called_once()
        call_args = mock_driver.cypher.call_args
        assert "MERGE (s)-[r:KNOWS]->(t)" in call_args[0][0]
        assert call_args[1]["params"] == ("alice", "bob")
        mock_driver.commit.assert_called_once()

    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_relation_with_properties(self, mock_load, store, mock_driver):
        rel = SimpleNamespace(
            source_id="a", target_id="b", label="LIKES",
            source_label="Person", target_label="Person",
            properties={"weight": 0.9},
        )
        store.upsert_relations([rel])
        call_args = mock_driver.cypher.call_args
        assert "SET" in call_args[0][0]
        assert call_args[1]["params"] == ("a", "b", 0.9)


# ---------------------------------------------------------------------------
# get
# ---------------------------------------------------------------------------


class TestGet:
    @patch("piggie._integrations.llamaindex_property._vertex_to_labelled_node")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_get_by_ids(self, mock_load, mock_convert, store, mock_driver):
        v = Vertex(id=1, label="Person", properties={"name": "Alice", "id": "alice"})
        mock_driver.cypher.return_value = Result([(v,)], ["n"])
        mock_convert.return_value = SimpleNamespace(id_="alice", label="Person")

        nodes = store.get(ids=["alice"])
        assert len(nodes) == 1
        mock_driver.cypher.assert_called_once()
        assert "%s" in mock_driver.cypher.call_args[0][0]

    @patch("piggie._integrations.llamaindex_property._vertex_to_labelled_node")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_get_by_properties(self, mock_load, mock_convert, store, mock_driver):
        v = Vertex(id=1, label="Person", properties={"name": "Alice"})
        mock_driver.cypher.return_value = Result([(v,)], ["n"])
        mock_convert.return_value = SimpleNamespace(id_="1")

        nodes = store.get(properties={"name": "Alice"})
        assert len(nodes) == 1
        call_args = mock_driver.cypher.call_args
        assert "WHERE" in call_args[0][0]
        assert "n.name = %s" in call_args[0][0]


# ---------------------------------------------------------------------------
# get_triplets
# ---------------------------------------------------------------------------


class TestGetTriplets:
    @patch("piggie._integrations.llamaindex_property._Relation")
    @patch("piggie._integrations.llamaindex_property._vertex_to_labelled_node")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_by_entity_names(self, mock_load, mock_convert, MockRelation,
                              store, mock_driver):
        s = Vertex(id=1, label="Person", properties={"name": "Alice"})
        o = Vertex(id=2, label="Person", properties={"name": "Bob"})
        r = Edge(id=10, label="KNOWS", start_id=1, end_id=2, properties={})
        mock_driver.cypher.return_value = Result([(s, r, o)], ["s", "r", "o"])
        mock_convert.side_effect = [
            SimpleNamespace(id_="1"), SimpleNamespace(id_="2"),
        ]
        MockRelation.return_value = SimpleNamespace(label="KNOWS")

        triplets = store.get_triplets(entity_names=["Alice"])
        assert len(triplets) == 1

    @patch("piggie._integrations.llamaindex_property._Relation")
    @patch("piggie._integrations.llamaindex_property._vertex_to_labelled_node")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_by_relation_names(self, mock_load, mock_convert, MockRelation,
                                store, mock_driver):
        s = Vertex(id=1, label="Person", properties={"name": "Alice"})
        o = Vertex(id=2, label="Person", properties={"name": "Bob"})
        r = Edge(id=10, label="KNOWS", start_id=1, end_id=2, properties={})
        mock_driver.cypher.return_value = Result([(s, r, o)], ["s", "r", "o"])
        mock_convert.side_effect = [
            SimpleNamespace(id_="1"), SimpleNamespace(id_="2"),
        ]
        MockRelation.return_value = SimpleNamespace(label="KNOWS")

        triplets = store.get_triplets(relation_names=["KNOWS"])
        assert len(triplets) == 1
        # Should use typed edge query
        call_args = mock_driver.cypher.call_args[0][0]
        assert "KNOWS" in call_args


# ---------------------------------------------------------------------------
# get_rel_map
# ---------------------------------------------------------------------------


class TestGetRelMap:
    @patch("piggie._integrations.llamaindex_property._Relation")
    @patch("piggie._integrations.llamaindex_property._vertex_to_labelled_node")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_basic_traversal(self, mock_load, mock_convert, MockRelation,
                              store, mock_driver):
        s = Vertex(id=1, label="Person", properties={"name": "Alice"})
        o = Vertex(id=2, label="Person", properties={"name": "Bob"})
        e = Edge(id=10, label="KNOWS", start_id=1, end_id=2, properties={})
        mock_driver.cypher.return_value = Result([(s, [e], o)], ["s", "r", "o"])
        mock_convert.side_effect = [
            SimpleNamespace(id_="1"), SimpleNamespace(id_="2"),
        ]
        MockRelation.return_value = SimpleNamespace(label="KNOWS")

        node = SimpleNamespace(id_="1")
        triplets = store.get_rel_map(graph_nodes=[node], depth=2, limit=10)
        assert len(triplets) == 1

    @patch("piggie._integrations.llamaindex_property._Relation")
    @patch("piggie._integrations.llamaindex_property._vertex_to_labelled_node")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_ignore_rels(self, mock_load, mock_convert, MockRelation,
                          store, mock_driver):
        s = Vertex(id=1, label="Person", properties={"name": "Alice"})
        o = Vertex(id=2, label="Person", properties={"name": "Bob"})
        e = Edge(id=10, label="KNOWS", start_id=1, end_id=2, properties={})
        mock_driver.cypher.return_value = Result([(s, [e], o)], ["s", "r", "o"])

        node = SimpleNamespace(id_="1")
        triplets = store.get_rel_map(
            graph_nodes=[node], ignore_rels=["KNOWS"],
        )
        assert len(triplets) == 0  # Filtered out


# ---------------------------------------------------------------------------
# delete
# ---------------------------------------------------------------------------


class TestDelete:
    def test_delete_by_ids(self, store, mock_driver):
        store.delete(ids=["alice", "bob"])
        # Should call cypher DETACH DELETE for each + sql DELETE from embedding table
        assert mock_driver.cypher.call_count == 2
        assert mock_driver.sql.call_count == 2
        mock_driver.commit.assert_called_once()

    def test_delete_by_entity_names(self, store, mock_driver):
        store.delete(entity_names=["Alice"])
        mock_driver.cypher.assert_called_once()
        assert "DETACH DELETE" in mock_driver.cypher.call_args[0][0]
        mock_driver.sql.assert_called_once()

    def test_delete_by_relation_names(self, store, mock_driver):
        store.delete(relation_names=["KNOWS"])
        mock_driver.cypher.assert_called_once()
        assert "DELETE r" in mock_driver.cypher.call_args[0][0]

    def test_delete_by_properties(self, store, mock_driver):
        store.delete(properties={"status": "inactive"})
        mock_driver.cypher.assert_called_once()
        call_args = mock_driver.cypher.call_args
        assert "WHERE" in call_args[0][0]
        assert "DETACH DELETE" in call_args[0][0]


# ---------------------------------------------------------------------------
# structured_query
# ---------------------------------------------------------------------------


class TestStructuredQuery:
    def test_passthrough(self, store, mock_driver):
        mock_driver.cypher.return_value = Result(
            [("Alice", 30)], ["name", "age"]
        )
        result = store.structured_query("MATCH (n) RETURN n.name, n.age")
        assert len(result) == 1
        assert result[0]["name"] == "Alice"

    def test_with_params(self, store, mock_driver):
        mock_driver.cypher.return_value = Result(
            [("Bob",)], ["name"]
        )
        result = store.structured_query(
            "MATCH (n {name: %s}) RETURN n.name",
            param_map={"name": "Bob"},
        )
        assert len(result) == 1
        assert mock_driver.cypher.call_args[1]["params"] == ("Bob",)


# ---------------------------------------------------------------------------
# vector_query
# ---------------------------------------------------------------------------


class TestVectorQuery:
    @patch("piggie._integrations.llamaindex_property._ChunkNode")
    @patch("piggie._integrations.llamaindex_property._EntityNode")
    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_returns_nodes_and_scores(self, mock_load, MockEntity, MockChunk,
                                       store, mock_driver):
        mock_driver.vector_search.return_value = Result(
            [("n1", "Chunk", "hello world", 0.15, '{}')],
            ["node_id", "node_label", "text", "_distance", "properties"],
        )
        MockChunk.return_value = SimpleNamespace(id_="n1", text="hello world")

        query = SimpleNamespace(query_embedding=[0.1, 0.2], similarity_top_k=5)
        nodes, scores = store.vector_query(query)
        assert len(nodes) == 1
        assert len(scores) == 1
        assert scores[0] == pytest.approx(0.15)

    @patch("piggie._integrations.llamaindex_property._load_llamaindex_types")
    def test_no_embedding(self, mock_load, store, mock_driver):
        query = SimpleNamespace(query_embedding=None, similarity_top_k=5)
        nodes, scores = store.vector_query(query)
        assert nodes == []
        assert scores == []


# ---------------------------------------------------------------------------
# Schema
# ---------------------------------------------------------------------------


class TestSchema:
    def test_get_schema(self, store, mock_driver):
        mock_driver.get_schema.return_value = {
            "vertex_labels": [{"name": "Person", "estimated_count": 10}],
            "edge_labels": [{"name": "KNOWS", "estimated_count": 5}],
        }
        schema = store.get_schema()
        assert "vertex_labels" in schema
        assert schema["vertex_labels"][0]["name"] == "Person"

    def test_get_schema_cached(self, store, mock_driver):
        mock_driver.get_schema.return_value = {"vertex_labels": [], "edge_labels": []}
        store.get_schema()
        store.get_schema()
        mock_driver.get_schema.assert_called_once()

    def test_get_schema_refresh(self, store, mock_driver):
        mock_driver.get_schema.return_value = {"vertex_labels": [], "edge_labels": []}
        store.get_schema()
        store.get_schema(refresh=True)
        assert mock_driver.get_schema.call_count == 2

    def test_get_schema_str(self, store, mock_driver):
        mock_driver.get_schema.return_value = {
            "vertex_labels": [{"name": "Person", "estimated_count": 10}],
            "edge_labels": [{"name": "KNOWS", "estimated_count": 5}],
        }
        s = store.get_schema_str()
        assert "Person" in s
        assert "KNOWS" in s
        assert "~10 rows" in s
