import queue
import logging
import threading
from contextlib import contextmanager
from typing import Callable, Generator
from ..backends.base import DatabaseBackend

logger = logging.getLogger(__name__)


class ConnectionPool:
    def __init__(
        self,
        backend_factory: Callable[[], DatabaseBackend],
        min_size: int = 1,
        max_size: int = 5,
    ):
        self._backend_factory = backend_factory
        self._min_size = min_size
        self._max_size = max_size
        self._pool: queue.Queue = queue.Queue(maxsize=max_size)
        self._created_count = 0
        self._lock = threading.Lock()

        for _ in range(min_size):
            self._create_and_add_connection()

    def _create_and_add_connection(self):
        """Create a new connection and add it to the pool."""
        # Check if we can create more connections
        with self._lock:
            if self._created_count >= self._max_size:
                return
            self._created_count += 1
            current_count = self._created_count

        # Create connection outside the lock to avoid blocking other operations
        try:
            conn = self._backend_factory()
            conn.connect()
            self._pool.put_nowait(conn)
            logger.info(
                "Created new database connection. Total connections: %d", current_count
            )
        except queue.Full:
            # Queue is full, shouldn't happen but handle gracefully
            with self._lock:
                self._created_count -= 1
            logger.warning("Connection pool queue is full, cannot add connection")
        except Exception as e:
            # If connection creation fails, decrement the count
            with self._lock:
                self._created_count -= 1
            logger.error(f"Failed to create database connection: {e}")
            raise

    def get_connection(
        self, block: bool = True, timeout: float = None
    ) -> DatabaseBackend:
        """Get a connection from the pool, creating a new one if needed."""
        try:
            conn = self._pool.get(block=False)
            logger.info("Acquired database connection from pool.")
            return conn
        except queue.Empty:
            # Check if we can create more connections
            with self._lock:
                can_create = self._created_count < self._max_size

            if can_create:
                self._create_and_add_connection()
                # Try to get the newly created connection
                try:
                    return self._pool.get(block=False)
                except queue.Empty:
                    # Another thread might have taken it, try with blocking behavior
                    return self._pool.get(block=block, timeout=timeout)
            else:
                # Pool is exhausted and max connections already created
                # Try to get from pool with the specified blocking behavior
                return self._pool.get(block=block, timeout=timeout)

    def release_connection(self, conn: DatabaseBackend):
        try:
            self._pool.put_nowait(conn)
            logger.info("Released database connection back to pool.")
        except queue.Full:
            logger.warning("Connection pool is full, closing connection instead")
            try:
                conn.close()
            except Exception as e:
                logger.error(f"Error closing connection: {e}")

    def close_all(self):
        while not self._pool.empty():
            conn = self._pool.get()
            try:
                conn.disconnect()
            except Exception as e:
                logger.error(f"Error disconnecting database connection: {e}")
        with self._lock:
            self._created_count = 0

    @contextmanager
    def connection(self) -> Generator[DatabaseBackend, None, None]:
        conn = self.get_connection()
        try:
            yield conn
        finally:
            self.release_connection(conn)
