// Copyright 2026 PingCAP, Inc.

use std::{
    sync::Arc,
    time::{Duration, Instant},
};

use bytes::Bytes;
use chrono::Utc;
#[cfg(not(test))]
use jemalloc_pprof::{JemallocProfCtl, PROF_CTL};
use kvengine::dfs::Dfs;
use serde::{Deserialize, Serialize};
use slog_global::{info, warn};
use tikv_util::{
    config::ReadableDuration,
    sys::{
        SysQuota, get_global_memory_usage, record_global_memory_usage,
        thread::StdThreadBuildWrapper,
    },
};

pub mod metrics;
pub use self::metrics::{MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC, MEMORY_USAGE_GAUGE_VEC};

pub const MEMORY_USAGE_REFRESH_INTERVAL: Duration = Duration::from_millis(1000);
pub const HEAP_PROFILE_PPROF_FILE_EXT: &str = "heap.pprof";
pub const HEAP_PROFILE_SVG_FILE_EXT: &str = "heap.svg";
pub const MEMORY_PROFILE_AUTO_DUMP_PATH: &str = "diagnostics/memory-profile";
pub const AUTO_DUMP_METRIC_COMPONENT_CLOUD_SERVER: &str = "cloud_server";
pub const AUTO_DUMP_METRIC_COMPONENT_CLOUD_WORKER: &str = "cloud_worker";

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct MemoryProfileAutoDumpConfig {
    pub enabled: bool,
    /// Trigger auto dump when process memory usage reaches this ratio.
    /// Value range: (0.0, 1.0], e.g. 0.8 means 80%.
    pub usage_threshold: f64,
    /// Re-arm auto dump only after usage drops to this ratio or lower.
    /// Value range: (0.0, usage_threshold), e.g. 0.5 means 50%.
    pub rearm_usage_threshold: f64,
    /// Memory usage must stay above threshold for this long before triggering
    /// dump.
    pub trigger_hold_duration: ReadableDuration,
    /// Minimum interval between two auto dumps.
    pub min_dump_interval: ReadableDuration,
}

impl Default for MemoryProfileAutoDumpConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            usage_threshold: 0.8,
            rearm_usage_threshold: 0.6,
            trigger_hold_duration: ReadableDuration::secs(10),
            min_dump_interval: ReadableDuration::minutes(10),
        }
    }
}

impl MemoryProfileAutoDumpConfig {
    pub fn validate(&self) -> Result<(), Box<dyn std::error::Error>> {
        if !(0.0..=1.0).contains(&self.usage_threshold) || self.usage_threshold == 0.0 {
            return Err(format!(
                "memory-profile-auto-dump.usage-threshold must be in (0.0, 1.0], got {}",
                self.usage_threshold
            )
            .into());
        }
        if !(0.0..self.usage_threshold).contains(&self.rearm_usage_threshold) {
            return Err(format!(
                "memory-profile-auto-dump.rearm-usage-threshold must be in (0.0, usage-threshold), got {} (usage-threshold={})",
                self.rearm_usage_threshold,
                self.usage_threshold,
            )
            .into());
        }
        Ok(())
    }
}

#[derive(Default)]
pub struct AutoDumpObserverState {
    pub dumped_while_high: bool,
    pub high_since: Option<Instant>,
    pub last_dump_at: Option<Instant>,
}

pub fn spawn_auto_dump_observer(
    component: &'static str,
    auto_dump_cfg: MemoryProfileAutoDumpConfig,
    dfs: Arc<dyn Dfs>,
) {
    let spawn_result = std::thread::Builder::new()
        .name("memory-profile-dump".to_string())
        .spawn_wrapper(move || {
            let mut state = AutoDumpObserverState::default();
            loop {
                std::thread::sleep(MEMORY_USAGE_REFRESH_INTERVAL);
                let Some((usage, memory_limit, usage_ratio)) = evaluate_auto_dump_and_prepare(
                    component,
                    &auto_dump_cfg,
                    &mut state,
                    Instant::now(),
                ) else {
                    continue;
                };

                let result = dump_heap_profile_to_dfs(dfs.clone(), component);
                match result {
                    Ok(paths) => {
                        MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
                            .with_label_values(&[component, "success"])
                            .inc();
                        info!(
                            "auto dump memory profile success";
                            "component" => component,
                            "usage" => usage,
                            "memory_limit" => memory_limit,
                            "usage_ratio" => usage_ratio,
                            "threshold" => auto_dump_cfg.usage_threshold,
                            "paths" => ?paths,
                        );
                    }
                    Err(err) => {
                        MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
                            .with_label_values(&[component, "failure"])
                            .inc();
                        warn!(
                            "auto dump memory profile failed";
                            "component" => component,
                            "usage" => usage,
                            "memory_limit" => memory_limit,
                            "usage_ratio" => usage_ratio,
                            "threshold" => auto_dump_cfg.usage_threshold,
                            "err" => err,
                        );
                    }
                }
            }
        });
    if let Err(err) = spawn_result {
        warn!(
            "spawn memory profile observer failed";
            "component" => component,
            "err" => ?err,
        );
    }
}

/// Refresh global process memory usage snapshot and return the latest value.
fn refresh_global_memory_usage(component: &'static str) -> u64 {
    record_global_memory_usage();
    let usage = get_global_memory_usage();
    MEMORY_USAGE_GAUGE_VEC
        .with_label_values(&[component])
        .set(usage as i64);
    usage
}

fn evaluate_auto_dump_and_prepare(
    component: &'static str,
    auto_dump_cfg: &MemoryProfileAutoDumpConfig,
    state: &mut AutoDumpObserverState,
    now: Instant,
) -> Option<(u64, u64, f64)> {
    let usage = refresh_global_memory_usage(component);
    let memory_limit = SysQuota::memory_limit_in_bytes();
    evaluate_auto_dump_with_snapshot(component, auto_dump_cfg, state, now, usage, memory_limit)
}

fn evaluate_auto_dump_with_snapshot(
    component: &'static str,
    auto_dump_cfg: &MemoryProfileAutoDumpConfig,
    state: &mut AutoDumpObserverState,
    now: Instant,
    usage: u64,
    memory_limit: u64,
) -> Option<(u64, u64, f64)> {
    if !auto_dump_cfg.enabled {
        MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
            .with_label_values(&[component, "skip_disabled"])
            .inc();
        return None;
    }

    if memory_limit == 0 {
        MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
            .with_label_values(&[component, "skip_no_memory_limit"])
            .inc();
        return None;
    }

    let usage_ratio = calculate_usage_ratio(usage, memory_limit);
    if usage_ratio <= auto_dump_cfg.rearm_usage_threshold {
        state.dumped_while_high = false;
    }
    if usage_ratio < auto_dump_cfg.usage_threshold {
        state.high_since = None;
        MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
            .with_label_values(&[component, "skip_below_usage_threshold"])
            .inc();
        return None;
    }
    if state.dumped_while_high {
        MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
            .with_label_values(&[component, "skip_dumped_while_high"])
            .inc();
        return None;
    }

    match state.high_since {
        Some(high_since)
            if now.duration_since(high_since) >= auto_dump_cfg.trigger_hold_duration.0 => {}
        Some(_) => {
            MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
                .with_label_values(&[component, "skip_waiting_hold_duration"])
                .inc();
            return None;
        }
        None => {
            state.high_since = Some(now);
            MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
                .with_label_values(&[component, "skip_waiting_hold_duration"])
                .inc();
            return None;
        }
    }

    if matches!(
        state.last_dump_at,
        Some(last_dump_at) if now.duration_since(last_dump_at) < auto_dump_cfg.min_dump_interval.0
    ) {
        MEMORY_PROFILE_AUTO_DUMP_ACTION_COUNTER_VEC
            .with_label_values(&[component, "skip_waiting_min_dump_interval"])
            .inc();
        return None;
    }

    state.last_dump_at = Some(now);
    state.dumped_while_high = true;
    Some((usage, memory_limit, usage_ratio))
}

#[cfg(not(test))]
fn with_jemalloc_prof_ctl<T>(
    f: impl FnOnce(&mut JemallocProfCtl) -> Result<T, String>,
) -> Result<T, String> {
    let prof_ctl_ref = std::panic::catch_unwind(|| PROF_CTL.as_ref())
        .map_err(|_| "jemalloc heap profiling is not enabled".to_owned())?;
    let Some(prof_ctl) = prof_ctl_ref else {
        return Err("jemalloc heap profiling is not enabled".to_owned());
    };
    let mut prof_ctl = prof_ctl
        .try_lock()
        .map_err(|_| "heap profiling is busy, try again later".to_owned())?;
    f(&mut prof_ctl)
}

/// Dump one heap profile as a symbolized flamegraph SVG.
#[cfg(not(test))]
pub fn dump_heap_profile_svg() -> Result<Vec<u8>, String> {
    with_jemalloc_prof_ctl(|prof_ctl| {
        prof_ctl
            .dump_flamegraph()
            .map_err(|e| format!("dump flamegraph with jemalloc_pprof fail: {}", e))
    })
}

#[cfg(test)]
pub fn dump_heap_profile_svg() -> Result<Vec<u8>, String> {
    Ok(Vec::new())
}

/// Dump one heap profile as symbolized pprof bytes.
#[cfg(not(test))]
pub fn dump_heap_profile_pprof() -> Result<Vec<u8>, String> {
    with_jemalloc_prof_ctl(|prof_ctl| {
        prof_ctl
            .dump_pprof()
            .map_err(|e| format!("dump pprof with jemalloc_pprof fail: {}", e))
    })
}

#[cfg(test)]
pub fn dump_heap_profile_pprof() -> Result<Vec<u8>, String> {
    Ok(Vec::new())
}

pub fn dump_heap_profile_to_dfs(
    dfs: Arc<dyn Dfs>,
    component: &'static str,
) -> Result<Vec<String>, String> {
    let heap_profile_pprof_data = dump_heap_profile_pprof()?;
    let heap_profile_svg_data = dump_heap_profile_svg()?;
    let upload_dfs = dfs.clone();
    let put_result = dfs.get_runtime().block_on(upload_heap_profile_to_dfs(
        upload_dfs,
        component,
        heap_profile_pprof_data,
        heap_profile_svg_data,
    ));
    put_result
}

pub fn calculate_usage_ratio(usage: u64, memory_limit: u64) -> f64 {
    if memory_limit == 0 {
        return 0.0;
    }
    usage as f64 / memory_limit as f64
}

async fn upload_heap_profile_to_dfs(
    dfs: Arc<dyn Dfs>,
    component: &'static str,
    heap_profile_pprof_data: Vec<u8>,
    heap_profile_svg_data: Vec<u8>,
) -> Result<Vec<String>, String> {
    let now = Utc::now();
    let pprof_object_key = build_heap_profile_object_key_with_ext(
        &dfs.get_prefix(),
        component,
        now,
        HEAP_PROFILE_PPROF_FILE_EXT,
    );
    dfs.put_object(
        pprof_object_key.clone(),
        Bytes::from(heap_profile_pprof_data),
        pprof_object_key.clone(),
    )
    .await
    .map_err(|err| format!("upload heap profile failed: {err:?}"))?;

    let svg_object_key = build_heap_profile_object_key_with_ext(
        &dfs.get_prefix(),
        component,
        now,
        HEAP_PROFILE_SVG_FILE_EXT,
    );
    dfs.put_object(
        svg_object_key.clone(),
        Bytes::from(heap_profile_svg_data),
        svg_object_key.clone(),
    )
    .await
    .map_err(|err| format!("upload heap profile svg failed: {err:?}"))?;
    Ok(vec![pprof_object_key, svg_object_key])
}

fn build_heap_profile_object_key_with_ext(
    dfs_prefix: &str,
    component: &'static str,
    now: chrono::DateTime<Utc>,
    file_ext: &str,
) -> String {
    let file_name = format!(
        "{}-{}.{}",
        component,
        now.format("%Y-%m-%d_%H-%M-%S"),
        file_ext
    );
    let clean_prefix = dfs_prefix.trim_end_matches('/');
    if clean_prefix.is_empty() {
        format!("{}/{}", MEMORY_PROFILE_AUTO_DUMP_PATH, file_name)
    } else {
        format!(
            "{}/{}/{}",
            clean_prefix, MEMORY_PROFILE_AUTO_DUMP_PATH, file_name
        )
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use async_trait::async_trait;
    use bytes::Bytes;
    use kvengine::dfs::{Dfs, Options};

    use super::*;

    struct FakeObjectDfs {
        runtime: tokio::runtime::Runtime,
        prefix: String,
        put_calls: Mutex<Vec<(String, Vec<u8>, String)>>,
    }

    impl FakeObjectDfs {
        fn new(prefix: &str) -> Self {
            Self {
                runtime: tokio::runtime::Builder::new_current_thread()
                    .enable_all()
                    .build()
                    .unwrap(),
                prefix: prefix.to_owned(),
                put_calls: Mutex::new(Vec::new()),
            }
        }
    }

    #[async_trait]
    impl Dfs for FakeObjectDfs {
        async fn exists(&self, _file_id: u64, _opts: Options) -> kvengine::dfs::Result<bool> {
            Err(kvengine::dfs::Error::Other("unused in test".to_owned()))
        }

        async fn size(&self, _file_id: u64, _opts: Options) -> kvengine::dfs::Result<u64> {
            Err(kvengine::dfs::Error::Other("unused in test".to_owned()))
        }

        async fn read_file(&self, _file_id: u64, _opts: Options) -> kvengine::dfs::Result<Bytes> {
            Err(kvengine::dfs::Error::Other("unused in test".to_owned()))
        }

        async fn create(
            &self,
            _file_id: u64,
            _data: Bytes,
            _opts: Options,
        ) -> kvengine::dfs::Result<()> {
            Err(kvengine::dfs::Error::Other("unused in test".to_owned()))
        }

        async fn remove(&self, _file_id: u64, _file_len: Option<u64>, _opts: Options) {
            // Not exercised by current unit tests.
        }

        async fn permanently_remove(
            &self,
            _file_id: u64,
            _opts: Options,
        ) -> kvengine::dfs::Result<()> {
            Err(kvengine::dfs::Error::Other("unused in test".to_owned()))
        }

        fn get_runtime(&self) -> &tokio::runtime::Runtime {
            &self.runtime
        }

        fn get_prefix(&self) -> String {
            self.prefix.clone()
        }

        async fn put_object(
            &self,
            key: String,
            data: Bytes,
            file_name: String,
        ) -> kvengine::dfs::Result<()> {
            self.put_calls
                .lock()
                .unwrap()
                .push((key, data.to_vec(), file_name));
            Ok(())
        }
    }

    fn test_auto_dump_cfg() -> MemoryProfileAutoDumpConfig {
        MemoryProfileAutoDumpConfig {
            enabled: true,
            usage_threshold: 0.8,
            rearm_usage_threshold: 0.5,
            trigger_hold_duration: ReadableDuration::secs(10),
            min_dump_interval: ReadableDuration::secs(60),
        }
    }

    #[test]
    fn test_auto_dump_config_validate() {
        let cfg = MemoryProfileAutoDumpConfig {
            enabled: true,
            usage_threshold: 0.6,
            rearm_usage_threshold: 0.5,
            trigger_hold_duration: ReadableDuration::secs(10),
            min_dump_interval: ReadableDuration::minutes(10),
        };
        cfg.validate()
            .expect("a valid auto dump config should pass validation");

        let invalid_cfg = MemoryProfileAutoDumpConfig {
            enabled: true,
            usage_threshold: 0.0,
            rearm_usage_threshold: 0.5,
            trigger_hold_duration: ReadableDuration::secs(10),
            min_dump_interval: ReadableDuration::minutes(10),
        };
        invalid_cfg
            .validate()
            .expect_err("usage_threshold=0.0 should fail validation");

        let invalid_rearm = MemoryProfileAutoDumpConfig {
            enabled: true,
            usage_threshold: 0.6,
            rearm_usage_threshold: 0.6,
            trigger_hold_duration: ReadableDuration::secs(10),
            min_dump_interval: ReadableDuration::minutes(10),
        };
        invalid_rearm
            .validate()
            .expect_err("rearm_usage_threshold>=usage_threshold should fail validation");
    }

    #[test]
    fn test_calculate_usage_ratio() {
        assert_eq!(calculate_usage_ratio(60, 100), 0.6);
        assert_eq!(calculate_usage_ratio(0, 0), 0.0);
    }

    #[test]
    fn test_evaluate_auto_dump_trigger_after_hold_and_skip_while_high() {
        let cfg = test_auto_dump_cfg();
        let mut state = AutoDumpObserverState::default();
        let base = Instant::now();

        // Below threshold should not start hold timer.
        assert!(
            evaluate_auto_dump_with_snapshot(
                AUTO_DUMP_METRIC_COMPONENT_CLOUD_WORKER,
                &cfg,
                &mut state,
                base,
                70,
                100,
            )
            .is_none()
        );
        assert!(state.high_since.is_none());
        assert!(!state.dumped_while_high);

        // First high usage starts the hold timer.
        assert!(
            evaluate_auto_dump_with_snapshot(
                AUTO_DUMP_METRIC_COMPONENT_CLOUD_WORKER,
                &cfg,
                &mut state,
                base + Duration::from_secs(1),
                85,
                100,
            )
            .is_none()
        );
        assert!(state.high_since.is_some());

        // Before hold duration is satisfied, still no trigger.
        assert!(
            evaluate_auto_dump_with_snapshot(
                AUTO_DUMP_METRIC_COMPONENT_CLOUD_WORKER,
                &cfg,
                &mut state,
                base + Duration::from_secs(5),
                90,
                100,
            )
            .is_none()
        );

        // After hold duration, one dump is triggered.
        let decision = evaluate_auto_dump_with_snapshot(
            AUTO_DUMP_METRIC_COMPONENT_CLOUD_WORKER,
            &cfg,
            &mut state,
            base + Duration::from_secs(11),
            90,
            100,
        );
        assert_eq!(decision, Some((90, 100, 0.9)));
        assert!(state.dumped_while_high);
        assert!(state.last_dump_at.is_some());

        // Stay high: no repeated dump until re-armed.
        assert!(
            evaluate_auto_dump_with_snapshot(
                AUTO_DUMP_METRIC_COMPONENT_CLOUD_WORKER,
                &cfg,
                &mut state,
                base + Duration::from_secs(20),
                92,
                100,
            )
            .is_none()
        );
    }

    #[test]
    fn test_evaluate_auto_dump_rearm_and_min_dump_interval() {
        let cfg = test_auto_dump_cfg();
        let mut state = AutoDumpObserverState::default();
        let base = Instant::now();

        // First trigger.
        assert!(
            evaluate_auto_dump_with_snapshot(
                AUTO_DUMP_METRIC_COMPONENT_CLOUD_SERVER,
                &cfg,
                &mut state,
                base,
                90,
                100,
            )
            .is_none()
        );
        let first = evaluate_auto_dump_with_snapshot(
            AUTO_DUMP_METRIC_COMPONENT_CLOUD_SERVER,
            &cfg,
            &mut state,
            base + Duration::from_secs(10),
            90,
            100,
        );
        assert_eq!(first, Some((90, 100, 0.9)));

        // Drop below re-arm threshold to allow next trigger.
        assert!(
            evaluate_auto_dump_with_snapshot(
                AUTO_DUMP_METRIC_COMPONENT_CLOUD_SERVER,
                &cfg,
                &mut state,
                base + Duration::from_secs(11),
                50,
                100,
            )
            .is_none()
        );
        assert!(!state.dumped_while_high);

        // Go high again within min interval: should not trigger even after hold.
        assert!(
            evaluate_auto_dump_with_snapshot(
                AUTO_DUMP_METRIC_COMPONENT_CLOUD_SERVER,
                &cfg,
                &mut state,
                base + Duration::from_secs(12),
                90,
                100,
            )
            .is_none()
        );
        assert!(
            evaluate_auto_dump_with_snapshot(
                AUTO_DUMP_METRIC_COMPONENT_CLOUD_SERVER,
                &cfg,
                &mut state,
                base + Duration::from_secs(25),
                90,
                100,
            )
            .is_none()
        );

        // After min interval passes, the existing high_since can trigger.
        let second = evaluate_auto_dump_with_snapshot(
            AUTO_DUMP_METRIC_COMPONENT_CLOUD_SERVER,
            &cfg,
            &mut state,
            base + Duration::from_secs(71),
            95,
            100,
        );
        assert_eq!(second, Some((95, 100, 0.95)));
    }

    #[test]
    fn test_build_heap_profile_object_key_trims_prefix_slash() {
        let now = Utc::now();
        let with_slash = build_heap_profile_object_key_with_ext(
            "abc/",
            "cloud_worker",
            now,
            HEAP_PROFILE_PPROF_FILE_EXT,
        );
        let without_slash = build_heap_profile_object_key_with_ext(
            "abc",
            "cloud_worker",
            now,
            HEAP_PROFILE_PPROF_FILE_EXT,
        );
        assert_eq!(with_slash, without_slash);
        assert!(with_slash.starts_with("abc/diagnostics/memory-profile/cloud_worker-"));
        assert!(with_slash.ends_with(".heap.pprof"));

        let no_prefix = build_heap_profile_object_key_with_ext(
            "",
            "cloud_server",
            now,
            HEAP_PROFILE_PPROF_FILE_EXT,
        );
        assert!(no_prefix.starts_with("diagnostics/memory-profile/cloud_server-"));
    }

    #[test]
    fn test_dump_heap_profile_to_dfs_uploads_objects() {
        let dfs = Arc::new(FakeObjectDfs::new("dfs-prefix/"));
        let paths = dump_heap_profile_to_dfs(dfs.clone(), "cloud_worker").unwrap();

        assert_eq!(paths.len(), 2);
        let pprof_path = &paths[0];
        let svg_path = &paths[1];
        assert!(pprof_path.starts_with("dfs-prefix/diagnostics/memory-profile/cloud_worker-"));
        assert!(pprof_path.ends_with(".heap.pprof"));
        assert!(svg_path.starts_with("dfs-prefix/diagnostics/memory-profile/cloud_worker-"));
        assert!(svg_path.ends_with(".heap.svg"));

        let calls = dfs.put_calls.lock().unwrap();
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].0, *pprof_path);
        assert_eq!(calls[0].2, *pprof_path);
        assert_eq!(calls[1].0, *svg_path);
        assert_eq!(calls[1].2, *svg_path);
        // In unit tests dump functions are stubbed to empty bytes.
        assert!(calls[0].1.is_empty());
        assert!(calls[1].1.is_empty());
    }
}