use std::sync::atomic::{AtomicUsize, Ordering};

use crate::error::ExportError;
use config::database::{DatabaseConfig, DatabaseRole as ConfigRole};
use sea_orm::prelude::async_trait::async_trait;
use sea_orm::{
    Database as SeaDatabase, DatabaseConnection, DbBackend, DbErr, ExecResult, QueryResult,
    Statement,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatabaseRole {
    Primary,
    Replica,
}

impl From<ConfigRole> for DatabaseRole {
    fn from(role: ConfigRole) -> Self {
        match role {
            ConfigRole::Primary => DatabaseRole::Primary,
            ConfigRole::Replica => DatabaseRole::Replica,
        }
    }
}

impl From<DatabaseRole> for ConfigRole {
    fn from(role: DatabaseRole) -> Self {
        match role {
            DatabaseRole::Primary => ConfigRole::Primary,
            DatabaseRole::Replica => ConfigRole::Replica,
        }
    }
}

pub trait Database: Send + Sync {
    fn config(&self) -> &DatabaseConfig;
    fn primary(&self) -> &DatabaseConnection;
    fn replicas(&self) -> &[DatabaseConnection];
    fn read(&self) -> &DatabaseConnection;
    fn write(&self) -> &DatabaseConnection;
}

impl Database for AppDatabase {
    fn config(&self) -> &DatabaseConfig {
        AppDatabase::config(self)
    }

    fn primary(&self) -> &DatabaseConnection {
        AppDatabase::primary(self)
    }

    fn replicas(&self) -> &[DatabaseConnection] {
        AppDatabase::replicas(self)
    }

    fn read(&self) -> &DatabaseConnection {
        AppDatabase::read(self)
    }

    fn write(&self) -> &DatabaseConnection {
        AppDatabase::write(self)
    }
}

pub struct AppDatabase {
    config: DatabaseConfig,
    primary: DatabaseConnection,
    replicas: Vec<DatabaseConnection>,
    round_robin: AtomicUsize,
}

impl AppDatabase {
    pub async fn connect(config: DatabaseConfig) -> Result<Self, ExportError> {
        config.validate().map_err(ExportError::InvalidConfig)?;

        let primary_opts = config
            .primary_connect_options()
            .map_err(ExportError::InvalidConfig)?;
        let primary = SeaDatabase::connect(primary_opts).await?;

        let mut replicas = Vec::with_capacity(config.replica_count());
        for i in 0..config.replica_count() {
            let opts = config
                .replica_connect_options(i)
                .map_err(ExportError::InvalidConfig)?;
            let conn = SeaDatabase::connect(opts).await?;
            replicas.push(conn);
        }

        Ok(Self {
            config,
            primary,
            replicas,
            round_robin: AtomicUsize::new(0),
        })
    }

    pub fn config(&self) -> &DatabaseConfig {
        &self.config
    }

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

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

    pub fn replica(&self, index: usize) -> Option<&DatabaseConnection> {
        self.replicas.get(index)
    }

    pub fn replica_count(&self) -> usize {
        self.replicas.len()
    }

    pub fn has_replicas(&self) -> bool {
        !self.replicas.is_empty()
    }

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

    pub fn read(&self) -> &DatabaseConnection {
        self.pick_replica().unwrap_or(&self.primary)
    }

    pub fn route_for(&self, sql: &str) -> &DatabaseConnection {
        if db::requires_primary(sql) {
            self.write()
        } else {
            self.read()
        }
    }

    fn pick_replica(&self) -> Option<&DatabaseConnection> {
        if self.replicas.is_empty() {
            return None;
        }
        let idx = self.round_robin.fetch_add(1, Ordering::Relaxed) % self.replicas.len();
        Some(&self.replicas[idx])
    }
}

#[async_trait]
impl sea_orm::ConnectionTrait for AppDatabase {
    fn get_database_backend(&self) -> DbBackend {
        self.primary.get_database_backend()
    }

    async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
        self.route_for(&stmt.sql).execute_raw(stmt).await
    }

    async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
        self.route_for(sql).execute_unprepared(sql).await
    }

    async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
        self.route_for(&stmt.sql).query_one_raw(stmt).await
    }

    async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
        self.route_for(&stmt.sql).query_all_raw(stmt).await
    }
}