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

use sea_orm::{DatabaseConnection, DbErr};

#[derive(Clone)]
pub struct DbConn {
    write: DatabaseConnection,
    read: Arc<[DatabaseConnection]>,
    read_index: Arc<AtomicUsize>,
    pub(crate) replica_failure_logged: Arc<AtomicBool>,
}

impl DbConn {
    pub fn new(write: DatabaseConnection, read: Vec<DatabaseConnection>) -> Self {
        Self {
            write,
            read: read.into(),
            read_index: Arc::new(AtomicUsize::new(0)),
            replica_failure_logged: Arc::new(AtomicBool::new(false)),
        }
    }

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

    pub fn read(&self) -> &Self {
        self
    }

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

    pub(crate) fn replica(&self) -> Option<&DatabaseConnection> {
        if self.read.is_empty() {
            return None;
        }
        let index = self.read_index.fetch_add(1, Ordering::Relaxed) % self.read.len();
        Some(&self.read[index])
    }

    pub async fn close(&self) -> Result<(), DbErr> {
        let mut result = self.write.close_by_ref().await;
        for connection in self.read.iter() {
            let replica_result = connection.close_by_ref().await;
            if result.is_ok() {
                result = replica_result;
            }
        }
        result
    }
}