// Copyright 2024 TiKV Project Authors. Licensed under Apache-2.0.

use tikv_alloc::error::ProfResult;

use crate::status_server::jemalloc_utils::{get_malloc_stats, issue_mallctl, issue_mallctl_args};

pub fn activate_prof() -> ProfResult<()> {
    {
        log::info!("activate_prof");
        let mut value: bool = true;
        let len = std::mem::size_of_val(&value) as u64;
        issue_mallctl_args(
            "prof.active",
            std::ptr::null_mut(),
            std::ptr::null_mut(),
            &mut value as *mut _ as *mut _,
            len,
        );
    }
    Ok(())
}

pub fn has_activate_prof() -> bool {
    let mut value: bool = false;
    let mut len = std::mem::size_of_val(&value) as u64;
    issue_mallctl_args(
        "prof.active",
        &mut value as *mut _ as *mut _,
        &mut len as *mut _ as *mut _,
        std::ptr::null_mut(),
        0,
    );
    value
}

pub fn deactivate_prof() -> ProfResult<()> {
    {
        log::info!("deactivate_prof");
        let mut value: bool = false;
        let len = std::mem::size_of_val(&value) as u64;
        issue_mallctl_args(
            "prof.active",
            std::ptr::null_mut(),
            std::ptr::null_mut(),
            &mut value as *mut _ as *mut _,
            len,
        );
    }
    Ok(())
}

extern crate libc;

#[allow(unused_variables)]
fn print_error_message(s: &str, r: ::std::os::raw::c_int) {
    #[cfg(not(any(target_os = "android", target_os = "dragonfly", target_os = "macos")))]
    {
        unsafe {
            let err = *libc::__errno_location();
            let err_msg = libc::strerror(err);
            let c_str = std::ffi::CStr::from_ptr(err_msg);
            let str_slice = c_str.to_str().unwrap_or("Unknown error");
            log::warn!(
                "{} returns non-zero {} error_code={} error_message={}",
                s,
                r,
                err,
                str_slice
            );
        }
    }
}

pub fn dump_prof(path: &str) -> tikv_alloc::error::ProfResult<()> {
    {
        let mut bytes = std::ffi::CString::new(path)?.into_bytes_with_nul();
        let mut ptr = bytes.as_mut_ptr() as *mut ::std::os::raw::c_char;
        let len = std::mem::size_of_val(&ptr) as u64;
        let r = issue_mallctl_args(
            "prof.dump",
            std::ptr::null_mut(),
            std::ptr::null_mut(),
            &mut ptr as *mut _ as *mut _,
            len,
        );
        if r != 0 {
            print_error_message("dump_prof", r);
        }
    }
    Ok(())
}

pub fn adhoc_dump(path: &str) -> tikv_alloc::error::ProfResult<()> {
    {
        let mut bytes = std::ffi::CString::new(path)?.into_bytes_with_nul();
        let mut ptr = bytes.as_mut_ptr() as *mut ::std::os::raw::c_char;
        let len = std::mem::size_of_val(&ptr) as u64;
        let r = issue_mallctl_args(
            "prof.dump",
            std::ptr::null_mut(),
            std::ptr::null_mut(),
            &mut ptr as *mut _ as *mut _,
            len,
        );
        if r != 0 {
            print_error_message("adhoc_dump", r);
        }
    }
    Ok(())
}

pub fn jeprof_purge_arena() {
    let narenas = issue_mallctl("arenas.narenas");
    log::info!("jeprof_purge_arena purge {} arenas", narenas);
    for i in 0..narenas {
        let a_string = format!("arena.{}.purge", i);
        let r = issue_mallctl_args(
            &a_string,
            std::ptr::null_mut(),
            std::ptr::null_mut(),
            std::ptr::null_mut(),
            0,
        );
        if r != 0 {
            log::warn!("jeprof_purge_arena purge {} return {}", a_string, r);
        }
    }
    log::info!("jeprof_purge_arena purge {} arenas done", narenas);
}

pub fn jeprof_memory_status() -> String {
    get_malloc_stats()
}

#[cfg(target_os = "linux")] // Only support linux for pprof heap test
#[test]
fn test_adhoc_dump_prof() {
    use tempfile::NamedTempFile;
    let prev_has_activate_prof = has_activate_prof();
    if !prev_has_activate_prof {
        let _ = activate_prof();
    }

    let _x = vec![1; 1000];
    let _y = vec![1; 1000];

    let f = NamedTempFile::new().unwrap();
    let path = f.path().to_str().unwrap();
    std::thread::sleep(std::time::Duration::from_millis(1000));
    adhoc_dump(path).unwrap();
    let target_path = std::path::Path::new(path);
    assert!(target_path.exists());
    log::info!(
        "std::fs::metadata(path).unwrap().len() {}",
        std::fs::metadata(path).unwrap().len()
    );
    assert!(std::fs::metadata(path).unwrap().len() > 100);

    if !prev_has_activate_prof {
        let _ = deactivate_prof();
    }
}
