// Copyright 2026 PingCAP, Inc.

use std::{
    collections::HashMap, env, io::Write as _, net::TcpListener, path::PathBuf, sync::Arc,
    time::Duration,
};

use anyhow::{Context, Result, bail};
use etcd_client::{Client as EtcdClient, GetOptions};
use native_br::{
    backup::BackupConfig,
    pd_meta::{
        PdMetaProvider, PdMetaSnapshot, PdMetaWatchCache, get_prefix_range_end_for_test,
        pd_keyspace_meta_prefixes_for_test,
    },
};
use pd_client::{PdClient, pd_control::PdScheduleConfig};
use rand::{Rng, prelude::SliceRandom};
use security::{SecurityConfig, SecurityManager};
use tempfile::TempDir;
use test_cloud_server::{
    oss::prepare_dfs,
    tidb::{PD_CLIENT_UPDATE_INTERVAL, PdServerMode, PdServers},
};
use tikv_util::{config::ReadableDuration, info, time::Instant, warn};

use crate::{env_param, init_logger, new_security_config, test_id};

const PD_BIN_ENV_KEY: &str = "PD_BIN";
const PD_PORT_ENV_KEY: &str = "PD_PORT";
const TEST_DURATION_ENV_KEY: &str = "PD_META_WATCH_TEST_DURATION";
const TEST_DURATION_DEFAULT_SECS: u64 = 45;

const PD_COUNT: u16 = 1;
const PD_HEALTHY_TIMEOUT: Duration = Duration::from_secs(30);
const CACHE_READY_TIMEOUT: Duration = Duration::from_secs(30);
const CACHE_PROGRESS_TIMEOUT: Duration = Duration::from_secs(5);
const ETCD_OP_TIMEOUT: Duration = Duration::from_secs(10);
const SCAN_BATCH_LIMIT: i64 = 128;

#[test]
fn test_random_pd_meta_watch_cache() {
    init_logger();
    let prepare_time = Instant::now_coarse();
    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .worker_threads(2)
        .thread_name("pd-meta-watch-random")
        .build()
        .unwrap();

    let security_conf = new_security_config();
    let pd = PdServerGuard::new(&security_conf)
        .unwrap_or_else(|err| panic!("prepare pd-server failed: {:?}", err));
    let (_dfs_temp_dir, mut oss, dfs_config) = prepare_dfs("pd_meta_watch_");
    let config = BackupConfig {
        security: security_conf,
        dfs: dfs_config,
        backup_pd_meta_watch_cache_enabled: true,
        backup_pd_meta_watch_cache_progress_timeout: ReadableDuration(CACHE_PROGRESS_TIMEOUT),
        ..Default::default()
    };

    let stats = runtime.block_on(run_random_pd_meta_watch_cache(&pd, config));
    pd.stop_all();
    oss.shutdown();
    let stats = stats.unwrap_or_else(|err| panic!("random pd meta watch cache failed: {:?}", err));

    let stdout = std::io::stdout();
    writeln!(
        stdout.lock(),
        "{} TEST SUCCEED: elapsed {:?}, prepare {:?}, {:?}",
        test_id(),
        stats.elapsed,
        prepare_time.saturating_elapsed(),
        stats
    )
    .unwrap();
    stdout.lock().flush().unwrap();
}

async fn run_random_pd_meta_watch_cache(
    pd: &PdServerGuard,
    mut config: BackupConfig,
) -> Result<TestStats> {
    let test_duration = Duration::from_secs(env_param(
        TEST_DURATION_ENV_KEY,
        env_param("TEST_DURATION", TEST_DURATION_DEFAULT_SECS),
    ));
    pd.start().await;
    config.pd = pd_client::Config::new(pd.servers.endpoints());
    config.pd.update_interval = PD_CLIENT_UPDATE_INTERVAL;

    let pd_client = Arc::new(pd.servers.get_client()) as Arc<dyn PdClient>;
    let cluster_id = pd_client
        .get_cluster_id()
        .context("get cluster id from pd")?;
    let prefixes = pd_keyspace_meta_prefixes_for_test(cluster_id);
    let mut etcd = EtcdWorkloadClient::new(config.pd.endpoints.clone());
    let mut cache = start_watch_cache(&config, pd_client.clone())?;

    let mut rng = rand::thread_rng();
    let mut stats = TestStats::default();
    let start_time = Instant::now_coarse();
    let mut next_cache_restart = Duration::from_secs(rng.gen_range(5..=8));
    let mut next_pd_restart = Duration::from_secs(rng.gen_range(10..=14));

    while start_time.saturating_elapsed() < test_duration {
        let batch_size = rng.gen_range(1..=8);
        for _ in 0..batch_size {
            let force_in_scope = match stats.puts {
                0 => Some(true),
                1 => Some(false),
                _ => None,
            };
            let kv = random_kv(cluster_id, &prefixes, stats.puts, force_in_scope, &mut rng);
            etcd.put(kv.key, kv.value).await?;
            stats.puts += 1;
            if kv.in_scope {
                stats.in_scope_puts += 1;
            } else {
                stats.out_of_scope_puts += 1;
            }
        }

        let elapsed = start_time.saturating_elapsed();
        if elapsed >= next_cache_restart {
            info!("restart pd meta watch cache"; "elapsed" => ?elapsed);
            drop(cache);
            tokio::time::sleep(Duration::from_millis(rng.gen_range(50..=200))).await;
            cache = start_watch_cache(&config, pd_client.clone())?;
            stats.cache_restarts += 1;
            next_cache_restart += Duration::from_secs(rng.gen_range(5..=8));
        }
        if elapsed >= next_pd_restart {
            info!("restart pd-server"; "elapsed" => ?elapsed);
            pd.restart(&mut etcd, &mut rng).await;
            stats.pd_restarts += 1;
            next_pd_restart += Duration::from_secs(rng.gen_range(10..=14));
        }

        let snapshot =
            wait_snapshot_matches_scan(&cache, &config, &prefixes, CACHE_READY_TIMEOUT).await?;
        stats.verifications += 1;
        stats.last_revision = snapshot.revision;
        stats.last_snapshot_kvs = snapshot.kvs.len();

        tokio::time::sleep(Duration::from_millis(rng.gen_range(20..=100))).await;
    }

    if stats.cache_restarts == 0 {
        info!("force one pd meta watch cache restart");
        drop(cache);
        cache = start_watch_cache(&config, pd_client.clone())?;
        let snapshot =
            wait_snapshot_matches_scan(&cache, &config, &prefixes, CACHE_READY_TIMEOUT).await?;
        stats.cache_restarts += 1;
        stats.verifications += 1;
        stats.last_revision = snapshot.revision;
        stats.last_snapshot_kvs = snapshot.kvs.len();
    }
    if stats.pd_restarts == 0 {
        info!("force one pd-server restart");
        pd.restart(&mut etcd, &mut rng).await;
        let snapshot =
            wait_snapshot_matches_scan(&cache, &config, &prefixes, CACHE_READY_TIMEOUT).await?;
        stats.pd_restarts += 1;
        stats.verifications += 1;
        stats.last_revision = snapshot.revision;
        stats.last_snapshot_kvs = snapshot.kvs.len();
    }

    assert!(stats.in_scope_puts > 0, "no in-scope puts were generated");
    assert!(
        stats.out_of_scope_puts > 0,
        "no out-of-scope puts were generated"
    );
    drop(cache);

    stats.elapsed = start_time.saturating_elapsed();
    Ok(stats)
}

fn start_watch_cache(
    config: &BackupConfig,
    pd_client: Arc<dyn PdClient>,
) -> Result<Arc<PdMetaWatchCache>> {
    PdMetaWatchCache::start(config.clone(), pd_client).context("start PdMetaWatchCache")
}

async fn wait_snapshot_matches_scan(
    cache: &Arc<PdMetaWatchCache>,
    config: &BackupConfig,
    prefixes: &[Vec<u8>],
    timeout: Duration,
) -> Result<PdMetaSnapshot> {
    let begin = Instant::now_coarse();
    let mut last_err = None;
    while begin.saturating_elapsed() < timeout {
        match cache.snapshot(CACHE_PROGRESS_TIMEOUT) {
            Ok(snapshot) => {
                let scanned = scan_pd_meta_at_revision(config, prefixes, snapshot.revision).await?;
                if snapshot.kvs != scanned.kvs {
                    bail!(
                        "PdMetaSnapshot differs from full etcd scan at revision {}: {}",
                        snapshot.revision,
                        diff_kvs(&snapshot.kvs, &scanned.kvs)
                    );
                }
                return Ok(snapshot);
            }
            Err(err) => {
                last_err = Some(format!("{err:?}"));
                tokio::time::sleep(Duration::from_millis(200)).await;
            }
        }
    }
    bail!(
        "PdMetaWatchCache did not become verifiable within {:?}, last error: {:?}",
        timeout,
        last_err
    );
}

async fn scan_pd_meta_at_revision(
    config: &BackupConfig,
    prefixes: &[Vec<u8>],
    revision: i64,
) -> Result<PdMetaSnapshot> {
    let mut client = connect_etcd(&config.pd.endpoints).await?;
    let mut kvs = HashMap::new();
    for prefix in prefixes {
        let get_option = GetOptions::new()
            .with_range(get_prefix_range_end_for_test(prefix))
            .with_revision(revision)
            .with_limit(SCAN_BATCH_LIMIT);
        let mut seek_key = prefix.clone();
        loop {
            let resp = client
                .get(seek_key.clone(), Some(get_option.clone()))
                .await
                .with_context(|| format!("scan etcd prefix {}", display_key(prefix)))?;
            let mut more = resp.more();
            for kv in resp.kvs() {
                if !kv.key().starts_with(prefix) {
                    more = false;
                    break;
                }
                kvs.insert(kv.key().to_vec(), kv.value().to_vec());
            }
            if !more {
                break;
            }
            let Some(last_key) = resp.kvs().last().map(|kv| kv.key()) else {
                break;
            };
            seek_key = last_key.to_vec();
            seek_key.push(0);
        }
    }
    Ok(PdMetaSnapshot { revision, kvs })
}

async fn connect_etcd(endpoints: &[String]) -> Result<EtcdClient> {
    EtcdClient::connect(endpoints.to_vec(), None)
        .await
        .with_context(|| format!("connect etcd through {:?}", endpoints))
}

struct EtcdWorkloadClient {
    endpoints: Vec<String>,
    client: Option<EtcdClient>,
}

impl EtcdWorkloadClient {
    fn new(endpoints: Vec<String>) -> Self {
        Self {
            endpoints,
            client: None,
        }
    }

    fn reset(&mut self) {
        self.client = None;
    }

    async fn put(&mut self, key: Vec<u8>, value: Vec<u8>) -> Result<()> {
        let begin = Instant::now_coarse();
        let key_for_log = display_key(&key);
        loop {
            if self.client.is_none() {
                self.client = Some(connect_etcd(&self.endpoints).await?);
            }
            let put_result = self
                .client
                .as_mut()
                .unwrap()
                .put(key.clone(), value.clone(), None)
                .await;
            match put_result {
                Ok(_) => return Ok(()),
                Err(err) if begin.saturating_elapsed() < ETCD_OP_TIMEOUT => {
                    warn!("put random pd meta key failed, reconnect etcd";
                        "key" => key_for_log.as_str(),
                        "err" => ?err);
                    self.reset();
                    tokio::time::sleep(Duration::from_millis(100)).await;
                }
                Err(err) => {
                    return Err(err).with_context(|| {
                        format!("put random pd meta key {key_for_log} within {ETCD_OP_TIMEOUT:?}")
                    });
                }
            }
        }
    }
}

struct RandomKv {
    key: Vec<u8>,
    value: Vec<u8>,
    in_scope: bool,
}

fn random_kv<R: Rng + ?Sized>(
    cluster_id: u64,
    prefixes: &[Vec<u8>],
    seq: usize,
    force_in_scope: Option<bool>,
    rng: &mut R,
) -> RandomKv {
    let in_scope = force_in_scope.unwrap_or_else(|| rng.gen_bool(0.7));
    let mut key = if in_scope {
        let prefix = prefixes.choose(rng).unwrap();
        let mut key = prefix.clone();
        key.extend_from_slice(format!("random-watch-cache/{seq:08}/").as_bytes());
        key
    } else {
        let candidates = [
            format!("/pd/{cluster_id}/unwatched/random-watch-cache/{seq:08}/"),
            format!(
                "/pd/{}/keyspaces/random-watch-cache/{seq:08}/",
                cluster_id + 1
            ),
            format!("resource_groups/random-watch-cache/{seq:08}/"),
            format!("/tidb/remote/not_worker/random-watch-cache/{seq:08}/"),
        ];
        candidates.choose(rng).unwrap().as_bytes().to_vec()
    };
    key.extend_from_slice(format!("{:016x}", rng.r#gen::<u64>()).as_bytes());
    assert_eq!(
        prefixes.iter().any(|prefix| key.starts_with(prefix)),
        in_scope,
        "generated key scope mismatch: {}",
        display_key(&key)
    );

    let mut value = format!("value-{seq:08}-").into_bytes();
    let random_len = rng.gen_range(1..=64);
    for _ in 0..random_len {
        value.push(rng.r#gen());
    }

    RandomKv {
        key,
        value,
        in_scope,
    }
}

struct PdServerGuard {
    servers: PdServers,
    _temp_dir: TempDir,
}

impl PdServerGuard {
    fn new(security_conf: &SecurityConfig) -> Result<Self> {
        let pd_bin = PathBuf::from(env::var(PD_BIN_ENV_KEY).context("env PD_BIN is not set")?);
        let temp_dir = tempfile::Builder::new()
            .prefix("pd_meta_watch_")
            .tempdir()
            .context("create pd temp dir")?;
        let security_mgr = Arc::new(
            SecurityManager::new(security_conf)
                .map_err(|err| anyhow::anyhow!("create security manager: {err}"))?,
        );
        let servers = PdServers::new(
            PdServerMode::Normal,
            pd_bin,
            temp_dir.path().to_owned(),
            alloc_pd_port_base(),
            0,
            PdScheduleConfig::default(),
            security_mgr,
        );
        Ok(Self {
            servers,
            _temp_dir: temp_dir,
        })
    }

    async fn start(&self) {
        self.servers.start(PD_COUNT);
        self.servers.must_healthy(PD_HEALTHY_TIMEOUT).await;
    }

    async fn restart<R: Rng + ?Sized>(&self, etcd: &mut EtcdWorkloadClient, rng: &mut R) {
        self.servers.stop_all();
        etcd.reset();
        tokio::time::sleep(Duration::from_millis(rng.gen_range(100..=400))).await;
        self.start().await;
    }

    fn stop_all(&self) {
        self.servers.stop_all();
    }
}

impl Drop for PdServerGuard {
    fn drop(&mut self) {
        self.servers.stop_all();
    }
}

fn alloc_pd_port_base() -> u16 {
    if let Ok(port) = env::var(PD_PORT_ENV_KEY) {
        return port.parse().expect("PD_PORT must be a valid u16");
    }

    let mut rng = rand::thread_rng();
    for _ in 0..1000 {
        let port = rng.gen_range(20_000..60_000);
        if pd_ports_available(port) {
            return port;
        }
    }
    panic!("failed to allocate free PD port pair");
}

fn pd_ports_available(port: u16) -> bool {
    let Ok(client_listener) = TcpListener::bind(("127.0.0.1", port)) else {
        return false;
    };
    let Ok(peer_listener) = TcpListener::bind(("127.0.0.1", port + 1)) else {
        return false;
    };
    drop(peer_listener);
    drop(client_listener);
    true
}

fn diff_kvs(actual: &HashMap<Vec<u8>, Vec<u8>>, expected: &HashMap<Vec<u8>, Vec<u8>>) -> String {
    let mut missing = expected
        .keys()
        .filter(|key| !actual.contains_key(*key))
        .map(|key| display_key(key))
        .collect::<Vec<_>>();
    let mut extra = actual
        .keys()
        .filter(|key| !expected.contains_key(*key))
        .map(|key| display_key(key))
        .collect::<Vec<_>>();
    let mut different = actual
        .iter()
        .filter(|(key, value)| {
            expected
                .get(*key)
                .is_some_and(|expected| expected != *value)
        })
        .map(|(key, _)| display_key(key))
        .collect::<Vec<_>>();
    missing.sort();
    extra.sort();
    different.sort();
    missing.truncate(5);
    extra.truncate(5);
    different.truncate(5);
    format!(
        "actual_kvs={}, expected_kvs={}, missing={:?}, extra={:?}, different={:?}",
        actual.len(),
        expected.len(),
        missing,
        extra,
        different
    )
}

fn display_key(key: &[u8]) -> String {
    String::from_utf8_lossy(key).into_owned()
}

#[derive(Default, Debug)]
struct TestStats {
    elapsed: Duration,
    puts: usize,
    in_scope_puts: usize,
    out_of_scope_puts: usize,
    verifications: usize,
    cache_restarts: usize,
    pd_restarts: usize,
    last_revision: i64,
    last_snapshot_kvs: usize,
}