cpumap 0.2.1

use crate::ordered_core_selections::OrderedCoreSelections;
use std::time::Duration;
use std::cell::RefCell;
use std::collections::HashMap;
use std::collections::HashSet;
use crate::util::CpuSetExt;
use crate::strings;
use sysinfo::UpdateKind;
use anyhow;
use fuzzy_matcher::FuzzyMatcher;
use fuzzy_matcher::skim::SkimMatcherV2;
use crate::topologycache::TopologyCache;
use log::*;

/// Information about a physical CPU (not a core or a logical core, but a full CPU possibly with multiple cores).
pub struct CPUInfo<'a> {
    /// Human-readable CPU name.
    pub name: &'a str
}

/// 'Thread orders' allow the user to automatically assign threads to cores with a single button click.
///
/// Each thread order specifies how to assign cores to individual threads.
#[derive(PartialEq, Clone, Copy, Debug)]
pub enum ThreadOrder {
    /// Bind threads to consecutive logical cores, one per thread.
    ByLCore,
    /// Bind threads to consecutive physical cores, each thread with all logical cores in a physical core
    ByCore,
    /// Bind threads to consecutive L3 caches, each thread to all logical cores under its L3 cache.
    ByL3,
}
impl ThreadOrder {
    /// 'Rotate' a thread order to the next value, or None added to the list.
    ///
    /// Used for TUI thread order 'radios'.
    pub fn rotate(order: &mut Option<ThreadOrder>) -> Option<ThreadOrder> {
        match order  {
            Some(Self::ByLCore) => *order = Some(Self::ByCore),
            Some(Self::ByCore)  => *order = Some(Self::ByL3),
            Some(Self::ByL3)    => *order = None,
            None          => *order = Some(Self::ByLCore),
        }
        *order
    }
}

/// Provides information about the system we're running on, and access to change that system (e.g.
/// control CPU affinities).
pub struct System {
    /// Portable access to process information (names and PIDs) and CPU affinity control.
    system: sysinfo::System,
    /// Stores cached information about topology objects to avoid recursive recomputation.
    cache:  TopologyCache,
    /// CPU (including CPU cache) topology.
    topo:   hwloc2::Topology,
    /// WORKAROUND: (TID, PID) pairs of threads that couldn't be found in `system.processes()`.
    ///
    /// May be a `sysinfo` bug?
    ///
    /// We store them here to avoid spamming the log with errors.
    missing_threads: RefCell<HashSet<(sysinfo::Pid, sysinfo::Pid)>>
}

impl System {
    /// Construct a System given system info an hardware topology.
    pub fn new(system: sysinfo::System, cache: TopologyCache, topo: hwloc2::Topology) -> System {
        System{
            system, 
            cache,
            topo,
            missing_threads: RefCell::new(HashSet::new())
        }
    }

    /// Get the logical core ('hardware thread') count of the system.
    pub fn logical_core_count(&self) -> usize {
        self.cache.cpusets_pu().len()
    }

    /// Get cached topology object information.
    pub fn cache(&self) -> &TopologyCache {
        &self.cache
    }

    /// Get a CpuSet with currently allowed cores for given process/thread ID (depending on `bind`).
    ///
    /// Note that the main thread has the same TID as the process  PID.
    ///
    /// Note that this directly gets data from the process in real time; even if the process was
    /// alive during the last refresh, it may not be alive *now* and this may fail.
    ///
    /// Returns `None` if the process is dead.
    pub fn get_cpu_affinity(&self, pid: sysinfo::Pid, bind: hwloc2::CpuBindFlags) -> Option<hwloc2::CpuSet>{
        self.topo.get_cpubind_for_process(pid.as_u32() as i32, bind)
    }

    /// Construct a CpuSet containing every PU on current system.
    pub fn affinity_full(&self) -> hwloc2::CpuSet {
        hwloc2::CpuSet::from_range(0, (self.logical_core_count() - 1).try_into().unwrap())
    }

    /// Set cores to allow process with given ID to run on.
    ///
    /// Returns error on failure to set affinity (e.g. due to permissions).
    pub fn set_cpu_affinity(&mut self, pid: sysinfo::Pid, bind: hwloc2::CpuBindFlags, cpuset: hwloc2::CpuSet)
        -> anyhow::Result<()> 
    {
        info!("set_cpu_affinity({}, {:?}, {})", pid, bind, cpuset);
        if let Err(e) = self.topo.set_cpubind_for_process(pid.as_u32() as i32, cpuset, bind) {
            return Err(anyhow::anyhow!("Failed to set CPU affinity for {} with flags {:?}: {:?}", pid, bind, e));
        } else {
            Ok(())
        }
    }

    /// Set CPU affinities for threads of a process (not the process itself).
    ///
    /// `affinities` are thread affinities indexed by Tid.
    ///
    /// Returns error on failure to set affinity (e.g. due to permissions).
    pub fn set_cpu_affinity_threads(&mut self, affinities: &HashMap<sysinfo::Pid, hwloc2::CpuSet>)
        -> anyhow::Result<()> 
    {
        for (tid, affinity) in affinities {
            self.set_cpu_affinity( *tid, hwloc2::CpuBindFlags::CPUBIND_THREAD, affinity.clone() )?;
        }
        Ok(())
    }

    /// Run a command in a shell and then set the affinity of the command's process(es).
    ///
    /// Sets affinties of: 
    /// - if `subprocess_pattern` is empty, any processes directly started by the shell command.
    /// - if `subprocess_pattern` is non-empty, any processes recursively started by the shell 
    ///   command with names/command-lines fuzzily matching `subprocess_pattern`
    ///
    /// Returns Pids of processes that had their affinity set.
    pub fn run_shell_command_and_set_process_affinity(
        &mut                    self,
        selected_cores_process: &hwloc2::CpuSet,
        command:                &String,
        subprocess_pattern:     &String ) -> anyhow::Result<Vec<sysinfo::Pid>>
    {
        // Recursively get any children of `shell_pid` that match `subprocess_pattern`
        fn get_matching_children(this: &System, ppid: sysinfo::Pid, subprocess_pattern: &str) -> Vec::<sysinfo::Pid> {
            let mut result = Vec::new();
            let matcher = SkimMatcherV2::default();
            for pid in this.get_processes_by_parent(ppid) {
                let process = this.get_process(pid)
                              .expect("no process info update since `get_processes_by_parent` - process must exist");

                let name = System::get_process_name(process);
                // In weird cases like AppImages, the process will not have the expected name
                // either in name() or exe(), but may still have that name in cmd(), so match both.
                let cmd_name = format!("{:?}", process.cmd());
                debug!("checking match with child process id {}, name '{}', cmd '{}'", pid, name, cmd_name);

                let cmd_match = matcher.fuzzy_match(cmd_name.as_str(), subprocess_pattern).is_some();
                let name_match = matcher.fuzzy_match(name, subprocess_pattern).is_some();
                if cmd_match || name_match {
                    result.push(pid);
                }
                result.extend(get_matching_children(this, pid, subprocess_pattern));
            }
            result
        }
        fn shell_err(shell: &mut std::process::Child) -> String {
            use std::io::Read;
            let mut err_buf = String::new();
            // Ignore any output read errors, will not append anything to err_buf.
            shell.stdout.as_mut().expect("should be piped").read_to_string(&mut err_buf).ok();
            shell.stderr.as_mut().expect("should be piped").read_to_string(&mut err_buf).ok();
            err_buf.truncate(256);
            err_buf
        }

        let mut shell = std::process::Command::new("sh")
                .arg("-c")
                .arg(command)
                .stdout(std::process::Stdio::piped())
                .stderr(std::process::Stdio::piped())
                .spawn()?;
        let shell_pid = sysinfo::Pid::from_u32(shell.id());
        // Shouldn't cause *too much* of user-visible lag.
        std::thread::sleep(Duration::from_millis(50));
        self.refresh_process_info();
        if self.get_process(shell_pid).is_none() {
            return Err(anyhow::anyhow!("After running command '{}': the process is not running; {}",
                                       command, shell_err(&mut shell)));
        }
        let shell_children = if subprocess_pattern.is_empty() {
            self.get_processes_by_parent(shell_pid)
        } else {
            get_matching_children(self, shell_pid, subprocess_pattern.as_str())
        };
        if shell_children.is_empty() {
            let err_base = if subprocess_pattern.is_empty() {
                "Shell didn't launch any process / took too long to launch?"
            } else {
                "No subprocess with matching name."
            };
            return Err(anyhow::anyhow!("{}; {}", err_base, shell_err(&mut shell)));
        } 

        // Set affinity of all spawned processes (usually just one process).
        for child_id in shell_children.iter() {
            if let Err(error) = self.set_cpu_affinity( *child_id, hwloc2::CpuBindFlags::CPUBIND_PROCESS, selected_cores_process.clone() ) {
                let child = self.get_process(*child_id)
                                .expect("no process info update since we got the children, this process should exist");
                return Err(anyhow::anyhow!("failed to set process affinity for the launched process {}/{} (permissions?): {:?}", 
                                           child_id, System::get_process_name(child), error));
            }
        }

        Ok(shell_children)
    }

    /// Update process information - has to be called sparingly to keep overhead low.
    pub fn refresh_process_info(&mut self) {
        info!("refreshing process info");
        self.system.refresh_processes_specifics(
            sysinfo::ProcessRefreshKind::new().with_cmd(UpdateKind::OnlyIfNotSet)
                                              .with_exe(UpdateKind::OnlyIfNotSet));
    }

    /// Get name of a process.
    ///
    /// Should be used for consistency instead of `sysinfo::Process::name()`.
    ///
    /// Tries to use process binary filename; if this is not possible, falls back to linux process name.
    pub fn get_process_name(process: &sysinfo::Process) -> &str {
        let opt_name = process.exe().and_then(|exe| exe.file_name().and_then(|name| name.to_str()));
        opt_name.unwrap_or(process.name())
    }

    /// Get processes matching given `filter` string.
    pub fn get_process_filter_matches<'a>(&'a self, filter: &str, allow_kernel: bool)
        -> Vec<(&'a sysinfo::Pid, &'a sysinfo::Process)> 
    {
        let matcher = SkimMatcherV2::default();
        let process_filter = |(pid,process): &(&sysinfo::Pid, &sysinfo::Process)| {
            // TODO: this is linux-specific and probably not even fully exhaustive. Improve the
            // implementation to work for other OS's
            if process.thread_kind().is_some() {
                return false;
            }
            if !allow_kernel {
                if cfg!(target_os = "linux") {
                    let kthreadd  = sysinfo::Pid::from(2);
                    let is_kernel = **pid == kthreadd || process.parent() == Some(kthreadd);
                    if is_kernel {
                        return false;
                    }
                }
            }
            let name = Self::get_process_name(process);
            return matcher.fuzzy_match(name, &filter).is_some() 
        };
        // Filter processes by `filter` and sort them by fuzzy match score.
        let mut matches: Vec<(&sysinfo::Pid, &sysinfo::Process)> = self.system.processes().iter().filter(process_filter).collect();
        matches.sort_unstable_by(|(_, process_a), (_, process_b)| {
            let name_a = Self::get_process_name(process_a);
            let name_b = Self::get_process_name(process_b);
            let (score_a, _) = matcher.fuzzy_indices(name_a, filter).unwrap();
            let (score_b, _) = matcher.fuzzy_indices(name_b, filter).unwrap();
            score_b.partial_cmp(&score_a).unwrap()
        });
        matches
    }

    /// Get process with specified PID, if it exists.
    pub fn get_process(&self, pid: sysinfo::Pid) -> Option<&sysinfo::Process> {
        self.system.processes().get(&pid)
    }

    /// Get Pids of all processes spawned by process with pid `ppid`.
    pub fn get_processes_by_parent(&self, ppid: sysinfo::Pid) -> Vec<sysinfo::Pid> {
        return self.system.processes()
                          .iter()
                          .filter(|(_pid, p)| p.parent() == Some(ppid) && p.thread_kind().is_none())
                          .map(|(pid, _p)| *pid)
                          .collect();
    }

    /// Get thread `tid` of process `pid`, if both exist.
    pub fn get_thread_of_process(&self, pid: sysinfo::Pid, tid: sysinfo::Pid) -> Option<&sysinfo::Process> {
        self.system.processes().get(&pid).and_then(|_process|{
            self.system.processes().get(&tid)
        })
    }

    /// Get name of thread `tid` of process `pid`, or `"[dead thread]"` if it doesn't exist.
    pub fn _get_thread_name(&self, pid: sysinfo::Pid, tid: sysinfo::Pid) -> &str {
        if let Some(thread) = self.get_thread_of_process(pid, tid) {
            Self::get_process_name(thread)
        } else {
            "[dead thread]"
        }
    }

    /// Get affinities of a process and all of its threads (including the main thread).
    ///
    /// Returns process affinity, affinities of its threads indexed by TID, and string to update
    /// the status message to. Affinities are always valid, even if we can't get affinity
    /// information (which returns empty affinity CpuSets).
    pub fn get_process_and_thread_affinities(&self, pid: sysinfo::Pid)
        -> (hwloc2::CpuSet, HashMap<sysinfo::Pid, hwloc2::CpuSet>, Option<String>) 
    {
        let get_affinity = |pid, bind| {
            if let Some(affinity) = self.get_cpu_affinity(pid, bind) {
                (affinity, None)
            } else {
                (hwloc2::CpuSet::new(), Some(strings::status_failed_to_get_affinity(&pid)))
            }
        };

        // Only get an error status if this fails for process, since threads die all the time.
        let (process_affinity, status) = get_affinity(pid, hwloc2::CpuBindFlags::CPUBIND_PROCESS);
        let nonmain_threads = 
            self.get_nonmain_threads_of(pid).expect("selected process sysinfo object should exist until next process info update");
        let thread_affinities: HashMap<sysinfo::Pid, hwloc2::CpuSet> = 
            std::iter::once(pid).chain(nonmain_threads.iter().map(|(tid, _)| *tid))
            .map(|tid| (tid, get_affinity(tid, hwloc2::CpuBindFlags::CPUBIND_THREAD).0)).collect();
        (process_affinity, thread_affinities, status)
    }

    /// Get non-main threads of process `pid`, if it exists; None otherwise.
    pub fn get_nonmain_threads_of(&self, pid: sysinfo::Pid) -> Option<Vec<(sysinfo::Pid, &sysinfo::Process)>> {
        let result = self.system.processes().get(&pid).and_then(|process| {
            process.tasks().and_then(|tasks|{
                Some(tasks
                     .iter()
                     .filter(|tid| { 
                         // For some reason, a TID may be in process.tasks() but not 
                         // in self.system.processes(). Maybe a sysinfo issue?
                         // We can only work around this by ignoring such threads.
                         let thread = self.system.processes().get(tid);
                         if !thread.is_some() {
                             let proc_name = Self::get_process_name(process);
                             let mut missing_threads = self.missing_threads.borrow_mut();
                             if !missing_threads.contains(&(**tid, pid))
                             {
                                 missing_threads.insert((**tid, pid));
                                 error!("TID {}, child of PID {}/{}, could not be found around system's processes", 
                                        tid, pid, proc_name);
                             }
                         } 
                         thread.is_some()
                     })
                     .map(|tid| { 
                         (*tid, self.system.processes().get(tid).expect("threads should be found in `processes"))
                     })
                     .collect())
            })
        });
        result
    }

    /// Get information about physical CPUs on the system.
    pub fn cpu_info<'a> (&'a self) -> Vec<CPUInfo<'a>> {
        self.topo.objects_with_type(&hwloc2::ObjectType::Package).unwrap().iter().map(|package| {
            let package = self.cache.get_by_obj(*package);
            CPUInfo{
                name: package.cpu_name.as_ref().expect("only valid for Package topology objects").as_str() 
            }
        }).collect()
    }

    /// Disable logical cores in `affinity` so that only one logical core per physical core is enabled.
    pub fn disable_affinity_smt(&self, affinity: &mut hwloc2::CpuSet) {
        for cpuset in self.cache.cpusets_core() {
            let mut bits = cpuset.clone().into_iter();
            bits.next();
            for bit in  bits {
                affinity.unset(bit);
            }
        }
    }

    /// Apply thread order `order` to `selected_cores_threads` and update `selected_cores_process` accordingly.
    ///
    /// Does not actually affect the process, `pid` is used to get thread info.
    ///
    /// `ignore_main_thread` determines whether the order applies to the main thread, and `pid` is
    /// used to get access to thread IDs in the process.
    pub fn apply_thread_order(
        &self,
        order:                  ThreadOrder,
        pid:                    sysinfo::Pid,
        selected_cores_process: &mut hwloc2::CpuSet,
        selected_cores_threads: &mut HashMap<sysinfo::Pid, hwloc2::CpuSet>,
        ignore_main_thread:     bool)
    {
        let threads = self.get_nonmain_threads_of(pid)
                          .expect("if a process does not exist in sysinfo data, it shouldn't be selected anymore");

        // If main thread is not included in the order, it keeps the processes' thread selection;
        // otherwise, setting all threads' affinities will remove unused PUs from process affinity
        // when applied, so remove them from process selection.
        if !ignore_main_thread {
            selected_cores_process.clear();
        }

        let mut thread_idx = 0;
        let mut apply_thread = |tid| {
            if let Some(selected_cores) = self.apply_thread_order_single_thread(order, thread_idx, selected_cores_process, ignore_main_thread) {
                selected_cores_threads.insert(tid, selected_cores);
            }
            thread_idx += 1;
        };
        apply_thread(pid);
        for (tid, _) in threads {
            apply_thread(tid);
        }
    }

    /// Use `thread_selections` made in UI to override `apply_affinities` items (usually generated
    /// by thread order) and return updated value of `apply_affinities`.
    ///
    /// Only overrides items for threads actually running in `pid`.
    ///
    /// Used to override affinities generated by thread order by core selections made explicitly by
    /// the user.
    ///
    /// * `apply_affinities`:  Thread affinities to set indexed by thread ID.
    /// * `thread_selections`: Core selections for threads made in the UI before running the
    ///                        process, in the order they are seen in the UI. Applied to 
    ///                        `apply_affinities` in the startup order of threads.
    pub fn override_core_selections(
        &self,
        pid:                  sysinfo::Pid,
        mut apply_affinities: HashMap<sysinfo::Pid, hwloc2::CpuSet>,
        thread_selections:    &OrderedCoreSelections)
        -> HashMap<sysinfo::Pid, hwloc2::CpuSet>
    {
        // We rely on threads in this hash table having consistent iteration order when accessed
        // from different functions (e.g. apply_thread_order).
        let threads = self.get_nonmain_threads_of(pid).expect("if a process does not exist in sysinfo data, it shouldn't be selected anymore");
        let mut thread_idx = 0;
        let mut apply_thread = |tid| {
            if let Some(affinity) = thread_selections.get(thread_idx)  {
                // Thread's core selection overrides any previous selection (e.g. from thread order)>
                apply_affinities.insert(tid, affinity.clone());
            }

            thread_idx += 1;
        };
        // main thread
        apply_thread(pid);
        // non-main threads
        for (tid, _) in threads {
            apply_thread(tid);
        }
        apply_affinities
    }

    /// Apply thread order `order` to thread with index `thread_idx` (0 is the main thread) and
    /// update `selected_cores_process` accordingly if a core is added to the threads' selection.
    ///
    /// If `ignore_main_thread == true`, the main thread will use process defaults.
    ///
    /// Used by `apply_thread_order` and when threads are 'added' in `Run` mode.
    ///
    /// Returns core selection to use for this thread according to specified order, or `None` for
    /// an ignored thread where the order does not apply.
    pub fn apply_thread_order_single_thread(
        &self,
        order:                  ThreadOrder,
        thread_idx:             usize,
        selected_cores_process: &mut hwloc2::CpuSet,
        ignore_main_thread:     bool) -> Option<hwloc2::CpuSet>
    {
        let thread_idx = if ignore_main_thread {
            if thread_idx == 0 {
                return None;
            }
            thread_idx - 1
        } else {
            thread_idx
        };

        let selected_cores_thread = self.select_cores_by_thread_order( order, thread_idx );
        *selected_cores_process = selected_cores_process.union(&selected_cores_thread);
        Some(selected_cores_thread)
    }

    /// Get a CPU set of all PUs for thread `thread_index` when using thread order `order`.
    fn select_cores_by_thread_order(&self, order: ThreadOrder, thread_index: usize) -> hwloc2::CpuSet {
        let select_cores = |cpuset: &[hwloc2::CpuSet]| {
            cpuset[thread_index % cpuset.len()].clone()
        };
        match order {
            ThreadOrder::ByLCore => select_cores(self.cache.cpusets_pu()),
            ThreadOrder::ByCore  => select_cores(self.cache.cpusets_core()), 
            ThreadOrder::ByL3    => select_cores(self.cache.cpusets_l3()),  
        }
    }
}

/// Extend OrderedCoreSelections with method/s that require System.
impl OrderedCoreSelections {
    /// Apply thread `order` to per-thread core selections in `self`, and update
    /// `selected_cores_process` accordingly to contain a union of all selected cores.
    ///
    /// No cores will be removed from `selected_cores_process`, since an actual running process may 
    /// have more threads than `self` specifies, which may use any cores.
    ///
    /// Separate from `system::apply_thread_order` which works on actual currently running threads,
    /// while this works on a sequence of thread indices from 0 to n, which can later be applied to 
    /// real threads in that order.
    pub fn apply_thread_order(
        &mut self,
        order:                  ThreadOrder,
        system:                 &System,
        selected_cores_process: &mut hwloc2::CpuSet,
        ignore_main_thread:     bool)
    {
        for (thread_idx, selected_cores_thread) in self.iter_mut().enumerate() {
            // Selecting cores for thread, and update  `selected_cores_process`. 
            *selected_cores_thread =
                system.apply_thread_order_single_thread( order, thread_idx, selected_cores_process, ignore_main_thread );
        }
    }
}