use std::{
    collections::BTreeMap,
    path::PathBuf,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::Instant,
};

use backon::{ConstantBuilder, Retryable};
use log::{info, warn};
use mysql_async::{Opts, OptsBuilder, Pool, SslOpts};
use tici_system::readable::ReadableDuration;
use tokio::sync::RwLock;

use crate::config::TiDBServerConfig;

struct TiDBConnectionPoolInner {
    ssl_opts: Option<SslOpts>,
    pools: BTreeMap<String, PoolState>,
    tidb_idx: AtomicUsize,
    config: TiDBServerConfig,
}

impl TiDBConnectionPoolInner {
    pub fn append_server(&mut self, url: &str) -> Result<(), mysql_async::Error> {
        let opts = OptsBuilder::from_opts(Opts::from_url(url)?).ssl_opts(self.ssl_opts.clone());
        let pool = Pool::new(opts);
        // We don't do anything if we already has this TiDB server.
        let _ = self.pools.try_insert(
            url.to_owned(),
            PoolState {
                pool,
                error_count: 0,
                last_error: None,
                blacklisted_until: None,
            },
        );
        Ok(())
    }

    pub fn demoting_servers(&mut self, urls: &[&str]) -> Result<(), mysql_async::Error> {
        let mut demoting_pools: BTreeMap<String, PoolState> = Default::default();
        for url in urls {
            if let Some(e) = self.pools.remove(*url) {
                info!("try demote tidb server {} success", url);
                let _ = demoting_pools.insert((*url).to_owned(), e);
            } else {
                warn!("try demote tidb server {} failed", url);
            }
        }
        tokio::spawn(async {
            let mut handles = vec![];
            let count = demoting_pools.len();
            for (_, state) in demoting_pools.into_iter() {
                handles.push(state.pool.disconnect());
            }
            futures::future::join_all(handles).await;
            log::info!("finished disconnect tidb servers {}", count);
        });
        Ok(())
    }

    async fn try_get_conn(&mut self) -> Result<mysql_async::Conn, mysql_async::Error> {
        if self.pools.is_empty() {
            return Err(mysql_async::Error::Other(
                "No available connections in the pool".into(),
            ));
        }

        let pool_count = self.pools.len();
        let mut last_err = None;
        let blacklist_tidb_server_threshold = self.config.blacklist_tidb_server_threshold;
        let blacklist_duration = self.config.blacklist_duration;

        // Try each pool in round-robin, skipping blacklisted ones
        for _ in 0..pool_count {
            let idx = self.tidb_idx.fetch_add(1, Ordering::Relaxed) % pool_count;
            let (name, state) = self.pools.iter_mut().nth(idx).unwrap();

            // Check blacklist
            if state.check_blacklist_state() {
                continue;
            }

            let mock_failed: bool = (|| {
                fail::fail_point!("tidb_backend_mock_get_conn_failed", |t| {
                    t.unwrap().parse::<bool>().unwrap()
                });
                false
            })();
            let res = if mock_failed {
                Err(mysql_async::Error::Other("mocked".into()))
            } else {
                Ok(state.pool.get_conn().await?)
            };
            // Try to get a connection
            match res {
                Ok(conn) => {
                    // Success, reset error count
                    state.on_get_conn_ok();
                    return Ok(conn);
                }
                Err(e) => {
                    // Record error
                    state.on_get_conn_error(
                        name,
                        blacklist_tidb_server_threshold,
                        blacklist_duration,
                    );
                    last_err = Some(e);
                    continue;
                }
            }
        }

        // All pools failed or are blacklisted
        Err(last_err.unwrap_or_else(|| {
            mysql_async::Error::Other("No available connections in the pool".into())
        }))
    }
}

struct PoolState {
    pool: mysql_async::Pool,
    error_count: usize,
    last_error: Option<Instant>,
    blacklisted_until: Option<Instant>,
}

impl PoolState {
    pub fn on_get_conn_ok(&mut self) {
        self.error_count = 0;
        self.last_error = None;
    }

    pub fn on_get_conn_error(&mut self, name: &str, threshold: usize, duration: ReadableDuration) {
        let now = Instant::now();
        self.error_count += 1;
        self.last_error = Some(now);

        if threshold > 0 && self.error_count >= threshold {
            self.blacklisted_until = Some(now + duration.0);
            log::warn!(
                "TiDBConnectionPool: pool {} blacklisted for {:?} after {} errors",
                name,
                duration,
                self.error_count
            );
        }
    }

    pub fn check_blacklist_state(&mut self) -> bool {
        let now = Instant::now();
        if let Some(until) = self.blacklisted_until {
            if now < until {
                // Still blacklisted, skip
                true
            } else {
                // Blacklist expired, reset
                self.error_count = 0;
                self.blacklisted_until = None;
                false
            }
        } else {
            false
        }
    }
}

#[derive(Clone)]
pub struct TiDBConnectionPool {
    inner: Arc<RwLock<TiDBConnectionPoolInner>>,
}

impl TiDBConnectionPool {
    pub fn new(config: &TiDBServerConfig) -> Self {
        let mut pools = BTreeMap::<String, PoolState>::new();
        let ssl_opts = if !config.cert_path.is_empty() {
            let path_buf = PathBuf::from(std::ffi::OsString::from(&config.cert_path.to_string()));
            Some(SslOpts::default().with_root_cert_path(Some(path_buf)))
        } else {
            None
        };
        for url in config.dsns.iter() {
            let ori_opts = Opts::from_url(url).unwrap();
            let opts = OptsBuilder::from_opts(ori_opts).ssl_opts(ssl_opts.clone());
            let pool = Pool::new(opts);
            pools.insert(
                url.clone(),
                PoolState {
                    pool,
                    error_count: 0,
                    last_error: None,
                    blacklisted_until: None,
                },
            );
        }
        Self {
            inner: Arc::new(RwLock::new(TiDBConnectionPoolInner {
                ssl_opts,
                pools,
                tidb_idx: AtomicUsize::new(0),
                config: config.clone(),
            })),
        }
    }

    pub async fn append_server(&mut self, url: &str) -> Result<(), mysql_async::Error> {
        let mut inner = self.inner.write().await;
        inner.append_server(url)?;
        Ok(())
    }

    pub async fn demoting_servers(&mut self, urls: &[&str]) -> Result<(), mysql_async::Error> {
        let mut inner = self.inner.write().await;
        inner.demoting_servers(urls)?;
        Ok(())
    }

    /// Try to get a connection from a non-blacklisted pool.
    /// If a pool fails too many times, it will be blacklisted for a period.
    pub async fn get_conn(&self) -> Result<mysql_async::Conn, mysql_async::Error> {
        let (retry_count, retry_interval) = self.get_retry_config().await;

        // Check if pool is empty (configuration error) - don't retry in this case
        {
            let inner = self.inner.read().await;
            if inner.pools.is_empty() {
                return Err(mysql_async::Error::Other(
                    "No available connections in the pool".into(),
                ));
            }
        }

        let builder = ConstantBuilder::default().with_delay(retry_interval);
        let builder = if retry_count >= 0 {
            builder.with_max_times(retry_count as usize)
        } else {
            builder.without_max_times()
        };

        (move || async move {
            let mut inner = self.inner.write().await;
            inner.try_get_conn().await
        })
        .retry(&builder)
        .notify(|err, dur| {
            warn!(
                "Failed to connect to TiDB: {}, retrying in {:?}...",
                err, dur
            );
        })
        .await
    }

    pub async fn stop(&mut self) -> Result<(), mysql_async::Error> {
        let mut inner = self.inner.write().await;
        let pools = std::mem::take(&mut inner.pools);
        let mut handles = vec![];
        for (_, state) in pools.into_iter() {
            handles.push(state.pool.disconnect());
        }
        futures::future::join_all(handles).await;
        Ok(())
    }

    pub async fn get_retry_config(&self) -> (i32, std::time::Duration) {
        let inner = self.inner.read().await;
        (
            inner.config.tidb_connection_retry_count,
            inner.config.tidb_connection_retry_interval.0,
        )
    }
}