"""Bulk import service handles bulk import session management."""

from __future__ import annotations

from typing import TYPE_CHECKING

import grpc

if TYPE_CHECKING:
    from ..types import (
        AbortBulkImportResult,
        BulkImportSession,
        BulkImportStatus,
        CheckpointResult,
        EndBulkImportResult,
    )
    from .service_context import ServiceContext


class BulkImportService:
    """Bulk import service for managing bulk import sessions."""

    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.BulkImportServiceStub(channel)
        return self._stub

    def start_bulk_import(
        self,
        graph_name: str,
        checkpoint_every: int,
        estimated_nodes: int,
        estimated_edges: int,
        memtable_size: int,
        max_memtables: int
    ) -> BulkImportSession:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..types import BulkImportSession

        try:
            request = gqldb_pb2.StartBulkImportRequest(
                graph_name=graph_name,
                checkpoint_every=checkpoint_every,
                estimated_nodes=estimated_nodes,
                estimated_edges=estimated_edges,
                memtable_size=memtable_size,
                max_memtables=max_memtables
            )
            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().StartBulkImport(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            if not response.success:
                raise GqldbError(f"Start bulk import failed: {response.message}")
            return BulkImportSession(
                success=response.success,
                session_id=response.session_id,
                message=response.message
            )
        except grpc.RpcError as e:
            raise GqldbError(f"Start bulk import failed: {e.details()}") from e

    def checkpoint(self, session_id: str) -> CheckpointResult:
        """Deprecated: Checkpoint is no longer needed."""
        from ..types import CheckpointResult
        return CheckpointResult(
            success=True,
            record_count=0,
            last_checkpoint_count=0,
            message="Checkpoint has been removed; use end_bulk_import which performs a final flush",
        )

    def end_bulk_import(self, session_id: str) -> EndBulkImportResult:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..types import EndBulkImportResult

        try:
            request = gqldb_pb2.EndBulkImportRequest(session_id=session_id)
            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().EndBulkImport(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            return EndBulkImportResult(
                success=response.success,
                total_records=response.total_records,
                message=response.message
            )
        except grpc.RpcError as e:
            raise GqldbError(f"End bulk import failed: {e.details()}") from e

    def abort_bulk_import(self, session_id: str) -> AbortBulkImportResult:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..types import AbortBulkImportResult

        try:
            request = gqldb_pb2.AbortBulkImportRequest(session_id=session_id)
            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().AbortBulkImport(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            return AbortBulkImportResult(
                success=response.success,
                message=response.message
            )
        except grpc.RpcError as e:
            raise GqldbError(f"Abort bulk import failed: {e.details()}") from e

    def get_bulk_import_status(self, session_id: str) -> BulkImportStatus:
        from ..errors import GqldbError
        from ..proto import gqldb_pb2
        from ..types import BulkImportStatus

        try:
            request = gqldb_pb2.GetBulkImportStatusRequest(session_id=session_id)
            metadata = self.ctx.get_session_metadata()
            response = self._get_stub().GetBulkImportStatus(
                request, metadata=metadata, timeout=self.ctx.config.timeout
            )
            return BulkImportStatus(
                is_active=response.is_active,
                graph_name=response.graph_name,
                record_count=response.record_count,
                last_checkpoint_count=response.last_checkpoint_count,
                created_at=response.created_at,
                last_activity=response.last_activity
            )
        except grpc.RpcError as e:
            raise GqldbError(f"Get bulk import status failed: {e.details()}") from e
