cpumap 0.2.1

use hwloc2;
use regex;
#[cfg(feature = "gui")]
use eframe::egui;

/// Convert a hwloc CPU 'brand' string into something short and human-readable.
pub fn cpu_brand_transform(brand: &str) -> String {
    // Most AMD CPUs around 2023 match this.
    let re_amd = regex::Regex::new(r"^(AMD.*?)\s(\d+)-Core Processor$").unwrap();
    if let Some(captures) = re_amd.captures(brand) {
        return String::from(captures.get(1).unwrap().as_str());
    }
    // Most Intel CPUs around 2023 match this.
    let re_intel_core = regex::Regex::new(r"^Intel\(R\)\sCore\(TM\)\s(.*?)\sCPU @ (\d+)\.(\d+)GHz$").unwrap();
    if let Some(captures) = re_intel_core.captures(brand) {
        return format!("Intel Core {}", captures.get(1).unwrap().as_str());
    }
    // Alder Lake and other heterogeneous Intel CPUs.
    let re_intel_hetero_core = regex::Regex::new(r"^\d+th Gen Intel\(R\)\sCore\(TM\)\s(i\d-\d+.?)$").unwrap();
    if let Some(captures) = re_intel_hetero_core.captures(brand) {
        return format!("Intel Core {}", captures.get(1).unwrap().as_str());
    }
    // TODO more regexes for more CPU brands/types.
    String::from(brand)
}

/// Generate a human-readable string with size of a cache.
pub fn format_cache_size(bytes: u64, precision: usize) -> String {
    let units = ["B", "k", "M", "G"];
    let pow = (bytes as f64).log(1024.0).floor() as usize;
    // Make sure we stay in range.
    let pow = usize::min(pow, units.len() - 1);
    // Number of kilo/mega/etc bytes.
    let number = bytes as f64 / usize::pow(1024, pow as u32) as f64;
    let number_str = format!("{:.prec$}", number, prec = precision);
    format!("{}{}", number_str.trim_end_matches('0').trim_end_matches('.'), units[pow])
}

/// Get the closest pair of factors of `count` where `a*b=count` and `a<=b`.
pub fn find_squarest_factors(count: usize) -> (usize, usize) {
    if count == 0 {
        return (0,0)
    }
    let sqrt_n = (count as f64).sqrt() as usize;
    for short_side in (1..=sqrt_n).rev() {
        if count % short_side == 0 {
            return (short_side, count / short_side);
        }
    }
    unreachable!("Any `count` is divisible by 1, so this will not be reached");
}

/// For a given core count, get `(rows,cols)` of the squarest grid where there is exactly 1 cell per core.
pub fn squarish_grid_dimensions(count: usize, horizontal: bool) -> (usize, usize) {
    let (shorter, longer) = find_squarest_factors(count);
    if horizontal {(shorter, longer)} else {(longer, shorter)}
}

/// Return `s` truncated to at most `max_width` characters, ended with '...' if truncated.
#[cfg(feature = "gui")] // not used in TUI build ATM
pub fn truncate_and_add_ellipsis(s: &str, max_width: usize) -> String {
    if s.chars().count() > max_width {
        let truncated = s.chars().take(max_width - 3).collect::<String>();
        format!("{}...", truncated)
    } else {
        s.to_string()
    }
}

/// Used to extend egui::Response to assign shortcuts to actuons on UI elements.
#[cfg(feature = "gui")]
pub trait ShorcutControllable {
    /// This element will be focused on given shortcut.
    fn focus_on_shortcut(self, ui: &egui::Ui, shortcut: egui::KeyboardShortcut) -> Self;
    /// Returns true if the element is clicked, or the given shortcut is pressed.
    fn clicked_or_shortcut(&self, ui: &egui::Ui, shortcut: egui::KeyboardShortcut) -> bool;
    /// Returns true if given shortcut is pressed.
    fn shortcut(&self, ui: &egui::Ui, shortcut: egui::KeyboardShortcut) -> bool;
}

#[cfg(feature = "gui")]
impl ShorcutControllable for egui::Response {
    fn focus_on_shortcut(self, ui: &egui::Ui, shortcut: egui::KeyboardShortcut) -> Self {
        if ui.input_mut(|i| i.consume_shortcut(&shortcut)) {
            self.request_focus();
        }
        self
    }

    fn clicked_or_shortcut(&self, ui: &egui::Ui, shortcut: egui::KeyboardShortcut) -> bool {
        self.clicked() || ui.input_mut(|i| i.consume_shortcut(&shortcut))
    }

    fn shortcut(&self, ui: &egui::Ui, shortcut: egui::KeyboardShortcut) -> bool {
        ui.input_mut(|i| i.consume_shortcut(&shortcut))
    }
}

/// Trait used to extend CpuSet with more operations (especially boolean set operations).
pub trait CpuSetExt {
    /// Set union - return a set with all items from `self` and `other`.
    fn union(&self, other: &Self) -> Self;

    /// Set subtraction - return a set with all items from `self` not found in `other`.
    fn subtract(&self, other: &Self) -> Self;

    /// Set intersection - return a set with all items found in both `self` and `other`.
    fn intersection(&self, other: &Self) -> Self;

    /// Toggle a bit (core) at specified index.
    fn toggle(&mut self, bit: u32);

    /// After removing cores from a process CpuSet, call this to synchronize that change with thread CpuSets.
    fn synchronize_removed_process_cores_into<'a, I>(&self, thread_cpusets: I) 
    where 
        I: Iterator<Item = &'a mut hwloc2::CpuSet>;

    /// Same as synchronize_remove_process_cores_into, but takes iterator of Option<CpuSet> and only
    /// synchronizes into CpuSets that are Some.
    fn synchronize_removed_process_cores_into_opt<'a, I>(&self, thread_cpusets: I) 
    where 
        I: Iterator<Item = &'a mut Option<hwloc2::CpuSet>>;

    /// After adding cores to a thread CpuSets, call this to synchronize that change with this process CpuSet.
    fn synchronize_added_thread_cores_from<'a, I>(&mut self, thread_cpusets: I) 
    where 
        I: Iterator<Item = &'a hwloc2::CpuSet>;

    /// Same as synchronize_added_cores_from, but takes iterator of Option<CpuSet> and only
    /// synchronizes from CpuSets that are Some.
    fn synchronize_added_thread_cores_from_opt<'a, I>(&mut self, thread_cpusets: I) 
    where
        I: Iterator<Item = &'a Option<hwloc2::CpuSet>>;
}

impl CpuSetExt for hwloc2::CpuSet {
    fn union(&self, other: &Self) -> Self {
        let mut result = self.clone();
        for bit in other.clone() {
            result.set(bit);
        }
        result
    }

    fn subtract(&self, other: &Self) -> Self {
        let mut result = self.clone();
        for bit in other.clone() {
            result.unset(bit);
        }
        result
    }

    fn intersection(&self, other: &Self) -> Self {
        let mut result = Self::new();
        for bit in self.clone() {
            if other.is_set(bit) {
                result.set(bit);
            }
        }
        result
    }

    fn toggle(&mut self, bit: u32) {
        if self.is_set(bit) {
            self.unset(bit);
        } else {
            self.set(bit);
        }
    }

    fn synchronize_removed_process_cores_into<'a, I>(&self, thread_cpusets: I) 
    where I: Iterator<Item = &'a mut hwloc2::CpuSet> 
    {
        for thread in thread_cpusets {
            *thread = self.intersection(thread);
        }
    }

    fn synchronize_removed_process_cores_into_opt<'a, I>(&self, thread_cpusets: I) 
    where I: Iterator<Item = &'a mut Option<hwloc2::CpuSet>> 
    {
        self.synchronize_removed_process_cores_into(thread_cpusets.filter_map(|t| t.as_mut()));
    }

    fn synchronize_added_thread_cores_from<'a, I>(&mut self, thread_cpusets: I) 
    where I: Iterator<Item = &'a hwloc2::CpuSet> 
    {
        for thread in thread_cpusets {
            *self = self.union(thread);
        }
    }

    fn synchronize_added_thread_cores_from_opt<'a, I>(&mut self, thread_cpusets: I) 
    where I: Iterator<Item = &'a Option<hwloc2::CpuSet>> 
    {
        self.synchronize_added_thread_cores_from(thread_cpusets.filter_map(|t| t.as_ref()));
    }
}