"""Connection pool implementation for database management."""

from __future__ import annotations

import threading
from collections import deque
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Callable, Generator

from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

from pytola_sim.lscsim.utils.logger import get_logger

if TYPE_CHECKING:
    from sqlalchemy.engine import Engine

logger = get_logger(__name__)


@dataclass
class ConnectionPoolConfig:
    """Configuration for connection pool."""

    max_connections: int = 10
    min_connections: int = 2
    connection_timeout: int = 30  # seconds
    max_idle_time: int = 300  # seconds
    max_usage_count: int = 1000  # reuse connections up to this count
    enable_logging: bool = True


@dataclass
class PooledConnection:
    """Wrapper for pooled database connection."""

    engine: Engine
    session_factory: sessionmaker
    created_at: datetime
    last_used: datetime
    usage_count: int = 0
    is_active: bool = True

    def is_expired(self, max_idle_time: int, max_usage_count: int) -> bool:
        """Check if connection should be expired."""
        now = datetime.now(timezone.utc)
        idle_time = (now - self.last_used).total_seconds()
        return (idle_time > max_idle_time) or (self.usage_count >= max_usage_count)


class ConnectionPool:
    """Thread-safe connection pool for database connections."""

    def __init__(self, database_url: str, config: ConnectionPoolConfig | None = None) -> None:
        self.database_url = database_url
        self.config = config or ConnectionPoolConfig()
        self._pool: deque[PooledConnection] = deque()
        self._active_connections: dict[int, PooledConnection] = {}
        self._lock = threading.RLock()
        self._current_size = 0

        # Pre-create minimum connections
        self._pre_create_connections()
        logger.info(
            f"Connection pool initialized with {self._current_size} connections",
        )

    def _pre_create_connections(self) -> None:
        """Pre-create minimum number of connections."""
        for _ in range(self.config.min_connections):
            self._create_connection()

    def _create_connection(self) -> PooledConnection | None:
        """Create a new database connection."""
        try:
            engine = create_engine(
                self.database_url,
                pool_pre_ping=True,
                pool_recycle=3600,
                echo=self.config.enable_logging,
            )

            session_factory = sessionmaker(bind=engine)

            connection = PooledConnection(
                engine=engine,
                session_factory=session_factory,
                created_at=datetime.now(timezone.utc),
                last_used=datetime.now(timezone.utc),
            )

            self._current_size += 1
        except (OSError, RuntimeError, ValueError):
            logger.exception("Failed to create database connection")
            return None
        else:
            return connection

    def acquire(self) -> PooledConnection | None:
        """Acquire a connection from the pool."""
        with self._lock:
            # Try to get an existing connection
            while self._pool:
                connection = self._pool.popleft()

                # Check if connection is still valid
                if connection.is_active and not connection.is_expired(
                    self.config.max_idle_time,
                    self.config.max_usage_count,
                ):
                    connection.usage_count += 1
                    connection.last_used = datetime.now(timezone.utc)
                    self._active_connections[id(connection)] = connection
                    return connection

                # Remove expired connection
                self._destroy_connection(connection)

            # Create new connection if under limit
            if self._current_size < self.config.max_connections:
                connection = self._create_connection()
                if connection:
                    connection.usage_count += 1
                    connection.last_used = datetime.now(timezone.utc)
                    self._active_connections[id(connection)] = connection
                    return connection

            return None

    def release(self, connection: PooledConnection) -> None:
        """Release a connection back to the pool."""
        with self._lock:
            conn_id = id(connection)
            if conn_id in self._active_connections:
                del self._active_connections[conn_id]

                if connection.is_active and not connection.is_expired(
                    self.config.max_idle_time,
                    self.config.max_usage_count,
                ):
                    self._pool.append(connection)
                else:
                    self._destroy_connection(connection)

    def _destroy_connection(self, connection: PooledConnection) -> None:
        """Destroy a connection and clean up resources."""
        try:
            connection.engine.dispose()
            self._current_size -= 1
            logger.debug("Connection destroyed")
        except (OSError, RuntimeError, ValueError):
            logger.exception("Error destroying connection")

    def cleanup_expired_connections(self) -> None:
        """Remove expired connections from the pool."""
        with self._lock:
            # Clean up pool
            expired_connections = []
            active_pool = deque()

            while self._pool:
                connection = self._pool.popleft()
                if connection.is_expired(
                    self.config.max_idle_time,
                    self.config.max_usage_count,
                ):
                    expired_connections.append(connection)
                else:
                    active_pool.append(connection)

            self._pool = active_pool

            # Destroy expired connections
            for connection in expired_connections:
                self._destroy_connection(connection)

            if expired_connections and self.config.enable_logging:
                logger.info(
                    f"Cleaned up {len(expired_connections)} expired connections",
                )

    def get_stats(self) -> dict[str, Any]:
        """Get pool statistics."""
        with self._lock:
            return {
                "current_pool_size": len(self._pool),
                "active_connections": len(self._active_connections),
                "total_connections": self._current_size,
                "max_connections": self.config.max_connections,
                "min_connections": self.config.min_connections,
            }

    def close_all(self) -> None:
        """Close all connections in the pool."""
        with self._lock:
            # Close pooled connections
            while self._pool:
                connection = self._pool.popleft()
                self._destroy_connection(connection)

            # Close active connections
            for connection in list(self._active_connections.values()):
                self._destroy_connection(connection)

            self._active_connections.clear()
            logger.info("All connections closed")


@contextmanager
def get_connection(pool: ConnectionPool) -> Generator[PooledConnection, None, None]:
    """Context manager for acquiring and releasing connections."""
    connection = pool.acquire()
    if not connection:
        msg = "Failed to acquire database connection from pool"
        raise RuntimeError(msg)

    try:
        yield connection
    finally:
        pool.release(connection)


class PooledDatabaseManager:
    """Database manager using connection pooling."""

    def __init__(
        self,
        database_url: str,
        pool_config: ConnectionPoolConfig | None = None,
    ) -> None:
        self.pool = ConnectionPool(database_url, pool_config)
        logger.info("Pooled database manager initialized")

    @contextmanager
    def get_session(self) -> Generator[Session, None, None]:
        """Get database session from connection pool."""
        with get_connection(self.pool) as pooled_conn:
            session = pooled_conn.session_factory()
            try:
                yield session
                session.commit()
            except (OSError, RuntimeError, ValueError):
                session.rollback()
                raise
            finally:
                session.close()

    def execute_query(self, query_func: Callable[[Session], object]) -> object:
        """Execute a query function with automatic connection management."""
        with self.get_session() as session:
            return query_func(session)

    def get_pool_stats(self) -> dict[str, Any]:
        """Get connection pool statistics."""
        return self.pool.get_stats()

    def cleanup(self) -> None:
        """Perform cleanup of expired connections."""
        self.pool.cleanup_expired_connections()

    def close(self) -> None:
        """Close all connections."""
        self.pool.close_all()
