foxhunt/database/src/pool.rs

use crate::error::{DatabaseError, DatabaseResult};
use config::PoolConfig;
use sqlx::postgres::{PgPool, PgPoolOptions};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};

// PoolConfig is now imported from the config crate

/// Extension trait for PoolConfig validation
trait PoolConfigValidation {
    fn validate(&self) -> DatabaseResult<()>;
}

impl PoolConfigValidation for PoolConfig {
    fn validate(&self) -> DatabaseResult<()> {
        if self.min_connections > self.max_connections {
            return Err(DatabaseError::Configuration {
                message: "min_connections cannot be greater than max_connections".to_string(),
            });
        }

        if self.max_connections == 0 {
            return Err(DatabaseError::Configuration {
                message: "max_connections must be greater than 0".to_string(),
            });
        }

        if self.acquire_timeout_secs == 0 {
            return Err(DatabaseError::Configuration {
                message: "acquire_timeout_secs must be greater than 0".to_string(),
            });
        }

        if !self.database_url.starts_with("postgres://")
            && !self.database_url.starts_with("postgresql://")
        {
            return Err(DatabaseError::Configuration {
                message: "database_url must be a valid PostgreSQL connection string".to_string(),
            });
        }

        Ok(())
    }
}

/// Connection pool statistics
#[derive(Debug, Clone)]
pub struct PoolStats {
    /// Total number of connections created
    pub total_connections_created: u64,
    /// Number of active connections
    pub active_connections: u32,
    /// Number of idle connections
    pub idle_connections: u32,
    /// Total number of connection acquisitions
    pub total_acquisitions: u64,
    /// Number of failed acquisitions
    pub failed_acquisitions: u64,
    /// Number of connections closed due to max lifetime
    pub connections_closed_max_lifetime: u64,
    /// Number of connections closed due to idle timeout
    pub connections_closed_idle_timeout: u64,
    /// Number of failed health checks
    pub failed_health_checks: u64,
}

impl Default for PoolStats {
    fn default() -> Self {
        Self {
            total_connections_created: 0,
            active_connections: 0,
            idle_connections: 0,
            total_acquisitions: 0,
            failed_acquisitions: 0,
            connections_closed_max_lifetime: 0,
            connections_closed_idle_timeout: 0,
            failed_health_checks: 0,
        }
    }
}

/// Enhanced database connection pool with monitoring and health checks
#[derive(Debug)]
pub struct DatabasePool {
    /// The underlying SQLx pool
    inner: PgPool,
    /// Pool configuration
    config: PoolConfig,
    /// Pool statistics
    stats: Arc<RwLock<PoolStats>>,
    /// Atomic counters for thread-safe metrics
    total_acquisitions: Arc<AtomicU64>,
    failed_acquisitions: Arc<AtomicU64>,
    total_connections_created: Arc<AtomicU64>,
}

impl DatabasePool {
    /// Create a new database pool with the given configuration
    pub async fn new(config: PoolConfig) -> DatabaseResult<Self> {
        // Validate configuration
        config.validate()?;

        info!(
            "Creating database pool with {} min connections, {} max connections",
            config.min_connections, config.max_connections
        );

        // Build the pool
        let pool = PgPoolOptions::new()
            .min_connections(config.min_connections)
            .max_connections(config.max_connections)
            .acquire_timeout(Duration::from_secs(config.acquire_timeout_secs))
            .max_lifetime(Duration::from_secs(config.max_lifetime_secs))
            .idle_timeout(Duration::from_secs(config.idle_timeout_secs))
            .test_before_acquire(config.test_before_acquire)
            .connect(&config.database_url)
            .await
            .map_err(|e| DatabaseError::ConnectionPool {
                message: format!("Failed to create connection pool: {}", e),
            })?;

        let database_pool = Self {
            inner: pool,
            config: config.clone(),
            stats: Arc::new(RwLock::new(PoolStats::default())),
            total_acquisitions: Arc::new(AtomicU64::new(0)),
            failed_acquisitions: Arc::new(AtomicU64::new(0)),
            total_connections_created: Arc::new(AtomicU64::new(0)),
        };

        // Start health check task if enabled
        if config.health_check_enabled {
            database_pool.start_health_check_task().await;
        }

        info!("Database pool created successfully");
        Ok(database_pool)
    }

    /// Get a connection from the pool
    pub async fn acquire(&self) -> DatabaseResult<sqlx::pool::PoolConnection<sqlx::Postgres>> {
        self.total_acquisitions.fetch_add(1, Ordering::Relaxed);

        debug!("Acquiring connection from pool");

        match self.inner.acquire().await {
            Ok(conn) => {
                debug!("Successfully acquired connection from pool");
                Ok(conn)
            }
            Err(e) => {
                self.failed_acquisitions.fetch_add(1, Ordering::Relaxed);
                error!("Failed to acquire connection from pool: {}", e);
                Err(DatabaseError::from(e))
            }
        }
    }

    /// Get a reference to the underlying pool for direct use
    pub fn inner(&self) -> &PgPool {
        &self.inner
    }

    /// Get pool statistics
    pub async fn stats(&self) -> PoolStats {
        let mut stats = self.stats.read().await.clone();

        // Update atomic counters
        stats.total_acquisitions = self.total_acquisitions.load(Ordering::Relaxed);
        stats.failed_acquisitions = self.failed_acquisitions.load(Ordering::Relaxed);
        stats.total_connections_created = self.total_connections_created.load(Ordering::Relaxed);

        // Get current pool state
        stats.active_connections = self.inner.size();
        stats.idle_connections = self.inner.num_idle() as u32;

        stats
    }

    /// Check if the pool is healthy
    pub async fn health_check(&self) -> DatabaseResult<bool> {
        debug!("Performing pool health check");

        match sqlx::query("SELECT 1").fetch_one(&self.inner).await {
            Ok(_) => {
                debug!("Pool health check passed");
                Ok(true)
            }
            Err(e) => {
                warn!("Pool health check failed: {}", e);
                let mut stats = self.stats.write().await;
                stats.failed_health_checks += 1;
                Ok(false)
            }
        }
    }

    /// Get the pool configuration
    pub fn config(&self) -> &PoolConfig {
        &self.config
    }

    /// Close the pool gracefully
    pub async fn close(&self) {
        info!("Closing database pool");
        self.inner.close().await;
        info!("Database pool closed");
    }

    /// Check if the pool is closed
    pub fn is_closed(&self) -> bool {
        self.inner.is_closed()
    }

    /// Start the health check background task
    async fn start_health_check_task(&self) {
        if !self.config.health_check_enabled {
            return;
        }

        let pool = self.inner.clone();
        let stats = self.stats.clone();
        let interval = Duration::from_secs(self.config.health_check_interval_secs);

        tokio::spawn(async move {
            let mut interval_timer = tokio::time::interval(interval);

            loop {
                interval_timer.tick().await;

                if pool.is_closed() {
                    debug!("Pool is closed, stopping health check task");
                    break;
                }

                match sqlx::query("SELECT 1").fetch_one(&pool).await {
                    Ok(_) => {
                        debug!("Scheduled health check passed");
                    }
                    Err(e) => {
                        warn!("Scheduled health check failed: {}", e);
                        let mut stats = stats.write().await;
                        stats.failed_health_checks += 1;
                    }
                }
            }
        });

        debug!(
            "Health check task started with {}s interval",
            self.config.health_check_interval_secs
        );
    }

    /// Execute a test query to validate the connection
    pub async fn ping(&self) -> DatabaseResult<()> {
        sqlx::query("SELECT 1")
            .fetch_one(&self.inner)
            .await
            .map_err(DatabaseError::from)?;
        Ok(())
    }

    /// Get current pool size
    pub fn size(&self) -> u32 {
        self.inner.size()
    }

    /// Get number of idle connections
    pub fn num_idle(&self) -> usize {
        self.inner.num_idle()
    }

    /// Reset pool statistics
    pub async fn reset_stats(&self) {
        let mut stats = self.stats.write().await;
        *stats = PoolStats::default();
        self.total_acquisitions.store(0, Ordering::Relaxed);
        self.failed_acquisitions.store(0, Ordering::Relaxed);
        self.total_connections_created.store(0, Ordering::Relaxed);
        info!("Pool statistics reset");
    }
}

impl Clone for DatabasePool {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            config: self.config.clone(),
            stats: self.stats.clone(),
            total_acquisitions: self.total_acquisitions.clone(),
            failed_acquisitions: self.failed_acquisitions.clone(),
            total_connections_created: self.total_connections_created.clone(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pool_config_validation() {
        let mut config = PoolConfig::default();
        assert!(config.validate().is_ok());

        // Test min > max
        config.min_connections = 10;
        config.max_connections = 5;
        assert!(config.validate().is_err());

        // Test max = 0
        config.min_connections = 0;
        config.max_connections = 0;
        assert!(config.validate().is_err());

        // Test invalid URL
        config.max_connections = 10;
        config.database_url = "invalid://url".to_string();
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_pool_config_default() {
        let config = PoolConfig::default();
        assert_eq!(config.min_connections, 5);
        assert_eq!(config.max_connections, 100);
        assert!(config.test_before_acquire);
        assert!(config.health_check_enabled);
    }

    #[tokio::test]
    async fn test_pool_stats_default() {
        let stats = PoolStats::default();
        assert_eq!(stats.total_connections_created, 0);
        assert_eq!(stats.active_connections, 0);
        assert_eq!(stats.failed_acquisitions, 0);
    }
}