use std::fmt::{Debug, Display};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering;

use sea_orm::prelude::async_trait::async_trait;
use sea_orm::{
    AccessMode, ConnectionTrait, DatabaseConnection, DatabaseTransaction, DbBackend, DbErr,
    ExecResult, IsolationLevel, QueryResult, Statement, TransactionError, TransactionOptions,
    TransactionTrait,
};

use crate::conn::DbConn;
use crate::routing;

impl DbConn {
    async fn route<'a, T, F, Fut>(&'a self, sql: &str, operation: F) -> Result<T, DbErr>
    where
        F: Fn(&'a DatabaseConnection) -> Fut + Send,
        Fut: Future<Output = Result<T, DbErr>> + Send,
        T: Send,
    {
        if !routing::requires_primary(sql)
            && let Some(replica) = self.replica()
        {
            match operation(replica).await {
                Ok(result) => return Ok(result),
                // 副本失败时静默回落会掩盖副本故障，把首次失败记进日志。
                Err(error) => {
                    if self
                        .replica_failure_logged
                        .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
                        .is_ok()
                    {
                        tracing::warn!(
                            %error,
                            "replica query failed; falling back to the primary"
                        );
                    }
                }
            }
        }
        operation(self.write()).await
    }
}

#[async_trait]
impl ConnectionTrait for DbConn {
    fn get_database_backend(&self) -> DbBackend {
        self.write().get_database_backend()
    }

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

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

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

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

    fn support_returning(&self) -> bool {
        self.write().support_returning()
    }

    fn is_mock_connection(&self) -> bool {
        self.write().is_mock_connection()
    }
}

#[async_trait]
impl TransactionTrait for DbConn {
    type Transaction = DatabaseTransaction;

    async fn begin(&self) -> Result<Self::Transaction, DbErr> {
        self.write().begin().await
    }

    async fn begin_with_config(
        &self,
        isolation_level: Option<IsolationLevel>,
        access_mode: Option<AccessMode>,
    ) -> Result<Self::Transaction, DbErr> {
        self.write()
            .begin_with_config(isolation_level, access_mode)
            .await
    }

    async fn begin_with_options(
        &self,
        options: TransactionOptions,
    ) -> Result<Self::Transaction, DbErr> {
        self.write().begin_with_options(options).await
    }

    async fn transaction<F, T, E>(&self, callback: F) -> Result<T, TransactionError<E>>
    where
        F: for<'c> FnOnce(
                &'c Self::Transaction,
            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'c>>
            + Send,
        T: Send,
        E: Display + Debug + Send,
    {
        self.write().transaction(callback).await
    }

    async fn transaction_with_config<F, T, E>(
        &self,
        callback: F,
        isolation_level: Option<IsolationLevel>,
        access_mode: Option<AccessMode>,
    ) -> Result<T, TransactionError<E>>
    where
        F: for<'c> FnOnce(
                &'c Self::Transaction,
            ) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'c>>
            + Send,
        T: Send,
        E: Display + Debug + Send,
    {
        self.write()
            .transaction_with_config(callback, isolation_level, access_mode)
            .await
    }
}