// Copyright 2025 PingCAP, Inc.

use std::sync::Arc;

use async_trait::async_trait;
use pd_client::PdClient;
use security::SecurityConfig;

use crate::{
    KeyspaceService, KeyspaceStates, ReplicationWorkerConfig, Result, bootstrap,
    util::new_keyspace_pd_client,
};

pub(crate) struct KeyspaceProvisionedService {
    keyspace_id: u32,
    conf: ReplicationWorkerConfig,
    sec_conf: SecurityConfig,
    pub(crate) states: KeyspaceStates,
    pd_client: Option<Arc<dyn PdClient>>,
}

impl KeyspaceProvisionedService {
    pub(crate) fn new(
        keyspace_id: u32,
        conf: &ReplicationWorkerConfig,
        sec_conf: &SecurityConfig,
        states: KeyspaceStates,
    ) -> Self {
        Self {
            keyspace_id,
            conf: conf.clone(),
            sec_conf: sec_conf.clone(),
            states,
            pd_client: None,
        }
    }
}

#[async_trait]
impl KeyspaceService for KeyspaceProvisionedService {
    fn keyspace_id(&self) -> u32 {
        self.keyspace_id
    }

    async fn start(&mut self) -> Result<()> {
        let pd_url = self.states.pd_url.clone();
        let pd_client = new_keyspace_pd_client(pd_url, &self.sec_conf).await;
        bootstrap(
            pd_client.clone(),
            self.conf.merged_engine.merged_store_id,
            self.conf.advertise_addr.clone(),
        )
        .await?;
        self.pd_client = Some(pd_client);
        Ok(())
    }

    async fn destroy(&mut self) -> Result<()> {
        Ok(())
    }

    fn get_states(&self) -> &KeyspaceStates {
        &self.states
    }

    fn get_states_mut(&mut self) -> &mut KeyspaceStates {
        &mut self.states
    }

    fn get_pd_client(&self) -> Arc<dyn PdClient> {
        self.pd_client.as_ref().unwrap().clone()
    }
}

#[cfg(feature = "testexport")]
pub mod local_provider {
    use std::{fmt::Display, fs, path::PathBuf, sync::Arc, time::Duration};

    use futures::executor::block_on;
    use pd_client::pd_control::PdControl;
    use security::{RestfulClient, SecurityConfig, SecurityManager};
    use serde_derive::{Deserialize, Serialize};
    use tikv_util::{
        box_err, info,
        retry::try_wait_result_async,
        sys::process::{
            ChildSupervisorHandle, ProcessSpawnSpec, StopReason, SupervisorResult, spawn_supervisor,
        },
    };

    use crate::Result;

    const LOCAL_PROVIDER_OWNER: &str = "local provider";

    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
    enum ProcessKind {
        Pd,
        Cdc,
        Tidb,
    }

    impl ProcessKind {
        fn process_name(self) -> &'static str {
            match self {
                Self::Pd => "rep-pd",
                Self::Cdc => "rep-cdc",
                Self::Tidb => "rep-tidb",
            }
        }

        fn thread_name(self) -> &'static str {
            match self {
                Self::Pd => "rep-proc-pd",
                Self::Cdc => "rep-proc-cdc",
                Self::Tidb => "rep-proc-tidb",
            }
        }

        fn spawn_spec(
            self,
            program: String,
            args: Vec<String>,
            log_file: PathBuf,
        ) -> ProcessSpawnSpec {
            ProcessSpawnSpec::new(
                LOCAL_PROVIDER_OWNER,
                self.process_name(),
                self.thread_name(),
                program,
                args,
                log_file,
            )
        }
    }

    /// Test-only process helper used by random replication tests in provisioned
    /// mode.
    pub struct LocalProvider {
        pub keyspace_id: u32,
        pub data_dir: PathBuf,
        pd_bin_path: String,
        cdc_bin_path: String,
        tidb_bin_path: String,
        base_port: u16,
        log_level: String, // E.g., "info", "debug".
        cdc_newarch: bool,

        pd: Option<ChildSupervisorHandle>,
        cdc: Option<ChildSupervisorHandle>,
        tidb: Option<ChildSupervisorHandle>,

        security_mgr: Arc<SecurityManager>,
    }

    impl LocalProvider {
        /// Creates a new local provider using the binary paths from the
        /// environment.
        pub fn new(
            keyspace_id: u32,
            data_dir: PathBuf,
            base_port: u16,
            log_level: String,
            cdc_newarch: bool,
        ) -> Self {
            let security_mgr = SecurityManager::new(&SecurityConfig::default()).unwrap();
            Self {
                keyspace_id,
                data_dir,
                pd_bin_path: std::env::var("REP_PD_BIN").unwrap(),
                cdc_bin_path: std::env::var("CDC_BIN").unwrap(),
                tidb_bin_path: std::env::var("TIDB_BIN").unwrap_or_default(),
                base_port,
                log_level,
                cdc_newarch,
                pd: None,
                cdc: None,
                tidb: None,
                security_mgr: Arc::new(security_mgr),
            }
        }

        /// Starts the local PD, TiCDC, and TiDB helper processes and waits for
        /// PD/TiDB health.
        ///
        /// Callers must install `tikv_util::set_panic_hook(false, ...)`
        /// first so watcher-thread failures terminate the whole test process
        /// instead of only surfacing later as channel or join errors.
        pub fn start(&mut self, timeout: Duration) {
            assert!(
                self.pd.is_none() && self.cdc.is_none() && self.tidb.is_none(),
                "LocalProvider::start called twice",
            );

            // PD -> CDC -> TiDB.
            self.pd = Some(unwrap_or_panic(self.spawn_pd_supervisor()));
            self.cdc = Some(unwrap_or_panic(self.spawn_cdc_supervisor()));
            self.tidb = Some(unwrap_or_panic(self.spawn_tidb_supervisor()));

            let wait_healthy = async {
                let (pd_res, tidb_res) = tokio::join!(
                    self.wait_pd_healthy(timeout),
                    self.wait_tidb_healthy(timeout)
                );
                pd_res.expect("Wait for PD health timeout");
                tidb_res.expect("Wait for TiDB health timeout");
            };
            block_on(wait_healthy);
        }

        fn pd_spawn_spec(&self) -> ProcessSpawnSpec {
            let name = format!("rep-pd-{}", self.keyspace_id);
            let data_dir = self.data_dir.join(&name);
            fs::create_dir_all(&data_dir).unwrap();
            let log_file = self.data_dir.join(format!("{}.log", name));
            let initial_cluster = format!("pd={}", self.local_pd_peer_url());
            let config_file = self.data_dir.join(format!("{}.toml", name));
            let pd_config = "[replication]\nmax-replicas = 1\n";
            fs::write(&config_file, pd_config).unwrap();
            ProcessKind::Pd.spawn_spec(
                self.pd_bin_path.clone(),
                vec![
                    "--name=pd".to_string(),
                    format!("--data-dir={}", data_dir.display()),
                    format!("--log-file={}", log_file.display()),
                    format!("--log-level={}", self.log_level),
                    format!("--client-urls={}", self.pd_client_url()),
                    format!("--peer-urls={}", self.local_pd_peer_url()),
                    format!("--initial-cluster={initial_cluster}"),
                    format!("--config={}", config_file.display()),
                ],
                log_file,
            )
        }

        fn cdc_spawn_spec(&self) -> ProcessSpawnSpec {
            let keyspace_id = self.keyspace_id;
            let data_dir = self.data_dir.join(format!("cdc-{keyspace_id}"));
            fs::create_dir_all(&data_dir).unwrap();
            let log_file = self.data_dir.join(format!("cdc-{keyspace_id}.log"));
            let mut args = vec![
                "server".to_string(),
                format!("--addr={}", self.cdc_server_addr()),
                format!("--log-file={}", log_file.display()),
                format!("--log-level={}", self.log_level),
                format!("--data-dir={}", data_dir.display()),
                format!("--pd={}", self.pd_client_url()),
            ];
            if self.cdc_newarch {
                args.push("--newarch".to_string());
            }
            ProcessKind::Cdc.spawn_spec(self.cdc_bin_path.clone(), args, log_file)
        }

        fn tidb_spawn_spec(&self) -> ProcessSpawnSpec {
            let keyspace_id = self.keyspace_id;
            let data_dir = self.data_dir.join(format!("tidb-{keyspace_id}"));
            fs::create_dir_all(&data_dir).unwrap();
            let log_file = self.data_dir.join(format!("rep-tidb-{keyspace_id}.log"));
            let slow_log_file = self
                .data_dir
                .join(format!("rep-tidb-slow-{keyspace_id}.log"));
            let config_file = self.data_dir.join(format!("rep-tidb-{keyspace_id}.toml"));
            let config = LocalTidbConfig::default();
            fs::write(&config_file, toml::to_string(&config).unwrap()).unwrap();
            ProcessKind::Tidb.spawn_spec(
                self.tidb_bin_path.clone(),
                vec![
                    format!("--P={}", self.local_tidb_port()),
                    format!("--log-file={}", log_file.display()),
                    format!("--log-slow-query={}", slow_log_file.display()),
                    "--store=unistore".to_string(),
                    format!("--path={}", data_dir.display()),
                    format!("--status={}", self.local_tidb_status_port()),
                    format!("--config={}", config_file.display()),
                ],
                log_file,
            )
        }

        fn spawn_pd_supervisor(&self) -> SupervisorResult<ChildSupervisorHandle> {
            spawn_supervisor(self.pd_spawn_spec())
        }

        fn spawn_cdc_supervisor(&self) -> SupervisorResult<ChildSupervisorHandle> {
            spawn_supervisor(self.cdc_spawn_spec())
        }

        fn spawn_tidb_supervisor(&self) -> SupervisorResult<ChildSupervisorHandle> {
            spawn_supervisor(self.tidb_spawn_spec())
        }

        /// Returns the replication PD client URL.
        pub fn pd_client_url(&self) -> String {
            Self::local_proc_url(self.local_pd_client_port())
        }

        /// Returns the TiCDC listen address.
        pub fn cdc_server_addr(&self) -> String {
            Self::local_proc_addr(self.local_cdc_server_port())
        }

        fn local_pd_peer_url(&self) -> String {
            Self::local_proc_url(self.local_pd_peer_port())
        }

        fn local_cdc_server_port(&self) -> u16 {
            self.base_port + 2000 + self.keyspace_id as u16
        }

        fn local_pd_peer_port(&self) -> u16 {
            self.base_port + 1000 + self.keyspace_id as u16
        }

        fn local_pd_client_port(&self) -> u16 {
            self.base_port + self.keyspace_id as u16
        }

        fn local_proc_addr(port: u16) -> String {
            format!("127.0.0.1:{}", port)
        }

        fn local_proc_url(port: u16) -> String {
            format!("http://{}", Self::local_proc_addr(port))
        }

        fn local_tidb_port(&self) -> u16 {
            self.base_port + 3000 + self.keyspace_id as u16
        }

        fn local_tidb_status_port(&self) -> u16 {
            self.base_port + 4000 + self.keyspace_id as u16
        }

        /// Stops all helper processes in `TiDB -> CDC -> PD` order.
        ///
        /// Any stop-time failure is treated as fatal and panics immediately so
        /// the panic hook can terminate the whole test process.
        pub fn destroy(&mut self) -> Result<()> {
            info!("destroy local provider");
            if let Some(tidb) = self.tidb.take() {
                unwrap_or_panic(tidb.stop_and_join(StopReason::Destroy));
            }
            if let Some(cdc) = self.cdc.take() {
                unwrap_or_panic(cdc.stop_and_join(StopReason::Destroy));
            }
            if let Some(pd) = self.pd.take() {
                unwrap_or_panic(pd.stop_and_join(StopReason::Destroy));
            }
            Ok(())
        }

        /// Restarts only the local replication PD process.
        ///
        /// Any restart-time failure is treated as fatal and panics immediately.
        pub fn restart_local_pd(&mut self) -> Result<()> {
            let pd = self
                .pd
                .as_ref()
                .unwrap_or_else(|| panic!("rep-pd is not started"));
            unwrap_or_panic(pd.request_restart());
            block_on(self.wait_pd_healthy(Duration::from_secs(10)))?;
            Ok(())
        }

        /// Returns a PD control client for the local replication PD.
        ///
        /// Panics if the provider has not been started yet.
        pub fn get_pd_control(&self) -> PdControl {
            let endpoints = vec![self.pd_client_url()];
            let cfg = pd_client::Config::new(endpoints);
            let mgr = self.security_mgr.clone();
            PdControl::new(cfg, mgr).unwrap()
        }

        /// Waits until the local replication PD health endpoint reports
        /// healthy.
        pub async fn wait_pd_healthy(&self, timeout: Duration) -> Result<()> {
            let pd_ctl = Arc::new(self.get_pd_control());
            try_wait_result_async(
                || {
                    let pd_ctl = pd_ctl.clone();
                    Box::pin(async move {
                        match pd_ctl.health().await {
                            Ok(health) if health => Ok(()),
                            Ok(_) => Err("PD unhealthy".to_string()),
                            Err(e) => {
                                let err = Err(format!("check PD healthy failed: {:?}", e));
                                info!("{:?}", err);
                                err
                            }
                        }
                    })
                },
                timeout,
                || Duration::from_millis(200),
            )
            .await
            .map_err(|s| box_err!("{}", s))
        }

        /// Waits until the local TiDB health endpoint reports healthy.
        pub async fn wait_tidb_healthy(&self, timeout: Duration) -> Result<()> {
            let status_url = Self::local_proc_url(self.local_tidb_status_port());
            let tidb_ctl = Arc::new(
                RestfulClient::new("tidb-ctl", vec![status_url], self.security_mgr.clone())
                    .unwrap(),
            );
            try_wait_result_async(
                || {
                    let tidb_ctl = tidb_ctl.clone();
                    Box::pin(async move {
                        match tidb_ctl.get::<TidbHealth>("health").await {
                            Ok(health) if health.status == "up" => Ok(()),
                            Ok(_) => Err("TiDB is not healthy".into()),
                            Err(e) => Err(format!("failed to get TiDB health: {:?}", e)),
                        }
                    })
                },
                timeout,
                || Duration::from_millis(200),
            )
            .await
            .map_err(|s| box_err!("{}", s))
        }
    }

    // Ref: https://github.com/tidbcloud/tidb-cse/blob/release-7.1-keyspace/server/health_handler.go
    // {"status":"up","token":"keyspace_a"}
    #[derive(Default, Deserialize, Debug)]
    #[serde(default)]
    struct TidbHealth {
        status: String, // "up" or "down"
        token: String,  // keyspace name
    }

    #[derive(Default, Serialize)]
    #[serde(rename_all = "kebab-case")]
    struct LocalTidbConfig {
        split_table: bool, // Set to false.
        security: TidbConfigSecurity,
    }

    #[derive(Default, Serialize)]
    #[serde(rename_all = "kebab-case")]
    struct TidbConfigSecurity {
        enable_sem: bool, // Disable "sem" to update `mysql.tidb`.
    }

    fn unwrap_or_panic<T, E>(res: std::result::Result<T, E>) -> T
    where
        E: Display,
    {
        match res {
            Ok(value) => value,
            Err(err) => panic!("{err}"),
        }
    }
}