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

//! This module provides unified APIs for accessing thread/process related
//! information. Only Linux platform is implemented correctly, for other
//! platform, it only guarantees successful compilation.

use std::{collections::HashMap, io, io::Result, thread};

use parking_lot::Mutex;
use tikv_alloc::{
    add_thread_memory_accessor, remove_thread_memory_accessor, thread_allocate_exclusive_arena,
};

/// A cross-platform CPU statistics data structure.
#[derive(Debug, Copy, Clone, Default, PartialEq)]
pub struct ThreadStat {
    // libc::clock_t is not used here because the definition of
    // clock_t is different on linux and bsd.
    pub s_time: i64,
    pub u_time: i64,
}

impl ThreadStat {
    /// Gets the cpu time in seconds.
    #[inline]
    pub fn total_cpu_time(&self) -> f64 {
        cpu_total(self.s_time, self.u_time)
    }
}

#[inline]
fn cpu_total(sys_time: i64, user_time: i64) -> f64 {
    (sys_time + user_time) as f64 / ticks_per_second() as f64
}

#[cfg(target_os = "linux")]
pub mod linux {
    #[inline]
    pub fn cpu_total(stat: &super::FullStat) -> f64 {
        super::cpu_total(stat.stime, stat.utime)
    }
}

#[cfg(target_os = "linux")]
mod imp {
    use std::{
        fs,
        io::{self, Error},
        iter::FromIterator,
    };

    use libc::c_int;
    pub use libc::pid_t as Pid;
    pub use procinfo::pid::{self, Stat as FullStat};

    lazy_static::lazy_static! {
        // getconf CLK_TCK
        static ref CLOCK_TICK: i64 = {
            unsafe {
                libc::sysconf(libc::_SC_CLK_TCK)
            }
        };

        static ref PROCESS_ID: Pid = unsafe { libc::getpid() };
    }

    #[inline]
    pub fn ticks_per_second() -> i64 {
        *CLOCK_TICK
    }

    /// Gets the ID of the current process.
    #[inline]
    pub fn process_id() -> Pid {
        *PROCESS_ID
    }

    /// Gets the ID of the current thread.
    #[inline]
    pub fn thread_id() -> Pid {
        thread_local! {
            static TID: Pid = unsafe { libc::syscall(libc::SYS_gettid) as Pid };
        }
        TID.with(|t| *t)
    }

    /// Gets thread ids of the given process id.
    /// WARN: Don't call this function frequently. Otherwise there will be a lot
    /// of memory fragments.
    pub fn thread_ids<C: FromIterator<Pid>>(pid: Pid) -> io::Result<C> {
        let dir = fs::read_dir(format!("/proc/{}/task", pid))?;
        Ok(dir
            .filter_map(|task| {
                let file_name = match task {
                    Ok(t) => t.file_name(),
                    Err(e) => {
                        log::error!("read task failed, pid {}, err {:?}", pid, e);
                        return None;
                    }
                };

                match file_name.to_str() {
                    Some(tid) => match tid.parse() {
                        Ok(tid) => Some(tid),
                        Err(e) => {
                            log::error!("read task failed, pid {}, err {:?}", pid, e);
                            None
                        }
                    },
                    None => {
                        log::error!("read task failed, pid {}", pid);
                        None
                    }
                }
            })
            .collect())
    }

    pub fn full_thread_stat(pid: Pid, tid: Pid) -> io::Result<FullStat> {
        pid::stat_task(pid, tid)
    }

    pub fn set_priority(pri: i32) -> io::Result<()> {
        // Unsafe due to FFI.
        unsafe {
            let tid = libc::syscall(libc::SYS_gettid);
            if libc::setpriority(libc::PRIO_PROCESS, tid as u32, pri) != 0 {
                let e = Error::last_os_error();
                return Err(e);
            }
            Ok(())
        }
    }

    pub fn get_priority() -> io::Result<i32> {
        // Unsafe due to FFI.
        unsafe {
            let tid = libc::syscall(libc::SYS_gettid);
            clear_errno();
            let ret = libc::getpriority(libc::PRIO_PROCESS, tid as u32);
            if ret == -1 {
                let e = Error::last_os_error();
                if let Some(errno) = e.raw_os_error() {
                    if errno != 0 {
                        return Err(e);
                    }
                }
            }
            Ok(ret)
        }
    }

    // Sadly the std lib does not have any support for setting `errno`, so we
    // have to implement this ourselves.
    extern "C" {
        #[link_name = "__errno_location"]
        fn errno_location() -> *mut c_int;
    }

    fn clear_errno() {
        // Unsafe due to FFI.
        unsafe {
            *errno_location() = 0;
        }
    }
}

#[cfg(target_os = "macos")]
#[allow(bad_style)]
mod imp {
    use std::{io, iter::FromIterator, mem::size_of, ptr::null_mut, slice};

    use libc::*;

    pub type Pid = mach_port_t;

    type task_inspect_t = mach_port_t;
    type thread_act_t = mach_port_t;
    type thread_act_array_t = *mut thread_act_t;

    extern "C" {
        fn task_threads(
            target_task: task_inspect_t,
            act_list: *mut thread_act_array_t,
            act_listCnt: *mut mach_msg_type_number_t,
        ) -> kern_return_t;
    }

    const MICRO_SEC_PER_SEC: i64 = 1_000_000;

    #[derive(Default)]
    pub struct FullStat {
        pub stime: i64,
        pub utime: i64,
        pub command: String,
    }

    /// Unlike Linux, the unit of `stime` and `utime` is microseconds instead of
    /// ticks. See [`full_thread_stat()`]
    #[inline]
    pub fn ticks_per_second() -> i64 {
        MICRO_SEC_PER_SEC
    }

    /// Gets the ID of the current process.
    #[inline]
    pub fn process_id() -> Pid {
        unsafe { mach2::traps::mach_task_self() }
    }

    /// Gets the ID of the current thread.
    #[inline]
    pub fn thread_id() -> Pid {
        unsafe { mach2::mach_init::mach_thread_self() }
    }

    pub fn thread_ids<C: FromIterator<Pid>>(pid: Pid) -> io::Result<C> {
        // https://www.gnu.org/software/hurd/gnumach-doc/Task-Information.html

        unsafe {
            let mut act_list: thread_act_array_t = null_mut();
            let mut act_count: mach_msg_type_number_t = 0;
            let ret = task_threads(pid, &mut act_list, &mut act_count);
            if ret != KERN_SUCCESS {
                return Err(io::Error::from_raw_os_error(ret));
            }

            let pids = slice::from_raw_parts_mut(act_list, act_count as _)
                .iter()
                .copied()
                .collect();

            vm_deallocate(
                pid,
                act_list as vm_address_t,
                size_of::<thread_act_t>() * act_count as usize,
            );

            Ok(pids)
        }
    }

    pub fn full_thread_stat(_pid: Pid, tid: Pid) -> io::Result<FullStat> {
        // https://www.gnu.org/software/hurd/gnumach-doc/Thread-Information.html

        unsafe {
            let flavor = THREAD_BASIC_INFO;
            let mut info = std::mem::zeroed::<thread_basic_info>();
            let mut thread_info_cnt = THREAD_BASIC_INFO_COUNT;

            let ret = thread_info(
                tid,
                flavor as task_flavor_t,
                (&mut info as *mut _) as thread_info_t,
                &mut thread_info_cnt,
            );
            if ret != KERN_SUCCESS {
                return Err(io::Error::from_raw_os_error(ret));
            }

            Ok(FullStat {
                stime: info.system_time.seconds as i64 * 1_000_000
                    + info.system_time.microseconds as i64,
                utime: info.user_time.seconds as i64 * 1_000_000
                    + info.user_time.microseconds as i64,
                ..Default::default()
            })
        }
    }

    pub fn set_priority(_: i32) -> io::Result<()> {
        Ok(())
    }

    pub fn get_priority() -> io::Result<i32> {
        Ok(0)
    }
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
mod imp {
    use std::{io, iter::FromIterator};

    pub type Pid = u32;

    #[derive(Default)]
    pub struct FullStat {
        pub stime: i64,
        pub utime: i64,
        pub command: String,
    }

    lazy_static::lazy_static! {
        static ref PROCESS_ID: Pid = std::process::id();
    }

    #[inline]
    pub fn ticks_per_second() -> i64 {
        1
    }

    /// Gets the ID of the current process.
    #[inline]
    pub fn process_id() -> Pid {
        *PROCESS_ID
    }

    /// Gets the ID of the current thread.
    #[inline]
    pub fn thread_id() -> Pid {
        u64::from(std::thread::current().id().as_u64()) as Pid
    }

    pub fn thread_ids<C: FromIterator<Pid>>(_pid: Pid) -> io::Result<C> {
        Ok(std::iter::once(thread_id()).collect())
    }

    pub fn full_thread_stat(_pid: Pid, _tid: Pid) -> io::Result<FullStat> {
        Ok(FullStat::default())
    }

    pub fn set_priority(_: i32) -> io::Result<()> {
        Ok(())
    }

    pub fn get_priority() -> io::Result<i32> {
        Ok(0)
    }
}

pub use self::imp::*;

pub fn thread_stat(pid: Pid, tid: Pid) -> io::Result<ThreadStat> {
    let full_stat = full_thread_stat(pid, tid)?;
    Ok(ThreadStat {
        s_time: full_stat.stime as _,
        u_time: full_stat.utime as _,
    })
}

pub fn current_thread_stat() -> io::Result<ThreadStat> {
    thread_stat(process_id(), thread_id())
}

pub trait StdThreadBuildWrapper {
    fn spawn_wrapper<F, T>(self, f: F) -> io::Result<thread::JoinHandle<T>>
    where
        F: FnOnce() -> T,
        F: Send + 'static,
        T: Send + 'static;
}

pub trait ThreadBuildWrapper {
    /// Register all system hooks along with a custom hook pair.
    fn with_sys_and_custom_hooks<F1, F2>(&mut self, after_start: F1, before_end: F2) -> &mut Self
    where
        F1: Fn() + Send + Sync + 'static,
        F2: Fn() + Send + Sync + 'static;

    /// Register some generic hooks like memory tracing or thread lifetime
    /// tracing.
    fn with_sys_hooks(&mut self) -> &mut Self {
        self.with_sys_and_custom_hooks(|| {}, || {})
    }
}

lazy_static::lazy_static! {
    pub static ref THREAD_NAME_HASHMAP: Mutex<HashMap<Pid, String>> = Mutex::new(HashMap::default());
    pub static ref THREAD_START_HOOKS: Mutex<Vec<Box<dyn Fn() + Sync + Send>>> = Mutex::new(Vec::new());
}

pub fn hook_thread_start(f: Box<dyn Fn() + Sync + Send>) {
    THREAD_START_HOOKS.lock().push(f);
}

pub(crate) fn call_thread_start_hooks() {
    for f in THREAD_START_HOOKS.lock().iter() {
        f();
    }
}

pub(crate) fn add_thread_name_to_map() {
    if let Some(name) = std::thread::current().name() {
        let tid = thread_id();
        THREAD_NAME_HASHMAP.lock().insert(tid, name.to_string());
        log::debug!("tid {} thread name is {}", tid, name);
    }
}

pub(crate) fn remove_thread_name_from_map() {
    let tid = thread_id();
    THREAD_NAME_HASHMAP.lock().remove(&tid);
}

impl StdThreadBuildWrapper for std::thread::Builder {
    fn spawn_wrapper<F, T>(self, f: F) -> Result<std::thread::JoinHandle<T>>
    where
        F: FnOnce() -> T,
        F: Send + 'static,
        T: Send + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.spawn(|| {
            call_thread_start_hooks();
            // SAFETY: we will call `remove_thread_memory_accessor` at defer.
            unsafe { add_thread_memory_accessor() };
            thread_allocate_exclusive_arena().unwrap();
            add_thread_name_to_map();
            scopeguard::defer! {{
                remove_thread_name_from_map();
                remove_thread_memory_accessor();
            }};
            f()
        })
    }
}

impl ThreadBuildWrapper for tokio::runtime::Builder {
    fn with_sys_and_custom_hooks<F1, F2>(&mut self, start: F1, end: F2) -> &mut Self
    where
        F1: Fn() + Send + Sync + 'static,
        F2: Fn() + Send + Sync + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.on_thread_start(move || {
            call_thread_start_hooks();
            // SAFETY: we will call `remove_thread_memory_accessor` at
            // `before-stop_wrapper`.
            // FIXME: What if the user only calls `after_start_wrapper`?
            unsafe { add_thread_memory_accessor() };
            thread_allocate_exclusive_arena().unwrap();
            add_thread_name_to_map();
            start();
        })
        .on_thread_stop(move || {
            end();
            remove_thread_name_from_map();
            remove_thread_memory_accessor();
        })
    }
}

impl ThreadBuildWrapper for futures::executor::ThreadPoolBuilder {
    fn with_sys_and_custom_hooks<F1, F2>(&mut self, start: F1, end: F2) -> &mut Self
    where
        F1: Fn() + Send + Sync + 'static,
        F2: Fn() + Send + Sync + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.after_start(move |_| {
            call_thread_start_hooks();
            // SAFETY: we will call `remove_thread_memory_accessor` at
            // `before-stop_wrapper`.
            // FIXME: What if the user only calls `after_start_wrapper`?
            unsafe { add_thread_memory_accessor() };
            thread_allocate_exclusive_arena().unwrap();
            add_thread_name_to_map();
            start();
        })
        .before_stop(move |_| {
            end();
            remove_thread_name_from_map();
            remove_thread_memory_accessor();
        })
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::HashSet, sync, sync::Arc};

    use futures::executor::block_on;
    use parking_lot::{Condvar, Mutex};

    use super::*;
    // use crate::yatp_pool::{DefaultTicker, YatpPoolBuilder};

    #[test]
    fn test_thread_id() {
        let id = thread_id();
        assert_ne!(id, 0);
        std::thread::spawn(move || {
            // Two threads should have different ids.
            assert_ne!(thread_id(), id);
        })
        .join()
        .unwrap();
    }

    #[test]
    // TODO: this is not a stable test under macOS
    // #[cfg(any(target_os = "linux", target_os = "macos"))]
    #[cfg(target_os = "linux")]
    fn test_thread_ids() {
        const THREAD_COUNT: usize = 10;
        let (tx, rx) = crossbeam::channel::bounded(THREAD_COUNT);

        #[allow(clippy::mutex_atomic)]
        let stop_threads_cvar = Arc::new((Mutex::new(false), Condvar::new()));

        let threads = (0..THREAD_COUNT)
            .map(|_| {
                let tx = tx.clone();
                let stop_threads_cvar = stop_threads_cvar.clone();
                std::thread::spawn(move || {
                    tx.send(thread_id()).unwrap();

                    let (lock, cvar) = &*stop_threads_cvar;
                    let mut guard = lock.lock();
                    let _guard = cvar.wait_while(&mut guard, |stop: &mut bool| !*stop);
                })
            })
            .collect::<Vec<_>>();

        let actual_tids = rx.iter().take(THREAD_COUNT).collect::<Vec<_>>();
        let ids = thread_ids::<HashSet<_>>(process_id()).unwrap();
        for tid in &actual_tids {
            assert!(ids.contains(tid));
        }

        {
            let (lock, cvar) = &*stop_threads_cvar;
            *lock.lock() = true;
            cvar.notify_all();
        }
        for thread in threads {
            thread.join().unwrap();
        }

        let stopped_tids = actual_tids;
        let ids = thread_ids::<HashSet<_>>(process_id()).unwrap();
        for tid in &stopped_tids {
            assert!(
                !ids.contains(tid),
                "ids={:?} tid={:?} stopped_tids={:?}",
                ids,
                tid,
                stopped_tids
            );
        }
    }

    #[test]
    fn test_thread_name_wrapper() {
        let thread_name = "thread_for_test";

        let (tx, rx) = sync::mpsc::sync_channel(10);

        let get_name = move || {
            let tid = thread_id();
            if let Some(name) = THREAD_NAME_HASHMAP.lock().get(&tid) {
                tx.clone().send(name.to_string()).unwrap();
            } else {
                panic!("thread not found");
            }
        };

        // test std thread builder
        std::thread::Builder::new()
            .name(thread_name.to_string())
            .spawn_wrapper(get_name.clone())
            .unwrap()
            .join()
            .unwrap();

        let name = rx.recv().unwrap();
        assert_eq!(name, thread_name);

        // test Yatp
        // let get_name_fn = get_name.clone();
        // block_on(
        //     YatpPoolBuilder::new(DefaultTicker {})
        //         .name_prefix(thread_name)
        //         .after_start(|| {})
        //         .before_stop(|| {})
        //         .build_future_pool()
        //         .spawn_handle(async move { get_name_fn() })
        //         .unwrap(),
        // )
        // .unwrap();

        // let name = rx.recv().unwrap();
        // assert!(name.contains(thread_name));

        // test tokio thread builder
        let get_name_fn = get_name;
        block_on(
            tokio::runtime::Builder::new_multi_thread()
                .thread_name(thread_name)
                .with_sys_hooks()
                .build()
                .unwrap()
                .spawn(async move { get_name_fn() }),
        )
        .unwrap();

        let name = rx.recv().unwrap();
        assert_eq!(name, thread_name);
    }
}