"""Integration tests for LlamaIndex PropertyGraphStore — requires live PG18+AGE."""

from __future__ import annotations

import uuid

import pytest

pytestmark = pytest.mark.integration

try:
    from llama_index.core.graph_stores.types import ChunkNode, EntityNode, Relation

    _HAS_LLAMAINDEX = True
except ImportError:
    _HAS_LLAMAINDEX = False

import piggie  # noqa: E402
from piggie._integrations.llamaindex_property import PiggiePropertyGraphStore  # noqa: E402

DSN = "host=/var/run/postgresql port=5433 dbname=piggie user=gregf"


@pytest.fixture()
def driver():
    graph_name = f"test_li_{uuid.uuid4().hex[:8]}"
    db = piggie.connect(DSN, graph=graph_name)
    yield db
    try:
        db.drop_graph(graph_name)
        db.commit()
    except Exception:
        db.rollback()
    db.close()


@pytest.fixture()
def store(driver):
    return PiggiePropertyGraphStore(
        driver,
        graph=driver.graph,
        embedding_dimensions=3,
        create_tables=True,
    )


@pytest.fixture()
def embedding_table(store):
    return store._embedding_table


class TestPropertyGraphStoreIntegration:
    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_upsert_and_get_entity_nodes(self, store):
        nodes = [
            EntityNode(id_="alice", label="Person", name="Alice",
                       properties={"age": 30}),
            EntityNode(id_="bob", label="Person", name="Bob",
                       properties={"age": 25}),
        ]
        store.upsert_nodes(nodes)

        result = store.get(ids=["alice"])
        assert len(result) == 1
        assert result[0].id_ == "alice"

    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_upsert_and_get_chunk_nodes(self, store):
        nodes = [
            ChunkNode(id_="c1", label="Chunk", text="Hello world"),
        ]
        store.upsert_nodes(nodes)

        result = store.get(ids=["c1"])
        assert len(result) == 1

    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_upsert_nodes_with_embeddings(self, store, driver, embedding_table):
        nodes = [
            EntityNode(id_="vec1", label="Doc", name="Doc 1",
                       properties={}, embedding=[0.1, 0.2, 0.3]),
        ]
        store.upsert_nodes(nodes)

        # Verify embedding was stored in companion table
        result = driver.sql(
            f"SELECT node_id, node_label FROM {embedding_table} WHERE node_id = %s",
            ("vec1",),
        )
        rows = result.to_dicts()
        assert len(rows) == 1
        assert rows[0]["node_id"] == "vec1"

    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_upsert_relations(self, store):
        nodes = [
            EntityNode(id_="alice", label="Person", name="Alice"),
            EntityNode(id_="bob", label="Person", name="Bob"),
        ]
        store.upsert_nodes(nodes)

        rels = [
            Relation(
                source_id="alice", target_id="bob",
                label="KNOWS",
                source_label="Person", target_label="Person",
            ),
        ]
        store.upsert_relations(rels)

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

    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_get_triplets_by_relation(self, store):
        nodes = [
            EntityNode(id_="a", label="Person", name="A"),
            EntityNode(id_="b", label="Person", name="B"),
        ]
        store.upsert_nodes(nodes)
        store.upsert_relations([
            Relation(
                source_id="a", target_id="b", label="FOLLOWS",
                source_label="Person", target_label="Person",
            ),
        ])

        triplets = store.get_triplets(relation_names=["FOLLOWS"])
        assert len(triplets) >= 1
        assert triplets[0][1].label == "FOLLOWS"

    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_structured_query(self, store):
        nodes = [
            EntityNode(id_="x", label="Thing", name="X"),
        ]
        store.upsert_nodes(nodes)

        result = store.structured_query(
            "MATCH (n:Thing {id: %s}) RETURN n.name",
            param_map={"id": "x"},
        )
        assert len(result) == 1

    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_vector_query(self, store):
        from llama_index.core.vector_stores.types import VectorStoreQuery

        nodes = [
            EntityNode(id_="v1", label="Doc", name="Doc A",
                       embedding=[1.0, 0.0, 0.0]),
            EntityNode(id_="v2", label="Doc", name="Doc B",
                       embedding=[0.0, 1.0, 0.0]),
        ]
        store.upsert_nodes(nodes)

        query = VectorStoreQuery(query_embedding=[1.0, 0.0, 0.0], similarity_top_k=2)
        found_nodes, scores = store.vector_query(query)
        assert len(found_nodes) >= 1
        assert len(scores) >= 1

    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_delete_removes_from_both(self, store, driver, embedding_table):
        nodes = [
            EntityNode(id_="del1", label="Person", name="DeleteMe",
                       embedding=[0.5, 0.5, 0.5]),
        ]
        store.upsert_nodes(nodes)

        # Verify exists
        assert len(store.get(ids=["del1"])) == 1

        # Delete
        store.delete(ids=["del1"])

        # Verify removed from graph
        assert len(store.get(ids=["del1"])) == 0

        # Verify removed from embedding table
        result = driver.sql(
            f"SELECT 1 FROM {embedding_table} WHERE node_id = %s",
            ("del1",),
        )
        assert len(result.to_dicts()) == 0

    @pytest.mark.skipif(not _HAS_LLAMAINDEX, reason="llama-index-core not installed")
    def test_get_schema(self, store):
        nodes = [EntityNode(id_="s1", label="Widget", name="W1")]
        store.upsert_nodes(nodes)

        schema = store.get_schema(refresh=True)
        assert "vertex_labels" in schema

        schema_str = store.get_schema_str(refresh=True)
        assert "Widget" in schema_str
