"""
Connection management for mongoengine read-write router.

Handles routing decisions for read and write operations based on context and configuration.
"""

from contextlib import contextmanager
from typing import Generator

from .config import get_config
from .context import has_write_operations, is_in_transaction


class MongoConnectionManager:
    """
    Manager for routing database operations to appropriate connections.

    This class implements the core routing logic for read-write separation,
    supporting multiple consistency strategies.
    """

    def __init__(self):
        self._force_primary_stack: list[bool] = []

    def get_read_db_alias(self, user_id: str | None = None) -> str:
        """
        Get the appropriate database alias for read operations.

        Routing decision is based on:
        1. If force_primary is active, use primary
        2. If in transaction and transaction_read_primary is enabled, use primary
        3. If read_after_write_consistency is enabled and write occurred, use primary
        4. Use strategy-specific routing

        Args:
            user_id: Optional user ID for consistent hashing

        Returns:
            The database alias to use for read operations
        """
        config = get_config()

        # Check if we're forcing primary reads
        if self._force_primary_stack and self._force_primary_stack[-1]:
            return config.primary_alias

        # Check if we should use primary based on context
        if self._should_use_primary_for_read():
            return config.primary_alias

        # Use strategy-specific routing
        return self._route_by_strategy(user_id)

    def get_write_db_alias(self) -> str:
        """
        Get the database alias for write operations.

        Writes always go to the primary database.

        Returns:
            The primary database alias
        """
        return get_config().primary_alias

    @contextmanager
    def force_primary(self) -> Generator[None, None, None]:
        """
        Context manager to force all operations to use the primary database.

        This is useful when you need to ensure fresh data for critical operations.

        Example:
            with connection_manager.force_primary():
                user = User.objects.first()  # Will read from primary
        """
        self._force_primary_stack.append(True)
        try:
            yield
        finally:
            self._force_primary_stack.pop()

    def _should_use_primary_for_read(self) -> bool:
        """
        Determine if reads should go to primary based on context.

        Returns:
            True if reads should go to primary, False otherwise
        """
        config = get_config()

        # Check transaction awareness
        if config.strategy == "transaction_aware" or config.transaction_read_primary:
            if is_in_transaction():
                return True

        # Check read-after-write consistency
        if config.read_after_write_consistency and has_write_operations():
            return True

        return False

    def _route_by_strategy(self, user_id: str | None) -> str:
        """
        Route read operation based on configured strategy.

        Args:
            user_id: Optional user ID for consistent hashing

        Returns:
            The database alias to use for read operations
        """
        config = get_config()

        if config.strategy == "random":
            return config.get_random_replica()

        elif config.strategy == "transaction_aware":
            # Already checked transaction state in _should_use_primary_for_read
            # So just route to a random replica
            return config.get_random_replica()

        elif config.strategy == "consistent_hash":
            if user_id:
                return config.get_replica_for_user(user_id)
            # Fall back to random if no user_id
            return config.get_random_replica()

        # Default fallback
        return config.replica_aliases[0] if config.replica_aliases else config.primary_alias


# Global connection manager instance
_connection_manager: MongoConnectionManager | None = None


def get_connection_manager() -> MongoConnectionManager:
    """Get the global connection manager instance."""
    global _connection_manager
    if _connection_manager is None:
        _connection_manager = MongoConnectionManager()
    return _connection_manager


def get_read_db_alias(user_id: str | None = None) -> str:
    """
    Get the appropriate database alias for read operations.

    Args:
        user_id: Optional user ID for consistent hashing

    Returns:
        The database alias to use for read operations
    """
    return get_connection_manager().get_read_db_alias(user_id)


def get_write_db_alias() -> str:
    """
    Get the database alias for write operations.

    Returns:
        The primary database alias
    """
    return get_connection_manager().get_write_db_alias()


@contextmanager
def force_primary() -> Generator[None, None, None]:
    """
    Context manager to force all operations to use the primary database.

    Example:
        with force_primary():
            user = User.objects.first()  # Will read from primary
    """
    with get_connection_manager().force_primary():
        yield
