import logging
import threading
import time
from contextlib import contextmanager
from pathlib import Path
from queue import Queue, Empty

import chromadb
from chromadb import Settings

logger = logging.getLogger(__name__)


class ChromaConnection:
    """Wrapper for a ChromaDB client connection with health tracking."""
    
    def __init__(self, client: chromadb.PersistentClient, created_at: float):
        self.client = client
        self.created_at = created_at
        self.last_used = created_at
        self.is_healthy = True
        self.use_count = 0
    
    def mark_used(self):
        """Mark this connection as recently used."""
        self.last_used = time.time()
        self.use_count += 1
    
    def check_health(self) -> bool:
        """Check if the connection is healthy."""
        try:
            # Simple health check - try to list collections
            self.client.list_collections()
            self.is_healthy = True
            return True
        except Exception as e:
            logger.warning(f"Connection health check failed: {e}")
            self.is_healthy = False
            return False
    
    def close(self):
        """Close the connection."""
        try:
            # ChromaDB PersistentClient doesn't have an explicit close method
            # but we can mark it as unhealthy
            self.is_healthy = False
        except Exception as e:
            logger.warning(f"Error closing connection: {e}")


class ChromaConnectionPool:
    """
    Connection pool for ChromaDB clients.
    
    Manages a pool of ChromaDB client connections with configurable size,
    timeout, and health monitoring capabilities.
    """
    
    def __init__(
        self,
        db_dir: Path,
        pool_size: int = 5,
        max_pool_size: int = 10,
        connection_timeout: float = 30.0,
        health_check_interval: float = 300.0,  # 5 minutes
        max_connection_age: float = 3600.0,    # 1 hour
    ):
        """
        Initialize the connection pool.
        
        Args:
            db_dir: Path to the ChromaDB database directory
            pool_size: Initial number of connections to create
            max_pool_size: Maximum number of connections in the pool
            connection_timeout: Timeout in seconds for getting a connection
            health_check_interval: Interval in seconds between health checks
            max_connection_age: Maximum age in seconds before connection is recycled
        """
        self.db_dir = db_dir
        self.pool_size = pool_size
        self.max_pool_size = max_pool_size
        self.connection_timeout = connection_timeout
        self.health_check_interval = health_check_interval
        self.max_connection_age = max_connection_age
        
        self._pool: Queue[ChromaConnection] = Queue(maxsize=max_pool_size)
        self._lock = threading.RLock()
        self._created_connections = 0
        self._last_health_check = 0.0
        self._closed = False
        
        # Initialize the pool with initial connections
        self._initialize_pool()
    
    def _create_connection(self) -> ChromaConnection:
        """Create a new ChromaDB connection."""
        try:
            client = chromadb.PersistentClient(
                path=self.db_dir,
                settings=Settings(allow_reset=True),
            )
            connection = ChromaConnection(client, time.time())
            logger.debug(f"Created new ChromaDB connection (total: {self._created_connections + 1})")
            return connection
        except Exception as e:
            logger.error(f"Failed to create ChromaDB connection: {e}")
            raise
    
    def _initialize_pool(self):
        """Initialize the pool with the configured number of connections."""
        with self._lock:
            for _ in range(self.pool_size):
                if self._created_connections < self.max_pool_size:
                    try:
                        connection = self._create_connection()
                        self._pool.put_nowait(connection)
                        self._created_connections += 1
                    except Exception as e:
                        logger.error(f"Failed to initialize connection: {e}")
                        break
    
    def _should_recycle_connection(self, connection: ChromaConnection) -> bool:
        """Check if a connection should be recycled due to age or health."""
        current_time = time.time()
        
        # Check age
        if current_time - connection.created_at > self.max_connection_age:
            logger.debug("Connection exceeded max age, recycling")
            return True
        
        # Check health if it's time
        if current_time - self._last_health_check > self.health_check_interval:
            if not connection.check_health():
                logger.debug("Connection failed health check, recycling")
                return True
        
        return False
    
    def _perform_health_checks(self):
        """Perform health checks on connections if needed."""
        current_time = time.time()
        if current_time - self._last_health_check > self.health_check_interval:
            self._last_health_check = current_time
            logger.debug("Performing connection pool health checks")
    
    @contextmanager
    def get_connection(self):
        """
        Get a connection from the pool as a context manager.
        
        Yields:
            ChromaConnection: A healthy connection from the pool
            
        Raises:
            RuntimeError: If pool is closed or no connection available within timeout
        """
        if self._closed:
            raise RuntimeError("Connection pool is closed")
        
        connection = None
        start_time = time.time()
        
        try:
            # Try to get a connection from the pool
            while time.time() - start_time < self.connection_timeout:
                try:
                    connection = self._pool.get(timeout=1.0)
                    
                    # Check if connection should be recycled
                    if self._should_recycle_connection(connection):
                        connection.close()
                        with self._lock:
                            self._created_connections -= 1
                        connection = None
                        continue
                    
                    # Mark connection as used and yield it
                    connection.mark_used()
                    self._perform_health_checks()
                    yield connection
                    return
                    
                except Empty:
                    # No connection available, try to create a new one if under limit
                    with self._lock:
                        if self._created_connections < self.max_pool_size:
                            try:
                                connection = self._create_connection()
                                self._created_connections += 1
                                connection.mark_used()
                                self._perform_health_checks()
                                yield connection
                                return
                            except Exception as e:
                                logger.error(f"Failed to create new connection: {e}")
                                continue
            
            raise RuntimeError(f"Failed to get connection within {self.connection_timeout} seconds")
            
        finally:
            # Return connection to pool if we got one
            if connection and not self._closed:
                try:
                    self._pool.put_nowait(connection)
                except Exception as e:
                    logger.warning(f"Failed to return connection to pool: {e}")
                    # If we can't return it, close it and decrement counter
                    connection.close()
                    with self._lock:
                        self._created_connections -= 1
    
    def get_stats(self) -> dict:
        """Get connection pool statistics."""
        with self._lock:
            return {
                "pool_size": self.pool_size,
                "max_pool_size": self.max_pool_size,
                "current_connections": self._created_connections,
                "available_connections": self._pool.qsize(),
                "connection_timeout": self.connection_timeout,
                "health_check_interval": self.health_check_interval,
                "max_connection_age": self.max_connection_age,
                "last_health_check": self._last_health_check,
                "closed": self._closed,
            }
    
    def close(self):
        """Close all connections in the pool."""
        if self._closed:
            return
        
        self._closed = True
        connections_closed = 0
        
        # Close all connections in the pool
        while not self._pool.empty():
            try:
                connection = self._pool.get_nowait()
                connection.close()
                connections_closed += 1
            except Empty:
                break
            except Exception as e:
                logger.warning(f"Error closing connection: {e}")
        
        with self._lock:
            self._created_connections = 0
        
        logger.info(f"Closed {connections_closed} connections in ChromaDB connection pool")
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()