"""Data Service integration tests."""

import pytest
import uuid

from gqldb.client import GqldbClient, QueryConfig
from gqldb.response import ExportConfig, ExportChunk
from gqldb.types import GraphType, NodeData, EdgeData
from tests.conftest import drop_test_graph


@pytest.mark.integration
class TestData:
    """Tests for data service (insert/export nodes and edges)."""

    @pytest.fixture
    def test_graph(self, connected_client: GqldbClient, request):
        """Create a temporary test graph."""
        graph_name = f"test_{request.node.name}_{uuid.uuid4().hex[:8]}"
        connected_client.create_graph(graph_name, GraphType.OPEN, "Test data graph")
        yield graph_name
        drop_test_graph(connected_client, graph_name)

    def test_insert_nodes(self, connected_client: GqldbClient, test_graph: str):
        """Test inserting nodes."""
        # Start bulk import session
        session = connected_client.start_bulk_import(test_graph)

        try:
            nodes = [
                NodeData(labels=["Person"], properties={"name": "Alice", "age": 30}),
                NodeData(labels=["Person"], properties={"name": "Bob", "age": 25}),
            ]

            result = connected_client.insert_nodes(
                test_graph,
                nodes,
                bulk_import_session_id=session.session_id
            )

            assert result.success
            assert result.node_count == 2
            print(f"Inserted {result.node_count} nodes")
        finally:
            connected_client.end_bulk_import(session.session_id)

    def test_insert_edges(self, connected_client: GqldbClient, test_graph: str):
        """Test inserting edges."""
        # Start bulk import session
        session = connected_client.start_bulk_import(test_graph)

        try:
            # First insert nodes
            nodes = [
                NodeData(labels=["Person"], properties={"name": "Alice"}),
                NodeData(labels=["Person"], properties={"name": "Bob"}),
            ]
            connected_client.insert_nodes(
                test_graph,
                nodes,
                bulk_import_session_id=session.session_id
            )

            # Query to get node IDs (bulk import doesn't return IDs immediately)
            query_config = QueryConfig(graph_name=test_graph)
            resp = connected_client.gql("MATCH (n:Person) RETURN id(n) AS node_id ORDER BY n.name", query_config)

            node_ids = []
            for i in range(resp.row_count):
                row = resp.rows[i]
                value = resp.get_by_name(row, "node_id")
                if value:
                    node_ids.append(value)

            # Then insert edges
            if len(node_ids) >= 2:
                edges = [
                    EdgeData(
                        label="KNOWS",
                        from_node_id=node_ids[0],
                        to_node_id=node_ids[1],
                        properties={"since": 2020}
                    ),
                ]
                result = connected_client.insert_edges(
                    test_graph,
                    edges,
                    bulk_import_session_id=session.session_id
                )

                assert result.success
                print(f"Inserted {result.edge_count} edges")
        finally:
            connected_client.end_bulk_import(session.session_id)

    def test_export(self, connected_client: GqldbClient, test_graph: str):
        """Test unified export with ExportConfig."""
        # Start bulk import session
        session = connected_client.start_bulk_import(test_graph)

        try:
            # First insert some nodes
            nodes = [
                NodeData(labels=["Person"], properties={"name": "Charlie"}),
                NodeData(labels=["Person"], properties={"name": "Diana"}),
            ]
            connected_client.insert_nodes(
                test_graph,
                nodes,
                bulk_import_session_id=session.session_id
            )

        finally:
            connected_client.end_bulk_import(session.session_id)

        # Export using new unified API
        chunks = []
        all_data = []

        def callback(chunk: ExportChunk):
            chunks.append(chunk)
            all_data.append(chunk.get_data_as_string())
            if chunk.is_final and chunk.has_stats():
                print(f"Export stats: nodes={chunk.stats.nodes_exported}, edges={chunk.stats.edges_exported}")

        config = ExportConfig(
            graph_name=test_graph,
            export_nodes=True,
            export_edges=True,
            batch_size=1000
        )
        connected_client.export(config, callback)
        print(f"Total chunks: {len(chunks)}, data length: {sum(len(d) for d in all_data)} bytes")

    def test_export_nodes_only(self, connected_client: GqldbClient, test_graph: str):
        """Test export with nodes only."""
        # Start bulk import session
        session = connected_client.start_bulk_import(test_graph)

        try:
            nodes = [
                NodeData(labels=["Person"], properties={"name": "ExportNode1"}),
            ]
            connected_client.insert_nodes(
                test_graph,
                nodes,
                bulk_import_session_id=session.session_id
            )
        finally:
            connected_client.end_bulk_import(session.session_id)

        chunks = []

        def callback(chunk: ExportChunk):
            chunks.append(chunk)

        config = ExportConfig(
            graph_name=test_graph,
            export_nodes=True,
            export_edges=False
        )
        connected_client.export(config, callback)
        print(f"Export nodes-only: {len(chunks)} chunks")

    def test_export_edges_only(self, connected_client: GqldbClient, test_graph: str):
        """Test export with edges only."""
        # Start bulk import session
        session = connected_client.start_bulk_import(test_graph)

        try:
            nodes = [
                NodeData(labels=["Person"], properties={"name": "EdgeNode1"}),
                NodeData(labels=["Person"], properties={"name": "EdgeNode2"}),
            ]
            connected_client.insert_nodes(
                test_graph,
                nodes,
                bulk_import_session_id=session.session_id
            )
            # Query to get node IDs
            query_config = QueryConfig(graph_name=test_graph)
            resp = connected_client.gql("MATCH (n:Person) WHERE n.name STARTS WITH 'EdgeNode' RETURN id(n) AS node_id ORDER BY n.name", query_config)

            node_ids = []
            for i in range(resp.row_count):
                row = resp.rows[i]
                value = resp.get_by_name(row, "node_id")
                if value:
                    node_ids.append(value)

            if len(node_ids) >= 2:
                edges = [
                    EdgeData(
                        label="EDGE_ONLY_TEST",
                        from_node_id=node_ids[0],
                        to_node_id=node_ids[1],
                        properties={}
                    ),
                ]
                connected_client.insert_edges(
                    test_graph,
                    edges,
                    bulk_import_session_id=session.session_id
                )
        finally:
            connected_client.end_bulk_import(session.session_id)

        chunks = []

        def callback(chunk: ExportChunk):
            chunks.append(chunk)

        config = ExportConfig(
            graph_name=test_graph,
            export_nodes=False,
            export_edges=True
        )
        connected_client.export(config, callback)
        print(f"Export edges-only: {len(chunks)} chunks")

    def test_delete_nodes(self, connected_client: GqldbClient, test_graph: str):
        """Test deleting nodes by ID."""
        # Start bulk import session
        session = connected_client.start_bulk_import(test_graph)

        try:
            # Insert nodes first
            nodes = [
                NodeData(labels=["ToDelete"], properties={"name": "ToDelete1"}),
                NodeData(labels=["ToDelete"], properties={"name": "ToDelete2"}),
            ]
            connected_client.insert_nodes(
                test_graph,
                nodes,
                bulk_import_session_id=session.session_id
            )

        finally:
            connected_client.end_bulk_import(session.session_id)

        # Query to get node IDs (bulk import doesn't return IDs immediately)
        query_config = QueryConfig(graph_name=test_graph)
        resp = connected_client.gql("MATCH (n:ToDelete) RETURN id(n) AS node_id ORDER BY n.name", query_config)

        node_ids = []
        for i in range(resp.row_count):
            row = resp.rows[i]
            value = resp.get_by_name(row, "node_id")
            if value:
                node_ids.append(value)

        assert len(node_ids) == 2
        print(f"Inserted nodes: {node_ids}")

        # Delete nodes by ID
        result = connected_client.delete_nodes(
            test_graph,
            node_ids=node_ids
        )
        assert result.success
        print(f"Deleted {result.deleted_count} nodes")

    def test_delete_edges(self, connected_client: GqldbClient, test_graph: str):
        """Test deleting edges by ID."""
        # Start bulk import session
        session = connected_client.start_bulk_import(test_graph)

        try:
            # Insert nodes first
            nodes = [
                NodeData(labels=["Person"], properties={"name": "Source"}),
                NodeData(labels=["Person"], properties={"name": "Target"}),
            ]
            connected_client.insert_nodes(
                test_graph,
                nodes,
                bulk_import_session_id=session.session_id
            )

            # Query to get node IDs
            query_config = QueryConfig(graph_name=test_graph)
            resp = connected_client.gql("MATCH (n:Person) RETURN id(n) AS node_id ORDER BY n.name", query_config)

            node_ids = []
            for i in range(resp.row_count):
                row = resp.rows[i]
                value = resp.get_by_name(row, "node_id")
                if value:
                    node_ids.append(value)

            if len(node_ids) >= 2:
                # Insert edge
                edges = [
                    EdgeData(
                        label="TO_DELETE",
                        from_node_id=node_ids[0],
                        to_node_id=node_ids[1],
                        properties={"temp": True}
                    ),
                ]
                connected_client.insert_edges(
                    test_graph,
                    edges,
                    bulk_import_session_id=session.session_id
                )

        finally:
            connected_client.end_bulk_import(session.session_id)

        # Query to get edge IDs
        edge_resp = connected_client.gql("MATCH ()-[e:TO_DELETE]->() RETURN id(e) AS edge_id", query_config)

        edge_ids = []
        for i in range(edge_resp.row_count):
            row = edge_resp.rows[i]
            value = edge_resp.get_by_name(row, "edge_id")
            if value:
                edge_ids.append(value)

        print(f"Inserted edge: {edge_ids}")

        # Delete edge by ID
        if edge_ids:
            result = connected_client.delete_edges(
                test_graph,
                edge_ids=edge_ids
            )
            assert result.success
            print(f"Deleted {result.deleted_count} edges")
