"""Main client for GQLDB Python driver."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

import grpc

from .config import GqldbConfig
from .connection import ConnectionPool
from .errors import EmptyQueryError, GqldbError, GraphNotFoundError
from .response import (
    DeleteResult,
    ExportChunk,
    ExportConfig,
    ExportedEdge,
    ExportedNode,
    ExportEdgesResult,
    ExportNodesResult,
    ExportStats,
    InsertEdgesResult,
    InsertNodesResult,
    Response,
    Row,
)
from .session import Session, SessionManager
from .transaction import Transaction, TransactionManager
from .types import (
    AbortBulkImportResult,
    ASTCacheStats,
    BulkCreateEdgesOptions,
    BulkCreateNodesOptions,
    BulkImportSession,
    BulkImportStatus,
    CacheStats,
    CacheType,
    CheckpointResult,
    CompactResult,
    EdgeData,
    EndBulkImportResult,
    GraphInfo,
    GraphType,
    HealthStatus,
    NodeData,
    PlanCacheStats,
    PropertyType,
    Statistics,
    TransactionInfo,
    TypedValue,
)


@dataclass
class QueryConfig:
    """Configuration for a query."""
    graph_name: str = ""
    parameters: Dict[str, Any] = field(default_factory=dict)
    transaction_id: int = 0
    timeout: int = 0
    read_only: bool = False
    max_path_results: int = 0  # Max paths to return from path queries (0 = unlimited)


class GqldbClient:
    """Main client for interacting with GQLDB."""

    def __init__(self, config: GqldbConfig):
        """Initialize the GQLDB client.

        Args:
            config: Configuration for the client.

        Raises:
            GqldbError: If initialization fails.
        """
        config.validate()
        self._config = config
        self._pool = ConnectionPool(config)
        self._sessions = SessionManager()
        self._tx_manager = TransactionManager()

        # Initialize service context
        from .services import (
            ServiceContext,
            SessionService,
            QueryService,
            GraphService,
            TransactionService,
            DataService,
            HealthService,
            AdminService,
            BulkImportService,
        )

        ctx = ServiceContext(config, self._pool, self._sessions, self._tx_manager)

        # Initialize service instances
        self._session_service = SessionService(ctx)
        self._query_service = QueryService(ctx)
        self._graph_service = GraphService(ctx)
        self._transaction_service = TransactionService(ctx)
        self._data_service = DataService(ctx)
        self._health_service = HealthService(ctx)
        self._admin_service = AdminService(ctx)
        self._bulk_import_service = BulkImportService(ctx)

    def close(self) -> None:
        """Close the client and all connections."""
        if self._sessions.is_logged_in():
            try:
                self.logout()
            except Exception:
                pass
        self._pool.close()

    def __enter__(self) -> "GqldbClient":
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Context manager exit."""
        self.close()

    # =========================================================================
    # Session Service
    # =========================================================================

    def login(self, username: str, password: str) -> Session:
        """Authenticate the user and create a session."""
        return self._session_service.login(username, password, self._config.default_graph)

    def logout(self) -> None:
        """Terminate the current session."""
        return self._session_service.logout()

    def ping(self) -> int:
        """Check the connection and return the latency in nanoseconds."""
        return self._session_service.ping()

    # =========================================================================
    # Query Service
    # =========================================================================

    def gql(self, query: str, config: Optional[QueryConfig] = None) -> Response:
        """Execute a GQL query and return the result.

        Args:
            query: GQL query string.
            config: Optional query configuration.

        Returns:
            Query response.

        Raises:
            EmptyQueryError: If query is empty.
            GqldbError: If query fails.
        """
        return self._query_service.gql(query, config)
    def gql_stream(
        self,
        query: str,
        config: Optional[QueryConfig] = None,
        callback: Optional[Callable[[Response], None]] = None
    ) -> None:
        """Execute a GQL query and stream the results.

        Args:
            query: GQL query string.
            config: Optional query configuration.
            callback: Callback function for each response chunk.

        Raises:
            EmptyQueryError: If query is empty.
            GqldbError: If query fails.
        """
        return self._query_service.gql_stream(query, config, callback)

    def explain(self, query: str, config: Optional[QueryConfig] = None) -> str:
        """Return the execution plan for a query.

        Args:
            query: GQL query string.
            config: Optional query configuration.

        Returns:
            Execution plan string.

        Raises:
            GqldbError: If explain fails.
        """
        return self._query_service.explain(query, config)
    def profile(self, query: str, config: Optional[QueryConfig] = None) -> str:
        """Execute a query with profiling and return statistics.

        Args:
            query: GQL query string.
            config: Optional query configuration.

        Returns:
            Profile statistics string.

        Raises:
            GqldbError: If profile fails.
        """
        return self._query_service.profile(query, config)
    # =========================================================================
    # Graph Service
    # =========================================================================

    def create_graph(
        self,
        name: str,
        graph_type: GraphType = GraphType.OPEN,
        description: str = ""
    ) -> None:
        """Create a new graph.

        Args:
            name: Graph name.
            graph_type: Type of graph.
            description: Optional description.

        Raises:
            GqldbError: If creation fails.
        """
        return self._graph_service.create_graph(name, graph_type, description)
    def drop_graph(self, name: str, if_exists: bool = False) -> None:
        """Delete a graph.

        Args:
            name: Graph name.
            if_exists: If True, don't error if graph doesn't exist.

        Raises:
            GqldbError: If deletion fails.
        """
        return self._graph_service.drop_graph(name, if_exists)
    def use_graph(self, name: str) -> None:
        """Set the current graph for the session.

        Args:
            name: Graph name.

        Raises:
            GqldbError: If operation fails.
        """
        return self._graph_service.use_graph(name)
    def list_graphs(self) -> List[GraphInfo]:
        """Return all available graphs.

        Returns:
            List of GraphInfo objects.

        Raises:
            GqldbError: If operation fails.
        """
        return self._graph_service.list_graphs()
    def get_graph_info(self, name: str) -> GraphInfo:
        """Return information about a specific graph.

        Args:
            name: Graph name.

        Returns:
            GraphInfo object.

        Raises:
            GraphNotFoundError: If graph not found.
            GqldbError: If operation fails.
        """
        return self._graph_service.get_graph_info(name)
    # =========================================================================
    # Transaction Service
    # =========================================================================

    def begin_transaction(
        self,
        graph_name: str,
        read_only: bool = False,
        timeout: int = 0
    ) -> Transaction:
        """Start a new transaction.

        Args:
            graph_name: Graph name for the transaction.
            read_only: Whether the transaction is read-only.
            timeout: Timeout in seconds.

        Returns:
            The created Transaction.

        Raises:
            GqldbError: If operation fails.
        """
        return self._transaction_service.begin_transaction(graph_name, read_only, timeout)
    def commit(self, transaction_id: int) -> bool:
        """Commit a transaction.

        Args:
            transaction_id: Transaction ID.

        Returns:
            True if successful.

        Raises:
            GqldbError: If operation fails.
        """
        return self._transaction_service.commit(transaction_id)
    def rollback(self, transaction_id: int) -> bool:
        """Rollback a transaction.

        Args:
            transaction_id: Transaction ID.

        Returns:
            True if successful.

        Raises:
            GqldbError: If operation fails.
        """
        return self._transaction_service.rollback(transaction_id)
    def list_transactions(self) -> List[TransactionInfo]:
        """Return active transactions.

        Returns:
            List of TransactionInfo objects.

        Raises:
            GqldbError: If operation fails.
        """
        return self._transaction_service.list_transactions()
    def with_transaction(
        self,
        graph_name: str,
        fn: Callable[[int], None],
        read_only: bool = False,
        timeout: Optional[int] = None
    ) -> None:
        """Execute a function within a transaction.

        Automatically commits on success or rolls back on error.

        Args:
            graph_name: Graph name for the transaction.
            fn: Function to execute with transaction ID.
            read_only: Whether the transaction is read-only.
            timeout: Optional timeout in seconds. Uses config timeout if not specified.

        Raises:
            GqldbError: If operation fails.
        """
        return self._transaction_service.with_transaction(
            graph_name, fn, read_only, timeout or self.config.timeout
        )
    # =========================================================================
    # Data Service
    # =========================================================================

    def insert_nodes(
        self,
        graph_name: str,
        nodes: List[NodeData],
        options: Optional[BulkCreateNodesOptions] = None,
        bulk_import_session_id: str = ""
    ) -> InsertNodesResult:
        """Insert multiple nodes into a graph.

        Args:
            graph_name: Graph name.
            nodes: List of node data.
            options: Optional bulk creation options.
            bulk_import_session_id: Optional bulk import session ID for auto-checkpoint.

        Returns:
            InsertNodesResult.

        Raises:
            GqldbError: If operation fails.
        """
        return self._data_service.insert_nodes(graph_name, nodes, options, bulk_import_session_id)
    def insert_edges(
        self,
        graph_name: str,
        edges: List[EdgeData],
        options: Optional[BulkCreateEdgesOptions] = None,
        bulk_import_session_id: str = ""
    ) -> InsertEdgesResult:
        """Insert multiple edges into a graph.

        Args:
            graph_name: Graph name.
            edges: List of edge data.
            options: Optional bulk creation options.
            bulk_import_session_id: Optional bulk import session ID for auto-checkpoint.

        Returns:
            InsertEdgesResult.

        Raises:
            GqldbError: If operation fails.
        """
        return self._data_service.insert_edges(graph_name, edges, options, bulk_import_session_id)
    def delete_nodes(
        self,
        graph_name: str,
        node_ids: Optional[List[str]] = None,
        labels: Optional[List[str]] = None,
        where: str = ""
    ) -> DeleteResult:
        """Delete nodes from a graph.

        Args:
            graph_name: Graph name.
            node_ids: Optional list of node IDs to delete.
            labels: Optional list of labels to filter by.
            where: Optional WHERE clause.

        Returns:
            DeleteResult.

        Raises:
            GqldbError: If operation fails.
        """
        return self._data_service.delete_nodes(graph_name, node_ids, labels, where)
    def delete_edges(
        self,
        graph_name: str,
        edge_ids: Optional[List[str]] = None,
        label: str = "",
        where: str = ""
    ) -> DeleteResult:
        """Delete edges from a graph.

        Args:
            graph_name: Graph name.
            edge_ids: Optional list of edge IDs to delete.
            label: Optional label to filter by.
            where: Optional WHERE clause.

        Returns:
            DeleteResult.

        Raises:
            GqldbError: If operation fails.
        """
        return self._data_service.delete_edges(graph_name, edge_ids, label, where)
    def export(
        self,
        config: ExportConfig,
        callback: Optional[Callable[[ExportChunk], None]] = 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.
        """
        return self._data_service.export(config, callback)

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

        Deprecated: Use export() with ExportConfig instead.

        Args:
            graph_name: Graph name.
            labels: Optional list of labels to filter by.
            limit: Maximum number of nodes (0 = no limit).
            callback: Callback for each batch.

        Raises:
            GqldbError: If operation fails.
        """
        return self._data_service.export_nodes(graph_name, labels, limit, callback)

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

        Deprecated: Use export() with ExportConfig instead.

        Args:
            graph_name: Graph name.
            labels: Optional list of labels to filter by.
            limit: Maximum number of edges (0 = no limit).
            callback: Callback for each batch.

        Raises:
            GqldbError: If operation fails.
        """
        return self._data_service.export_edges(graph_name, labels, limit, callback)
    # =========================================================================
    # Health Service
    # =========================================================================

    def health_check(self, service: str = "") -> HealthStatus:
        """Check the health of a service.

        Args:
            service: Service name (empty for overall health).

        Returns:
            HealthStatus.

        Raises:
            GqldbError: If operation fails.
        """
        return self._health_service.health_check(service)
    def _proto_to_health_status(self, proto_status: int) -> HealthStatus:
        """Convert proto health status to HealthStatus enum."""
        # Proto enum values: HEALTH_STATUS_UNKNOWN=0, SERVING=1, NOT_SERVING=2, SERVICE_UNKNOWN=3
        if proto_status == 1:
            return HealthStatus.SERVING
        elif proto_status == 2:
            return HealthStatus.NOT_SERVING
        elif proto_status == 3:
            return HealthStatus.SERVICE_UNKNOWN
        return HealthStatus.UNKNOWN

    def watch(
        self,
        service: str = "",
        callback: Optional[Callable[[HealthStatus], None]] = None
    ) -> None:
        """Watch health status changes via server-side streaming.

        Args:
            service: Service name to watch (empty for overall health).
            callback: Callback function called with each HealthStatus update.

        Raises:
            GqldbError: If operation fails.

        Example:
            >>> def on_status(status: HealthStatus):
            ...     print(f"Health status: {status}")
            >>> client.watch(callback=on_status)
        """
        return self._health_service.watch(service, callback)
    def watch_iter(self, service: str = ""):
        """Watch health status changes and yield HealthStatus values.

        This is a generator version of watch() for iteration.

        Args:
            service: Service name to watch (empty for overall health).

        Yields:
            HealthStatus: Each health status update from the server.

        Example:
            >>> for status in client.watch_iter():
            ...     print(f"Health status: {status}")
            ...     if status != HealthStatus.SERVING:
            ...         break
        """
        return self._health_service.watch_iter(service)
    # =========================================================================
    # Admin Service
    # =========================================================================

    def warmup_parser(self, count: int) -> None:
        """Pre-allocate parser instances.

        Args:
            count: Number of parser instances.

        Raises:
            GqldbError: If operation fails.
        """
        return self._admin_service.warmup_parser(count)
    def get_cache_stats(
        self,
        cache_type: CacheType = CacheType.ALL
    ) -> CacheStats:
        """Return cache statistics.

        Args:
            cache_type: Type of cache to query.

        Returns:
            CacheStats.

        Raises:
            GqldbError: If operation fails.
        """
        return self._admin_service.get_cache_stats(cache_type)
    def clear_cache(self, cache_type: CacheType = CacheType.ALL) -> None:
        """Clear the specified cache.

        Args:
            cache_type: Type of cache to clear.

        Raises:
            GqldbError: If operation fails.
        """
        return self._admin_service.clear_cache(cache_type)
    def get_statistics(self, graph_name: str = "") -> Statistics:
        """Return database statistics.

        Args:
            graph_name: Optional graph name (empty for current graph).

        Returns:
            Statistics.

        Raises:
            GqldbError: If operation fails.
        """
        return self._admin_service.get_statistics(graph_name)
    def invalidate_permission_cache(self, username: str = "") -> None:
        """Invalidate the RBAC permission cache.

        Args:
            username: Optional username (empty for all users).

        Raises:
            GqldbError: If operation fails.
        """
        return self._admin_service.invalidate_permission_cache(username)
    def compact(self) -> CompactResult:
        """Trigger manual compaction of the database storage.

        Returns:
            CompactResult.

        Raises:
            GqldbError: If operation fails.
        """
        from .proto import gqldb_pb2

        if self._admin_stub is None:
            channel = self._get_channel()
            self._init_stubs(channel)

        try:
            request = gqldb_pb2.CompactRequest()
            metadata = self._get_session_metadata()
            response = self._admin_stub.Compact(
                request,
                metadata=metadata,
                timeout=self._config.timeout
            )
            return CompactResult(
                success=response.success,
                message=response.message
            )
        except grpc.RpcError as e:
            raise GqldbError(f"Compact failed: {e.details()}") from e

    # =========================================================================
    # Bulk Import Service
    # =========================================================================

    def start_bulk_import(
        self,
        graph_name: str,
        checkpoint_every: int = 0,
        estimated_nodes: int = 0,
        estimated_edges: int = 0,
        memtable_size: int = 0,
        max_memtables: int = 0
    ) -> BulkImportSession:
        """Start a bulk import session for optimized high-throughput inserts.

        Args:
            graph_name: Graph name.
            checkpoint_every: Deprecated. No longer used, server ignores this value.
            estimated_nodes: Hint for pre-allocating node ID cache.
            estimated_edges: Hint for edge batch sizing.
            memtable_size: Memtable size in bytes before flush (default: 64MB).
            max_memtables: Max immutable memtables before stall (default: 4).

        Returns:
            BulkImportSession.

        Raises:
            GqldbError: If operation fails.
        """
        return self._bulk_import_service.start_bulk_import(graph_name, checkpoint_every, estimated_nodes, estimated_edges, memtable_size, max_memtables)
    def checkpoint(self, session_id: str) -> CheckpointResult:
        """Flush all accumulated data to disk for durability.

        .. deprecated::
            Checkpoint is no longer needed. Use end_bulk_import() which performs a final flush.

        Args:
            session_id: Bulk import session ID.

        Returns:
            CheckpointResult.
        """
        import warnings
        warnings.warn(
            "checkpoint() is deprecated. Use end_bulk_import() which performs a final flush.",
            DeprecationWarning,
            stacklevel=2,
        )
        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:
        """End the bulk import session with a final checkpoint.

        Args:
            session_id: Bulk import session ID.

        Returns:
            EndBulkImportResult.

        Raises:
            GqldbError: If operation fails.
        """
        return self._bulk_import_service.end_bulk_import(session_id)
    def abort_bulk_import(self, session_id: str) -> AbortBulkImportResult:
        """Cancel the bulk import session without final sync.

        Args:
            session_id: Bulk import session ID.

        Returns:
            AbortBulkImportResult.

        Raises:
            GqldbError: If operation fails.
        """
        return self._bulk_import_service.abort_bulk_import(session_id)
    def get_bulk_import_status(self, session_id: str) -> BulkImportStatus:
        """Return the current status of a bulk import session.

        Args:
            session_id: Bulk import session ID.

        Returns:
            BulkImportStatus.

        Raises:
            GqldbError: If operation fails.
        """
        return self._bulk_import_service.get_bulk_import_status(session_id)
    # =========================================================================
    # Convenience Methods
    # =========================================================================

    def get_session(self) -> Optional[Session]:
        """Return the current session."""
        return self._sessions.get_session()

    def is_logged_in(self) -> bool:
        """Check if there is an active session."""
        return self._sessions.is_logged_in()

    @property
    def config(self) -> GqldbConfig:
        """Return the client configuration."""
        return self._config
