use anyhow::Result;
use async_trait::async_trait;
use tici_common::index::IndexKey;

#[async_trait]
pub trait FragIDGenerator: Send + Sync {
    async fn gen_id(&self) -> anyhow::Result<u64>;
}

pub struct TSOGenerator {
    #[cfg(feature = "with_pd_client")]
    pd_client: Box<dyn pd_client::PdClient>,
}

impl TSOGenerator {
    pub async fn new(
        endpoints: Vec<String>,
        security_config: tici_common::security::SecurityConfig,
    ) -> anyhow::Result<Self> {
        #[cfg(feature = "with_pd_client")]
        {
            let pd_client = Self::create_pd_client(endpoints, security_config).await?;
            Ok(Self { pd_client })
        }

        #[cfg(not(feature = "with_pd_client"))]
        {
            anyhow::bail!("Enable feature 'with_pd_client' to use TSOGenerator");
        }
    }

    pub async fn gen_tso(&self) -> anyhow::Result<u64> {
        #[cfg(feature = "with_pd_client")]
        {
            let tso = self.pd_client.get_tso().await?;
            Ok(tso.into_inner())
        }

        #[cfg(not(feature = "with_pd_client"))]
        {
            anyhow::bail!("Enable feature 'with_pd_client' to use TSOGenerator");
        }
    }

    #[cfg(feature = "with_pd_client")]
    async fn create_pd_client(
        endpoints: Vec<String>,
        security_config: tici_common::security::SecurityConfig,
    ) -> anyhow::Result<Box<dyn pd_client::PdClient>> {
        use anyhow::Context;
        use log::info;
        let pd_config = pd_client::Config {
            endpoints,
            ..Default::default()
        };

        info!("create_pd_client, cfg={:?}", security_config);
        let pd_sec_cfg = pd_security::SecurityConfig {
            ca_path: security_config.ca_path.clone(),
            cert_path: security_config.cert_path.clone(),
            key_path: security_config.key_path.clone(),
            ..Default::default()
        };
        let security_mgr = pd_security::SecurityManager::new(&pd_sec_cfg)
            .map_err(|e| anyhow::Error::msg(e.to_string()))
            .with_context(|| "Failed to create SecurityManager")?;
        let security_mgr = std::sync::Arc::new(security_mgr);

        let pd_client = pd_client::RpcClient::new_async(&pd_config, None, security_mgr)
            .await
            .with_context(|| {
                format!(
                    "Failed to create pd client, addr={}",
                    pd_config.endpoints.join(",")
                )
            })?;

        info!("create_pd_client success");
        Ok(Box::new(pd_client))
    }
}

#[async_trait]
impl FragIDGenerator for TSOGenerator {
    async fn gen_id(&self) -> anyhow::Result<u64> {
        self.gen_tso().await
    }
}

pub struct FragPathGenerator {
    frag_base_path: String,
    frag_id_gen: Box<dyn FragIDGenerator>,
}

impl FragPathGenerator {
    pub fn new(frag_base_path: &str, frag_id_gen: Box<dyn FragIDGenerator>) -> Self {
        Self {
            frag_base_path: frag_base_path.to_string(),
            frag_id_gen,
        }
    }

    // Frag path format: {local_data_path}/t_{table_id}/i_{index_id}/f_{tso}
    // - Separate the frag data by table_id and index_id to facilitate DROP INDEX
    //   and DROP TABLE operations.
    // - In cases like index creation where shard IDs may not exist, we don't add
    //   shard_id to path.
    pub async fn new_path(&self, index_key: &IndexKey) -> Result<String> {
        let frag_id = self.frag_id_gen.gen_id().await?;
        Ok(format!(
            "{}/{}/f_{}",
            self.frag_base_path,
            index_key.to_s3_prefix(),
            frag_id
        ))
    }

    pub fn frag_base_path(&self) -> &str {
        &self.frag_base_path
    }

    pub fn extract_timestamp(path: &str) -> Option<u64> {
        let parts: Vec<&str> = path.split('/').collect();
        if parts.len() < 5 {
            return None;
        }
        if parts.get(2).unwrap().starts_with("ks_") {
            // {prefix}/frags/ks_1/t_110/i_1001/f_xxx
            if parts.len() < 6 {
                return None;
            }
            let ts = match parts[5].replace("f_", "").parse::<u64>() {
                Ok(ts) => ts,
                Err(_) => return None,
            };

            Some(ts)
        } else {
            // {prefix}/frags/t_110/i_1001/f_xxx
            let ts = match parts[4].replace("f_", "").parse::<u64>() {
                Ok(ts) => ts,
                Err(_) => return None,
            };

            Some(ts)
        }
    }

    pub fn extract_base(path: &str) -> Option<String> {
        let parts: Vec<&str> = path.split('/').collect();
        if parts.len() < 5 {
            return None;
        }
        if parts.get(2).unwrap().starts_with("ks_") {
            // {prefix}/frags/ks_1/t_110/i_1001/f_xxx
            if parts.len() < 6 {
                return None;
            }
            Some(parts[..6].join("/"))
        } else {
            Some(parts[..5].join("/"))
        }
    }

    pub fn path_is_fragment_file(path: &str) -> bool {
        let parts: Vec<&str> = path.split('/').filter(|x| !x.is_empty()).collect();
        if parts.len() < 5 {
            return false;
        }
        if parts.get(2).unwrap().starts_with("ks_") {
            parts.len() >= 6
        } else {
            parts.len() >= 5
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::LocalTSGenerator;

    #[tokio::test]
    async fn test_path_gen() {
        let generator = LocalTSGenerator::new();
        let path_gen = FragPathGenerator::new("aaa/frags", Box::new(generator));
        let new_path = path_gen
            .new_path(&IndexKey::new_nullspace(1, 2))
            .await
            .unwrap();
        assert!(
            FragPathGenerator::extract_base(new_path.as_str())
                .unwrap()
                .starts_with("aaa/frags/t_1/i_2/")
        );
        assert!(FragPathGenerator::extract_timestamp(new_path.as_str()).unwrap() > 0);
    }

    #[test]
    fn test_extract() {
        assert_eq!(
            FragPathGenerator::extract_timestamp(
                "testprefix/frags/ks_1/t_110/i_1001/f_123456/merged.1"
            )
            .unwrap(),
            123456
        );
        assert_eq!(
            FragPathGenerator::extract_timestamp(
                "tici_default_prefix/frags/t_127/i_2/f_462251698917736457/multi.0"
            )
            .unwrap(),
            462251698917736457
        );
        assert_eq!(
            FragPathGenerator::extract_base(
                "tici_default_prefix/frags/t_127/i_2/f_462251698917736457/multi.0"
            )
            .unwrap(),
            "tici_default_prefix/frags/t_127/i_2/f_462251698917736457"
        );
        assert_eq!(
            FragPathGenerator::extract_base(
                "tici_default_prefix/frags/ks_1/t_127/i_2/f_462251698917736457/multi.0"
            )
            .unwrap(),
            "tici_default_prefix/frags/ks_1/t_127/i_2/f_462251698917736457"
        );
        assert!(FragPathGenerator::path_is_fragment_file(
            "tici_default_prefix/frags/t_127/i_2/f_462251698917736457/multi.0"
        ),);
        assert!(FragPathGenerator::path_is_fragment_file(
            "tici_default_prefix/frags/ks_1/t_127/i_2/f_462251698917736457/multi.0"
        ));
        assert!(!FragPathGenerator::path_is_fragment_file(
            "tici_default_prefix/frags/ks_1/t_127/i_2/"
        ),);
    }
}