foxhunt/database/src/transaction.rs

use crate::error::{DatabaseError, DatabaseResult, ErrorContext};
use crate::pool::DatabasePool;
use config::database::TransactionConfig;
use sqlx::FromRow;
use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

// TransactionConfig is now imported from the config crate

/// Transaction manager for handling database transactions
#[derive(Debug)]
pub struct TransactionManager {
    pool: DatabasePool,
    config: TransactionConfig,
    active_transactions: Arc<AtomicU64>,
    total_transactions: Arc<AtomicU64>,
    failed_transactions: Arc<AtomicU64>,
    retry_attempts: Arc<AtomicU64>,
}

impl TransactionManager {
    /// Create a new transaction manager
    pub fn new(pool: DatabasePool, config: TransactionConfig) -> Self {
        Self {
            pool,
            config,
            active_transactions: Arc::new(AtomicU64::new(0)),
            total_transactions: Arc::new(AtomicU64::new(0)),
            failed_transactions: Arc::new(AtomicU64::new(0)),
            retry_attempts: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Begin a new transaction with default timeout
    pub async fn begin(&self) -> DatabaseResult<DatabaseTransaction> {
        self.begin_with_timeout(Duration::from_secs(self.config.default_timeout_secs))
            .await
    }

    /// Begin a new transaction with custom timeout
    pub async fn begin_with_timeout(
        &self,
        timeout_duration: Duration,
    ) -> DatabaseResult<DatabaseTransaction> {
        let start_time = Instant::now();
        self.total_transactions.fetch_add(1, Ordering::Relaxed);
        self.active_transactions.fetch_add(1, Ordering::Relaxed);

        debug!(
            "Beginning new database transaction with {}s timeout",
            timeout_duration.as_secs()
        );

        // Get the transaction directly from the pool to avoid lifetime issues
        let transaction = self.pool.inner().begin().await?;
        let transaction_id = Uuid::new_v4();

        debug!("Transaction {} started successfully", transaction_id);
        Ok(DatabaseTransaction {
            inner: Some(transaction),
            id: transaction_id,
            start_time,
            timeout: timeout_duration,
            savepoints: Vec::new(),
            config: self.config.clone(),
            manager_stats: TransactionManagerStats {
                active_transactions: self.active_transactions.clone(),
                failed_transactions: self.failed_transactions.clone(),
            },
        })
    }

    /// Execute a closure within a transaction with retry logic
    pub async fn with_transaction<F, R, Fut>(&self, f: F) -> DatabaseResult<R>
    where
        F: Fn(DatabaseTransaction) -> Fut + Send + Sync,
        Fut: Future<Output = DatabaseResult<(R, DatabaseTransaction)>> + Send,
        R: Send,
    {
        self.with_transaction_timeout(f, Duration::from_secs(self.config.default_timeout_secs))
            .await
    }

    /// Execute a closure within a transaction with custom timeout and retry logic
    pub async fn with_transaction_timeout<F, R, Fut>(
        &self,
        f: F,
        timeout_duration: Duration,
    ) -> DatabaseResult<R>
    where
        F: Fn(DatabaseTransaction) -> Fut + Send + Sync,
        Fut: Future<Output = DatabaseResult<(R, DatabaseTransaction)>> + Send,
        R: Send,
    {
        let mut attempts = 0;
        let max_attempts = if self.config.enable_retry {
            self.config.max_retries + 1
        } else {
            1
        };

        loop {
            attempts += 1;

            let tx = self.begin_with_timeout(timeout_duration).await?;
            let tx_id = tx.id();

            match f(tx).await {
                Ok((result, tx)) => match tx.commit().await {
                    Ok(_) => {
                        debug!(
                            "Transaction {} completed successfully on attempt {}",
                            tx_id, attempts
                        );
                        return Ok(result);
                    },
                    Err(e) if e.is_retryable() && attempts < max_attempts => {
                        self.retry_attempts.fetch_add(1, Ordering::Relaxed);
                        warn!(
                            "Transaction {} commit failed (attempt {}): {}. Retrying...",
                            tx_id, attempts, e
                        );
                        self.delay_retry(attempts).await;
                        continue;
                    },
                    Err(e) => {
                        error!(
                            "Transaction {} commit failed on attempt {}: {}",
                            tx_id, attempts, e
                        );
                        return Err(e);
                    },
                },
                Err(e) if e.is_retryable() && attempts < max_attempts => {
                    self.retry_attempts.fetch_add(1, Ordering::Relaxed);
                    warn!(
                        "Transaction {} failed (attempt {}): {}. Retrying...",
                        tx_id, attempts, e
                    );
                    self.delay_retry(attempts).await;
                    continue;
                },
                Err(e) => {
                    error!(
                        "Transaction {} failed on attempt {}: {}",
                        tx_id, attempts, e
                    );
                    return Err(e);
                },
            }
        }
    }

    /// Get transaction statistics
    pub fn stats(&self) -> TransactionStats {
        TransactionStats {
            active_transactions: self.active_transactions.load(Ordering::Relaxed),
            total_transactions: self.total_transactions.load(Ordering::Relaxed),
            failed_transactions: self.failed_transactions.load(Ordering::Relaxed),
            retry_attempts: self.retry_attempts.load(Ordering::Relaxed),
        }
    }

    /// Reset transaction statistics
    pub fn reset_stats(&self) {
        self.active_transactions.store(0, Ordering::Relaxed);
        self.total_transactions.store(0, Ordering::Relaxed);
        self.failed_transactions.store(0, Ordering::Relaxed);
        self.retry_attempts.store(0, Ordering::Relaxed);
        info!("Transaction manager statistics reset");
    }

    /// Delay before retry with exponential backoff
    async fn delay_retry(&self, attempt: u32) {
        let delay = Duration::from_millis(
            self.config.retry_delay_ms * (2_u64.pow(attempt.saturating_sub(1))),
        );
        debug!(
            "Waiting {}ms before retry attempt {}",
            delay.as_millis(),
            attempt + 1
        );
        tokio::time::sleep(delay).await;
    }
}

impl Clone for TransactionManager {
    fn clone(&self) -> Self {
        Self {
            pool: self.pool.clone(),
            config: self.config.clone(),
            active_transactions: self.active_transactions.clone(),
            total_transactions: self.total_transactions.clone(),
            failed_transactions: self.failed_transactions.clone(),
            retry_attempts: self.retry_attempts.clone(),
        }
    }
}

/// A database transaction wrapper with enhanced features
pub struct DatabaseTransaction {
    inner: Option<sqlx::Transaction<'static, sqlx::Postgres>>,
    id: Uuid,
    start_time: Instant,
    timeout: Duration,
    savepoints: Vec<String>,
    config: TransactionConfig,
    manager_stats: TransactionManagerStats,
}

/// Statistics tracking for transaction manager
#[derive(Debug, Clone)]
struct TransactionManagerStats {
    active_transactions: Arc<AtomicU64>,
    failed_transactions: Arc<AtomicU64>,
}

impl DatabaseTransaction {
    /// Get the transaction ID
    pub fn id(&self) -> Uuid {
        self.id
    }

    /// Get the elapsed time since transaction start
    pub fn elapsed(&self) -> Duration {
        self.start_time.elapsed()
    }

    /// Check if the transaction has timed out
    pub fn is_timed_out(&self) -> bool {
        self.elapsed() > self.timeout
    }

    /// Get remaining time before timeout
    pub fn remaining_time(&self) -> Duration {
        self.timeout.saturating_sub(self.elapsed())
    }

    /// Create a savepoint
    pub async fn savepoint(&mut self, name: &str) -> DatabaseResult<()> {
        if self.savepoints.len() >= self.config.max_savepoints as usize {
            return Err(DatabaseError::Transaction {
                message: format!(
                    "Maximum number of savepoints ({}) exceeded",
                    self.config.max_savepoints
                ),
            });
        }

        if self.is_timed_out() {
            return Err(DatabaseError::Timeout {
                operation: format!("savepoint_{}", name),
                timeout_secs: self.timeout.as_secs(),
            });
        }

        let savepoint_name = format!("sp_{}_{}", self.id.simple(), name);

        sqlx::query(&format!("SAVEPOINT {}", savepoint_name))
            .execute(&mut **self.inner.as_mut().expect("Transaction already consumed"))
            .await
            .with_context(&format!("Failed to create savepoint {}", savepoint_name))?;

        self.savepoints.push(savepoint_name.clone());
        debug!(
            "Created savepoint {} for transaction {}",
            savepoint_name, self.id
        );
        Ok(())
    }

    /// Rollback to a savepoint
    pub async fn rollback_to_savepoint(&mut self, name: &str) -> DatabaseResult<()> {
        let savepoint_name = format!("sp_{}_{}", self.id.simple(), name);

        if !self.savepoints.contains(&savepoint_name) {
            return Err(DatabaseError::Transaction {
                message: format!("Savepoint {} not found", savepoint_name),
            });
        }

        if self.is_timed_out() {
            return Err(DatabaseError::Timeout {
                operation: format!("rollback_to_savepoint_{}", name),
                timeout_secs: self.timeout.as_secs(),
            });
        }

        sqlx::query(&format!("ROLLBACK TO SAVEPOINT {}", savepoint_name))
            .execute(&mut **self.inner.as_mut().expect("Transaction already consumed"))
            .await
            .with_context(&format!(
                "Failed to rollback to savepoint {}",
                savepoint_name
            ))?;

        // Remove this savepoint and all subsequent ones
        if let Some(pos) = self.savepoints.iter().position(|sp| sp == &savepoint_name) {
            self.savepoints.truncate(pos);
        }

        debug!(
            "Rolled back to savepoint {} for transaction {}",
            savepoint_name, self.id
        );
        Ok(())
    }

    /// Release a savepoint
    pub async fn release_savepoint(&mut self, name: &str) -> DatabaseResult<()> {
        let savepoint_name = format!("sp_{}_{}", self.id.simple(), name);

        if !self.savepoints.contains(&savepoint_name) {
            return Err(DatabaseError::Transaction {
                message: format!("Savepoint {} not found", savepoint_name),
            });
        }

        sqlx::query(&format!("RELEASE SAVEPOINT {}", savepoint_name))
            .execute(&mut **self.inner.as_mut().expect("Transaction already consumed"))
            .await
            .with_context(&format!("Failed to release savepoint {}", savepoint_name))?;

        self.savepoints.retain(|sp| sp != &savepoint_name);
        debug!(
            "Released savepoint {} for transaction {}",
            savepoint_name, self.id
        );
        Ok(())
    }

    /// Execute a query within the transaction
    pub async fn execute(&mut self, query: &str) -> DatabaseResult<u64> {
        if self.is_timed_out() {
            return Err(DatabaseError::Timeout {
                operation: "execute_query".to_string(),
                timeout_secs: self.timeout.as_secs(),
            });
        }

        let result = sqlx::query(query)
            .execute(&mut **self.inner.as_mut().expect("Transaction already consumed"))
            .await
            .with_query_context(query)?;

        Ok(result.rows_affected())
    }

    /// Fetch all rows from a query
    pub async fn fetch_all<T>(&mut self, query: &str) -> DatabaseResult<Vec<T>>
    where
        T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
    {
        if self.is_timed_out() {
            return Err(DatabaseError::Timeout {
                operation: "fetch_all".to_string(),
                timeout_secs: self.timeout.as_secs(),
            });
        }

        let rows = sqlx::query_as(query)
            .fetch_all(&mut **self.inner.as_mut().expect("Transaction already consumed"))
            .await
            .with_query_context(query)?;

        Ok(rows)
    }

    /// Fetch one row from a query
    pub async fn fetch_one<T>(&mut self, query: &str) -> DatabaseResult<T>
    where
        T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
    {
        if self.is_timed_out() {
            return Err(DatabaseError::Timeout {
                operation: "fetch_one".to_string(),
                timeout_secs: self.timeout.as_secs(),
            });
        }

        let row = sqlx::query_as(query)
            .fetch_one(&mut **self.inner.as_mut().expect("Transaction already consumed"))
            .await
            .with_query_context(query)?;

        Ok(row)
    }

    /// Fetch optional row from a query
    pub async fn fetch_optional<T>(&mut self, query: &str) -> DatabaseResult<Option<T>>
    where
        T: for<'r> FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
    {
        if self.is_timed_out() {
            return Err(DatabaseError::Timeout {
                operation: "fetch_optional".to_string(),
                timeout_secs: self.timeout.as_secs(),
            });
        }

        let row = sqlx::query_as(query)
            .fetch_optional(&mut **self.inner.as_mut().expect("Transaction already consumed"))
            .await
            .with_query_context(query)?;

        Ok(row)
    }

    /// Commit the transaction
    pub async fn commit(mut self) -> DatabaseResult<()> {
        if self.is_timed_out() {
            return Err(DatabaseError::Timeout {
                operation: "commit_transaction".to_string(),
                timeout_secs: self.timeout.as_secs(),
            });
        }

        let elapsed = self.elapsed();
        let tx_id = self.id;

        match self.inner.take().unwrap().commit().await {
            Ok(_) => {
                self.manager_stats
                    .active_transactions
                    .fetch_sub(1, Ordering::Relaxed);
                info!(
                    "Transaction {} committed successfully in {}ms",
                    tx_id,
                    elapsed.as_millis()
                );
                Ok(())
            },
            Err(e) => {
                self.manager_stats
                    .active_transactions
                    .fetch_sub(1, Ordering::Relaxed);
                self.manager_stats
                    .failed_transactions
                    .fetch_add(1, Ordering::Relaxed);
                error!(
                    "Transaction {} commit failed after {}ms: {}",
                    tx_id,
                    elapsed.as_millis(),
                    e
                );
                Err(DatabaseError::from(e))
            },
        }
    }

    /// Rollback the transaction
    pub async fn rollback(mut self) -> DatabaseResult<()> {
        let elapsed = self.elapsed();
        let tx_id = self.id;

        match self.inner.take().unwrap().rollback().await {
            Ok(_) => {
                self.manager_stats
                    .active_transactions
                    .fetch_sub(1, Ordering::Relaxed);
                info!(
                    "Transaction {} rolled back successfully in {}ms",
                    tx_id,
                    elapsed.as_millis()
                );
                Ok(())
            },
            Err(e) => {
                self.manager_stats
                    .active_transactions
                    .fetch_sub(1, Ordering::Relaxed);
                self.manager_stats
                    .failed_transactions
                    .fetch_add(1, Ordering::Relaxed);
                error!(
                    "Transaction {} rollback failed after {}ms: {}",
                    tx_id,
                    elapsed.as_millis(),
                    e
                );
                Err(DatabaseError::from(e))
            },
        }
    }
}

impl Drop for DatabaseTransaction {
    fn drop(&mut self) {
        // Only warn if transaction is still active (not None)
        if self.inner.is_some() {
            self.manager_stats
                .active_transactions
                .fetch_sub(1, Ordering::Relaxed);
            warn!("Transaction {} was dropped without explicit commit/rollback - automatic rollback will occur", self.id);
        }
    }
}

/// Transaction statistics
#[derive(Debug, Clone)]
pub struct TransactionStats {
    /// Number of currently active transactions
    pub active_transactions: u64,
    /// Total number of transactions started
    pub total_transactions: u64,
    /// Number of failed transactions
    pub failed_transactions: u64,
    /// Number of retry attempts
    pub retry_attempts: u64,
}

impl TransactionStats {
    /// Calculate success rate as percentage
    pub fn success_rate(&self) -> f64 {
        if self.total_transactions == 0 {
            0.0
        } else {
            ((self.total_transactions - self.failed_transactions) as f64
                / self.total_transactions as f64)
                * 100.0
        }
    }

    /// Calculate retry rate as percentage
    pub fn retry_rate(&self) -> f64 {
        if self.total_transactions == 0 {
            0.0
        } else {
            (self.retry_attempts as f64 / self.total_transactions as f64) * 100.0
        }
    }
}

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

    #[test]
    fn test_transaction_config_default() {
        let config = TransactionConfig::default();
        assert_eq!(config.default_timeout_secs, 30);
        assert_eq!(config.max_savepoints, 10);
        assert!(config.enable_retry);
        assert_eq!(config.max_retries, 3);
    }

    #[test]
    fn test_transaction_stats_calculation() {
        let stats = TransactionStats {
            active_transactions: 2,
            total_transactions: 100,
            failed_transactions: 5,
            retry_attempts: 10,
        };

        assert_eq!(stats.success_rate(), 95.0);
        assert_eq!(stats.retry_rate(), 10.0);
    }

    #[test]
    fn test_transaction_stats_zero_transactions() {
        let stats = TransactionStats {
            active_transactions: 0,
            total_transactions: 0,
            failed_transactions: 0,
            retry_attempts: 0,
        };

        assert_eq!(stats.success_rate(), 0.0);
        assert_eq!(stats.retry_rate(), 0.0);
    }
}