pub use sea_orm;

use crate::config::{AppDatabaseConnConfig, AppDatabaseSettings};
pub use sea_orm::DbErr as SqlxError;
use sea_orm::{
    ConnectOptions, ConnectionTrait, Database, DatabaseBackend, DatabaseConnection, DbErr,
    Statement,
};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use tracing::error;

#[derive(Debug, thiserror::Error)]
pub enum DbError {
    #[error("invalid database configuration: {0}")]
    InvalidConfig(String),

    #[error("failed to build primary pool: {0}")]
    PrimaryPool(#[source] DbErr),

    #[error("failed to build replica pool ({index}): {source}")]
    ReplicaPool {
        index: usize,
        #[source]
        source: DbErr,
    },

    #[error("failed to reach {pool}: {source}")]
    ConnectivityCheck {
        pool: String,
        #[source]
        source: DbErr,
    },
}

#[derive(Clone)]
pub struct Db {
    inner: Arc<DbInner>,
}

struct DbInner {
    primary: DatabaseConnection,
    replicas: Vec<DatabaseConnection>,
    rr_counter: AtomicUsize,
}

impl Db {
    pub async fn from_config(cfg: &AppDatabaseSettings) -> Result<Self, DbError> {
        validate_config(cfg)?;
        let primary = build_connection(
            &cfg.primary,
            cfg.sqlx_logging,
            cfg.query_timeout,
            cfg.slow_query_threshold_ms,
        )
        .await
        .map_err(DbError::PrimaryPool)?;
        check_connectivity(&primary, "primary").await?;

        let mut replicas = Vec::with_capacity(cfg.replicas.len());
        for (i, rc) in cfg.replicas.iter().enumerate() {
            match build_connection(
                rc,
                cfg.sqlx_logging,
                cfg.query_timeout,
                cfg.slow_query_threshold_ms,
            )
            .await
            {
                Ok(conn) => {
                    check_connectivity(&conn, &format!("replica[{i}]")).await?;
                    replicas.push(conn);
                }
                Err(e) => {
                    error!(index = i, error = %e, "failed to build replica pool");
                    return Err(DbError::ReplicaPool {
                        index: i,
                        source: e,
                    });
                }
            }
        }

        Ok(Self {
            inner: Arc::new(DbInner {
                primary,
                replicas,
                rr_counter: AtomicUsize::new(0),
            }),
        })
    }

    pub fn from_connection(connection: DatabaseConnection) -> Self {
        Self {
            inner: Arc::new(DbInner {
                primary: connection,
                replicas: Vec::new(),
                rr_counter: AtomicUsize::new(0),
            }),
        }
    }

    pub fn writer(&self) -> DatabaseConnection {
        self.inner.primary.clone()
    }

    pub fn reader(&self) -> DatabaseConnection {
        if self.inner.replicas.is_empty() {
            return self.inner.primary.clone();
        }

        let idx = next_reader_index(self.inner.replicas.len(), &self.inner.rr_counter);
        self.inner.replicas[idx].clone()
    }

    pub fn primary_connection(&self) -> &DatabaseConnection {
        &self.inner.primary
    }

    pub fn replica_connections(&self) -> &[DatabaseConnection] {
        &self.inner.replicas
    }

    pub async fn close(&self) {
        let _ = self.inner.primary.clone().close().await;
        for replica in &self.inner.replicas {
            let _ = replica.clone().close().await;
        }
    }
}

fn validate_config(cfg: &AppDatabaseSettings) -> Result<(), DbError> {
    let connections = std::iter::once(&cfg.primary).chain(&cfg.replicas);
    let mut total = 0u32;
    for (index, connection) in connections.enumerate() {
        if connection.url.is_empty() {
            return Err(DbError::InvalidConfig(format!(
                "pool {index} url must be configured"
            )));
        }
        if connection.max_connections == 0 {
            return Err(DbError::InvalidConfig(format!(
                "pool {index} max_connections must be greater than zero"
            )));
        }
        if connection.min_connections > connection.max_connections {
            return Err(DbError::InvalidConfig(format!(
                "pool {index} min_connections exceeds max_connections"
            )));
        }
        total = total
            .checked_add(connection.max_connections)
            .ok_or_else(|| {
                DbError::InvalidConfig("total connection count overflows u32".to_string())
            })?;
    }
    if total > cfg.max_connections_total {
        return Err(DbError::InvalidConfig(format!(
            "configured pools require {total} connections, exceeding max_connections_total {}",
            cfg.max_connections_total
        )));
    }
    Ok(())
}

#[allow(clippy::cast_possible_truncation)]
async fn build_connection(
    cfg: &AppDatabaseConnConfig,
    sqlx_logging: bool,
    query_timeout: u64,
    slow_query_threshold_ms: u64,
) -> Result<DatabaseConnection, DbErr> {
    let app_name = cfg.application_name.clone();
    let schema = cfg.schema.clone();

    let mut opts = ConnectOptions::new(&cfg.url);

    opts.connect_lazy(true);

    opts.max_connections(cfg.max_connections)
        .min_connections(cfg.min_connections);

    let acquire_timeout = cfg.acquire_timeout.unwrap_or(5);
    opts.acquire_timeout(Duration::from_secs(acquire_timeout));
    if let Some(lt) = cfg.max_lifetime {
        opts.max_lifetime(Duration::from_secs(lt));
    }
    if let Some(it) = cfg.idle_timeout {
        opts.idle_timeout(Duration::from_secs(it));
    }

    opts.sqlx_logging(sqlx_logging);
    opts.test_before_acquire(true);

    if sqlx_logging && slow_query_threshold_ms > 0 {
        opts.sqlx_slow_statements_logging_settings(
            log::LevelFilter::Warn,
            Duration::from_millis(slow_query_threshold_ms),
        );
    }

    opts.map_sqlx_postgres_opts(move |pg_opts: sea_orm::sqlx::postgres::PgConnectOptions| {
        let mut pg_opts = pg_opts;
        if let Some(ref name) = app_name {
            pg_opts = pg_opts.application_name(name);
        }
        if let Some(ref sch) = schema {
            pg_opts = pg_opts.options([("search_path", sch.as_str())]);
        }
        if query_timeout > 0 {
            let stmt_timeout = format!("{query_timeout}s");
            pg_opts = pg_opts.options([("statement_timeout", stmt_timeout.as_str())]);
        }
        pg_opts
    });

    Database::connect(opts).await
}

async fn check_connectivity(conn: &DatabaseConnection, pool: &str) -> Result<(), DbError> {
    let stmt = Statement::from_string(DatabaseBackend::Postgres, "SELECT 1");
    conn.execute_raw(stmt)
        .await
        .map(|_| ())
        .map_err(|source| DbError::ConnectivityCheck {
            pool: pool.to_string(),
            source,
        })
}

fn next_reader_index(len: usize, counter: &AtomicUsize) -> usize {
    debug_assert!(len > 0);
    counter.fetch_add(1, Ordering::Relaxed) % len
}

pub async fn rollback_or_log(txn: sea_orm::DatabaseTransaction) {
    if let Err(error) = txn.rollback().await {
        error!(error = %error, "transaction_rollback_failed");
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{AppDatabaseConnConfig, AppDatabaseSettings};

    fn conn_config(url: &str, max: u32, min: u32) -> AppDatabaseConnConfig {
        AppDatabaseConnConfig {
            url: url.to_string(),
            max_connections: max,
            min_connections: min,
            ..AppDatabaseConnConfig::default()
        }
    }

    fn settings(
        primary: AppDatabaseConnConfig,
        replicas: Vec<AppDatabaseConnConfig>,
        total: u32,
    ) -> AppDatabaseSettings {
        AppDatabaseSettings {
            primary,
            replicas,
            max_connections_total: total,
            ..AppDatabaseSettings::default()
        }
    }

    #[test]
    fn validate_config_rejects_empty_url() {
        let cfg = settings(conn_config("", 20, 2), vec![], 100);
        let err = validate_config(&cfg).unwrap_err().to_string();
        assert!(err.contains("url must be configured"), "{err}");
    }

    #[test]
    fn validate_config_rejects_zero_max_connections() {
        let cfg = settings(conn_config("postgres://db", 0, 0), vec![], 100);
        let err = validate_config(&cfg).unwrap_err().to_string();
        assert!(
            err.contains("max_connections must be greater than zero"),
            "{err}"
        );
    }

    #[test]
    fn validate_config_rejects_min_over_max() {
        let cfg = settings(conn_config("postgres://db", 5, 10), vec![], 100);
        let err = validate_config(&cfg).unwrap_err().to_string();
        assert!(err.contains("min_connections exceeds"), "{err}");
    }

    #[test]
    fn validate_config_rejects_total_overflow() {
        let primary = conn_config("postgres://db", u32::MAX, 2);
        let replica = conn_config("postgres://db2", u32::MAX, 2);
        let cfg = settings(primary, vec![replica], u32::MAX);
        let err = validate_config(&cfg).unwrap_err().to_string();
        assert!(err.contains("overflows u32"), "{err}");
    }

    #[test]
    fn validate_config_rejects_total_exceeding_budget() {
        let primary = conn_config("postgres://db", 80, 2);
        let replica = conn_config("postgres://db2", 30, 2);
        let cfg = settings(primary, vec![replica], 100);
        let err = validate_config(&cfg).unwrap_err().to_string();
        assert!(err.contains("exceeding max_connections_total"), "{err}");
    }

    #[test]
    fn validate_config_accepts_valid_configuration() {
        let primary = conn_config("postgres://db", 20, 2);
        let replica = conn_config("postgres://db2", 10, 2);
        let cfg = settings(primary, vec![replica], 100);
        assert!(validate_config(&cfg).is_ok());
    }

    #[test]
    fn next_reader_index_round_robins() {
        let counter = AtomicUsize::new(0);
        let idxs: Vec<usize> = (0..6).map(|_| next_reader_index(3, &counter)).collect();
        assert_eq!(idxs, vec![0, 1, 2, 0, 1, 2]);
    }

    #[test]
    fn next_reader_index_single_pool_is_stable() {
        let counter = AtomicUsize::new(0);
        let idxs: Vec<usize> = (0..5).map(|_| next_reader_index(1, &counter)).collect();
        assert_eq!(idxs, vec![0, 0, 0, 0, 0]);
    }
}