"""Data service handles node and edge insert/delete/export operations."""

from __future__ import annotations

from typing import TYPE_CHECKING, Callable, List, Optional

import grpc

if TYPE_CHECKING:
    from ..response import (
        DeleteResult,
        ExportChunk,
        ExportConfig,
        ExportEdgesResult,
        ExportNodesResult,
        InsertEdgesResult,
        InsertNodesResult,
    )
    from ..types import (
        BulkCreateEdgesOptions,
        BulkCreateNodesOptions,
        EdgeData,
        NodeData,
    )
    from .service_context import ServiceContext


class DataService:
    """Data service for inserting, deleting, and exporting graph data."""

    def __init__(self, ctx: ServiceContext):
        self.ctx = ctx
        self._stub = None

    def _get_stub(self):
        if self._stub is None:
            from ..proto import gqldb_pb2_grpc
            channel = self.ctx.get_channel()
            self._stub = gqldb_pb2_grpc.DataServiceStub(channel)
        return self._stub

    def insert_nodes(
        self,
        graph_name: str,
        nodes: List[NodeData],
        options: Optional[BulkCreateNodesOptions],
        bulk_import_session_id: str
    ) -> InsertNodesResult:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..response import InsertNodesResult
        from ..types import TypedValue

        try:
            pb_nodes = []
            for node in nodes:
                pb_props = {}
                for k, v in node.properties.items():
                    tv = TypedValue.from_value(v)
                    pb_props[k] = gqldb_pb2.TypedValue(
                        type=tv.type.value,
                        data=tv.data,
                        is_null=tv.is_null
                    )
                pb_nodes.append(gqldb_pb2.NodeData(
                    labels=node.labels,
                    properties=pb_props
                ))

            request = gqldb_pb2.InsertNodesRequest(
                graph_name=graph_name,
                nodes=pb_nodes,
                bulk_import_session_id=bulk_import_session_id
            )
            # ALWAYS set options (even when None) to avoid deprecated bulk operations error
            request.options.CopyFrom(gqldb_pb2.BulkCreateNodesOptions(
                overwrite=options.overwrite if options else False
            ))

            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().InsertNodes(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            return InsertNodesResult(
                success=response.success,
                node_ids=list(response.node_ids),
                node_count=response.node_count,
                message=response.message
            )
        except grpc.RpcError as e:
            raise GqldbError(f"Insert nodes failed: {e.details()}") from e

    def insert_edges(
        self,
        graph_name: str,
        edges: List[EdgeData],
        options: Optional[BulkCreateEdgesOptions],
        bulk_import_session_id: str
    ) -> InsertEdgesResult:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..response import InsertEdgesResult
        from ..types import TypedValue

        try:
            pb_edges = []
            for edge in edges:
                pb_props = {}
                for k, v in edge.properties.items():
                    tv = TypedValue.from_value(v)
                    pb_props[k] = gqldb_pb2.TypedValue(
                        type=tv.type.value,
                        data=tv.data,
                        is_null=tv.is_null
                    )
                pb_edges.append(gqldb_pb2.EdgeData(
                    label=edge.label,
                    from_node_id=edge.from_node_id,
                    to_node_id=edge.to_node_id,
                    properties=pb_props
                ))

            request = gqldb_pb2.InsertEdgesRequest(
                graph_name=graph_name,
                edges=pb_edges,
                bulk_import_session_id=bulk_import_session_id
            )
            # ALWAYS set options (even when None) to avoid deprecated bulk operations error
            request.options.CopyFrom(gqldb_pb2.BulkCreateEdgesOptions(
                skip_invalid_nodes=options.skip_invalid_nodes if options else False
            ))

            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().InsertEdges(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            return InsertEdgesResult(
                success=response.success,
                edge_ids=list(response.edge_ids),
                edge_count=response.edge_count,
                message=response.message,
                skipped_count=response.skipped_count
            )
        except grpc.RpcError as e:
            raise GqldbError(f"Insert edges failed: {e.details()}") from e

    def delete_nodes(
        self,
        graph_name: str,
        node_ids: Optional[List[str]],
        labels: Optional[List[str]],
        where: str
    ) -> DeleteResult:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..response import DeleteResult

        try:
            request = gqldb_pb2.DeleteNodesRequest(
                graph_name=graph_name,
                node_ids=node_ids or [],
                labels=labels or [],
                where=where
            )
            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().DeleteNodes(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            return DeleteResult(
                success=response.success,
                deleted_count=response.deleted_count,
                message=response.message
            )
        except grpc.RpcError as e:
            raise GqldbError(f"Delete nodes failed: {e.details()}") from e

    def delete_edges(
        self,
        graph_name: str,
        edge_ids: Optional[List[str]],
        label: str,
        where: str
    ) -> DeleteResult:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..response import DeleteResult

        try:
            request = gqldb_pb2.DeleteEdgesRequest(
                graph_name=graph_name,
                edge_ids=edge_ids or [],
                label=label,
                where=where
            )
            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().DeleteEdges(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            return DeleteResult(
                success=response.success,
                deleted_count=response.deleted_count,
                message=response.message
            )
        except grpc.RpcError as e:
            raise GqldbError(f"Delete edges failed: {e.details()}") from e

    def export(
        self,
        config: ExportConfig,
        callback: Optional[Callable[[ExportChunk], None]]
    ) -> None:
        """Export graph data in JSON Lines format (streaming).

        Args:
            config: Export configuration.
            callback: Callback for each exported chunk.

        Raises:
            GqldbError: If operation fails.
        """
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..response import ExportChunk, ExportStats

        try:
            request = gqldb_pb2.ExportRequest(
                graph_name=config.graph_name,
                batch_size=config.batch_size,
                export_nodes=config.export_nodes,
                export_edges=config.export_edges,
                node_labels=config.node_labels or [],
                edge_labels=config.edge_labels or [],
                include_metadata=config.include_metadata
            )
            metadata = self.ctx.get_session_metadata()
            stream = self._get_stub().Export(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )

            for response in stream:
                stats = None
                if response.HasField('stats'):
                    stats = ExportStats(
                        nodes_exported=response.stats.nodes_exported,
                        edges_exported=response.stats.edges_exported,
                        bytes_written=response.stats.bytes_written,
                        duration_ms=response.stats.duration_ms
                    )

                chunk = ExportChunk(
                    data=response.data,
                    is_final=response.is_final,
                    stats=stats
                )

                if callback:
                    callback(chunk)

                if response.is_final:
                    break

        except grpc.RpcError as e:
            raise GqldbError(f"Export failed: {e.details()}") from e

    def export_nodes(
        self,
        graph_name: str,
        labels: Optional[List[str]],
        limit: int,
        callback: Optional[Callable[[ExportNodesResult], None]]
    ) -> None:
        """Stream nodes from a graph.

        Deprecated: Use export() with ExportConfig instead.
        """
        import warnings
        warnings.warn(
            "export_nodes is deprecated, use export() with ExportConfig instead",
            DeprecationWarning,
            stacklevel=2
        )

        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..response import ExportedNode, ExportNodesResult
        from ..types import PropertyType, TypedValue

        try:
            request = gqldb_pb2.ExportNodesRequest(
                graph_name=graph_name,
                labels=labels or [],
                limit=limit
            )
            metadata = self.ctx.get_session_metadata()
            stream = self._get_stub().ExportNodes(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )

            for response in stream:
                nodes = []
                for pb_node in response.nodes:
                    props = {}
                    for k, v in pb_node.properties.items():
                        tv = TypedValue(
                            type=PropertyType(v.type),
                            data=v.data,
                            is_null=v.is_null
                        )
                        props[k] = tv.to_python()
                    nodes.append(ExportedNode(
                        id=pb_node.id,
                        labels=list(pb_node.labels),
                        properties=props
                    ))

                result = ExportNodesResult(nodes=nodes, has_more=response.has_more)
                if callback:
                    callback(result)
        except grpc.RpcError as e:
            raise GqldbError(f"Export nodes failed: {e.details()}") from e

    def export_edges(
        self,
        graph_name: str,
        labels: Optional[List[str]],
        limit: int,
        callback: Optional[Callable[[ExportEdgesResult], None]]
    ) -> None:
        """Stream edges from a graph.

        Deprecated: Use export() with ExportConfig instead.
        """
        import warnings
        warnings.warn(
            "export_edges is deprecated, use export() with ExportConfig instead",
            DeprecationWarning,
            stacklevel=2
        )

        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..response import ExportedEdge, ExportEdgesResult
        from ..types import PropertyType, TypedValue

        try:
            request = gqldb_pb2.ExportEdgesRequest(
                graph_name=graph_name,
                labels=labels or [],
                limit=limit
            )
            metadata = self.ctx.get_session_metadata()
            stream = self._get_stub().ExportEdges(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )

            for response in stream:
                edges = []
                for pb_edge in response.edges:
                    props = {}
                    for k, v in pb_edge.properties.items():
                        tv = TypedValue(
                            type=PropertyType(v.type),
                            data=v.data,
                            is_null=v.is_null
                        )
                        props[k] = tv.to_python()
                    edges.append(ExportedEdge(
                        id=pb_edge.id,
                        label=pb_edge.label,
                        from_node_id=pb_edge.from_node_id,
                        to_node_id=pb_edge.to_node_id,
                        properties=props
                    ))

                result = ExportEdgesResult(edges=edges, has_more=response.has_more)
                if callback:
                    callback(result)
        except grpc.RpcError as e:
            raise GqldbError(f"Export edges failed: {e.details()}") from e
