"""Tests for PyTorch Geometric integration — mocked where possible."""

from __future__ import annotations

from unittest.mock import MagicMock

import pytest

from piggie._integrations.pyg import (
    _HAS_PYG,
    EdgeAttr,
    PiggieFeatureStore,
    PiggieGraphStore,
    TensorAttr,
)
from piggie.result import Result

# Skip all tests if PyG is not installed — these are still valuable as
# structural tests that validate the code paths with mocked drivers.
pytestmark = pytest.mark.skipif(not _HAS_PYG, reason="PyTorch Geometric not installed")


@pytest.fixture()
def mock_driver():
    driver = MagicMock()
    driver.graph = "test_graph"
    # Make sql() and commit() no-ops by default
    driver.sql.return_value = Result([], [])
    return driver


# ---------------------------------------------------------------------------
# FeatureStore
# ---------------------------------------------------------------------------


class TestPiggieFeatureStore:
    def test_put_and_get_tensor(self, mock_driver):
        import torch

        fs = PiggieFeatureStore(mock_driver, graph="test_graph")
        tensor = torch.randn(3, 4)
        attr = TensorAttr(group_name="Person", attr_name="x", index=0)

        result = fs._put_tensor(tensor, attr)
        assert result is True
        mock_driver.sql.assert_called()
        mock_driver.commit.assert_called()

    def test_remove_tensor(self, mock_driver):
        fs = PiggieFeatureStore(mock_driver, graph="test_graph")
        attr = TensorAttr(group_name="Person", attr_name="x", index=0)

        result = fs._remove_tensor(attr)
        assert result is True

    def test_get_all_tensor_attrs(self, mock_driver):
        mock_driver.sql.side_effect = [
            Result([], []),  # CREATE TABLE
            Result([("Person", "x"), ("Person", "y")], ["node_type", "feature_name"]),
        ]
        fs = PiggieFeatureStore(mock_driver, graph="test_graph")
        attrs = fs.get_all_tensor_attrs()
        assert len(attrs) == 2
        assert attrs[0].group_name == "Person"
        assert attrs[0].attr_name == "x"


# ---------------------------------------------------------------------------
# GraphStore
# ---------------------------------------------------------------------------


class TestPiggieGraphStore:
    def test_build_node_id_map(self, mock_driver):
        mock_driver.cypher.return_value = Result(
            [(100,), (200,), (300,)], ["nid"]
        )
        gs = PiggieGraphStore(mock_driver, graph="test_graph")
        id_map = gs._build_node_id_map("Person")
        assert id_map == {100: 0, 200: 1, 300: 2}

    def test_get_edge_index(self, mock_driver):

        # First call: build source node map
        # Second call: build dest node map
        # Third call: get edges
        mock_driver.cypher.side_effect = [
            Result([(100,), (200,)], ["nid"]),  # src map
            Result([(300,), (400,)], ["nid"]),  # dst map
            Result([(100, 300), (200, 400)], ["src", "dst"]),  # edges
        ]
        gs = PiggieGraphStore(mock_driver, graph="test_graph")
        edge_attr = EdgeAttr(edge_type=("Person", "KNOWS", "Person"))
        edge_index = gs._get_edge_index(edge_attr)
        assert edge_index is not None
        assert edge_index.shape == (2, 2)
        assert edge_index[0].tolist() == [0, 1]  # remapped src IDs
        assert edge_index[1].tolist() == [0, 1]  # remapped dst IDs

    def test_get_edge_index_same_label(self, mock_driver):

        # Same label for src and dst — node map called once, cached
        mock_driver.cypher.side_effect = [
            Result([(100,), (200,), (300,)], ["nid"]),  # node map (cached for both)
            Result([(100, 200), (200, 300)], ["src", "dst"]),  # edges
        ]
        gs = PiggieGraphStore(mock_driver, graph="test_graph")
        edge_attr = EdgeAttr(edge_type=("Person", "KNOWS", "Person"))
        edge_index = gs._get_edge_index(edge_attr)
        assert edge_index is not None
        assert edge_index.shape == (2, 2)

    def test_put_edge_index_noop(self, mock_driver):
        gs = PiggieGraphStore(mock_driver, graph="test_graph")
        assert gs._put_edge_index(None, EdgeAttr()) is True

    def test_remove_edge_index_noop(self, mock_driver):
        gs = PiggieGraphStore(mock_driver, graph="test_graph")
        assert gs._remove_edge_index(EdgeAttr()) is True

    def test_invalid_edge_type(self, mock_driver):
        from piggie.exceptions import IntegrationError
        gs = PiggieGraphStore(mock_driver, graph="test_graph")
        with pytest.raises(IntegrationError, match="edge_type must be"):
            gs._get_edge_index(EdgeAttr(edge_type=None))

    def test_get_all_edge_attrs(self, mock_driver):
        mock_driver.get_schema.return_value = {
            "vertex_labels": [{"name": "Person"}],
            "edge_labels": [{"name": "KNOWS"}],
        }
        mock_driver.cypher.return_value = Result(
            [("Person", "Person")], ["src_label", "dst_label"]
        )
        gs = PiggieGraphStore(mock_driver, graph="test_graph")
        attrs = gs.get_all_edge_attrs()
        assert len(attrs) >= 1
