leenfetch_core/modules/

utils.rs

#![allow(clippy::collapsible_if, clippy::useless_vec)]

use image::{ImageFormat, ImageReader};
use std::{
    collections::HashMap,
    env, fs,
    io::{Cursor, Write},
    path::Path,
};

use super::{ascii::get_builtin_ascii_art, colors::get_builtin_distro_colors};

pub const DEFAULT_ANSI_ALL_COLORS: [&str; 16] = [
    "\x1b[1;30m", // Black
    "\x1b[1;31m", // Red
    "\x1b[1;32m", // Green
    "\x1b[1;33m", // Yellow
    "\x1b[1;34m", // Blue
    "\x1b[1;35m", // Magenta
    "\x1b[1;36m", // Cyan
    "\x1b[1;37m", // White
    "\x1b[1;90m", // Bright Black
    "\x1b[1;91m", // Bright Red
    "\x1b[1;92m", // Bright Green
    "\x1b[1;93m", // Bright Yellow
    "\x1b[1;94m", // Bright Blue
    "\x1b[1;95m", // Bright Magenta
    "\x1b[1;96m", // Bright Cyan
    "\x1b[1;97m", // Bright White
];

/// Generates a visual bar representation of a percentage using Unicode blocks.
///
/// # Arguments
///
/// * `percent` - An unsigned 8-bit integer representing the percentage (0-100) to be visualized.
///
/// # Returns
///
/// * A `String` containing a visual bar representation, with filled blocks (`█`) indicating the
///   percentage, and empty blocks (`░`) for the remainder. The total length of the bar is 14
///   characters, enclosed in square brackets.
pub fn get_bar(percent: u8) -> String {
    let total_blocks = 14;
    let filled_blocks = (percent as usize * total_blocks) / 100;
    let empty_blocks = total_blocks - filled_blocks;

    // Pre-computed block strings for common values (avoid allocations)
    const BLOCKS: &[&str] = &[
        "",
        "█",
        "██",
        "███",
        "████",
        "█████",
        "██████",
        "███████",
        "████████",
        "█████████",
        "██████████",
        "███████████",
        "████████████",
        "█████████████",
        "██████████████",
    ];
    const EMPTY_BLOCKS: &[&str] = &[
        "",
        "░",
        "░░",
        "░░░",
        "░░░░",
        "░░░░░",
        "░░░░░░",
        "░░░░░░░",
        "░░░░░░░░",
        "░░░░░░░░░",
        "░░░░░░░░░░",
        "░░░░░░░░░░░",
        "░░░░░░░░░░░░",
        "░░░░░░░░░░░░░",
        "░░░░░░░░░░░░░░",
    ];

    let filled = BLOCKS[filled_blocks.min(14)];
    let empty = EMPTY_BLOCKS[empty_blocks.min(14)];

    format!("[{}{}]", filled, empty)
}

/// Generates a vector of 2 strings, each containing a row of 8 blocks
/// colored with different ANSI foreground colors. The first string has
/// normal colors, the second has bold colors.
///
/// The input string `color_blocks` should contain 8 identical block characters
/// (e.g. █, ░, ▓, ▒, etc.). The output strings will have these blocks
/// colored with different ANSI colors.
pub fn get_terminal_color(color_blocks: &str) -> String {
    let color_codes: [u8; 8] = [30, 31, 32, 33, 34, 35, 36, 37]; // ANSI foreground colors

    let mut normal = Vec::with_capacity(8);
    // let mut bold = Vec::with_capacity(8);

    for &code in &color_codes {
        normal.push(format!("\x1b[{}m{}\x1b[0m", code, color_blocks)); // normal
        // bold.push(format!("\x1b[1;{}m{}\x1b[0m", code, color_blocks)); // bold
    }

    // vec![normal.join(""), bold.join("")]
    normal.join("")
}

// ---------------------------------
//        ASCII ART Functions
// ---------------------------------

/// Reads a file at the given custom path and returns its content as a string.
/// If there is an error while reading the file, an empty string is returned.
///
/// # Arguments
///
/// * `custom_path`: The path to the custom ASCII art file.
///
/// # Returns
///
/// A string containing the content of the file, or an empty string if there was an error.
pub fn get_custom_ascii(custom_path: &str) -> String {
    if let Ok(content) = fs::read_to_string(Path::new(custom_path)) {
        return content;
    }

    "".to_string()
}

pub fn is_image_path(path: &str) -> bool {
    Path::new(path)
        .extension()
        .and_then(|ext| ext.to_str())
        .map(|ext| matches!(ext.to_ascii_lowercase().as_str(), "png" | "jpg" | "jpeg"))
        .unwrap_or(false)
}

pub fn terminal_supports_inline_images() -> bool {
    matches!(env::var("TERM"), Ok(term) if term == "xterm-kitty")
        || env::var_os("KITTY_WINDOW_ID").is_some()
}

pub fn render_inline_image(path: &str, columns: usize) -> Result<(), String> {
    let reader = ImageReader::open(path)
        .map_err(|err| format!("Failed to open image {path}: {err}"))?
        .with_guessed_format()
        .map_err(|err| format!("Failed to detect image format for {path}: {err}"))?;

    let image = reader
        .decode()
        .map_err(|err| format!("Failed to decode image {path}: {err}"))?;

    let mut png_bytes = Cursor::new(Vec::new());
    image
        .write_to(&mut png_bytes, ImageFormat::Png)
        .map_err(|err| format!("Failed to encode image {path} as PNG: {err}"))?;

    let encoded = base64_encode(png_bytes.get_ref());
    let mut stdout = std::io::stdout();
    if columns > 0 {
        write!(
            &mut stdout,
            "\x1b_Ga=T,C=1,f=100,c={columns};{encoded}\x1b\\"
        )
    } else {
        write!(&mut stdout, "\x1b_Ga=T,C=1,f=100;{encoded}\x1b\\")
    }
    .map_err(|err| format!("Failed to write inline image escape sequence: {err}"))?;
    stdout
        .flush()
        .map_err(|err| format!("Failed to flush inline image: {err}"))?;

    Ok(())
}

fn base64_encode(data: &[u8]) -> String {
    const TABLE: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

    let mut out = String::with_capacity(data.len().div_ceil(3) * 4);
    let mut chunks = data.chunks_exact(3);

    for chunk in &mut chunks {
        let n = ((chunk[0] as u32) << 16) | ((chunk[1] as u32) << 8) | chunk[2] as u32;
        out.push(TABLE[((n >> 18) & 0x3f) as usize] as char);
        out.push(TABLE[((n >> 12) & 0x3f) as usize] as char);
        out.push(TABLE[((n >> 6) & 0x3f) as usize] as char);
        out.push(TABLE[(n & 0x3f) as usize] as char);
    }

    match chunks.remainder() {
        [b0] => {
            let n = (*b0 as u32) << 16;
            out.push(TABLE[((n >> 18) & 0x3f) as usize] as char);
            out.push(TABLE[((n >> 12) & 0x3f) as usize] as char);
            out.push('=');
            out.push('=');
        }
        [b0, b1] => {
            let n = ((*b0 as u32) << 16) | ((*b1 as u32) << 8);
            out.push(TABLE[((n >> 18) & 0x3f) as usize] as char);
            out.push(TABLE[((n >> 12) & 0x3f) as usize] as char);
            out.push(TABLE[((n >> 6) & 0x3f) as usize] as char);
            out.push('=');
        }
        _ => {}
    }

    out
}

/// Given a distro name, returns a string of its corresponding ASCII art.
/// If the distro isn't found, an empty string is returned.
/// If the distro is "off", an empty string is returned.
pub fn get_ascii_and_colors(ascii_distro: &str) -> String {
    if ascii_distro == "off" {
        return "".to_string();
    }

    let ascii_art = resolve_ascii_art(ascii_distro);

    ascii_art.to_string()
}

/// Resolves ASCII art for a given distro name.
/// Tries direct match first, then falls back to ID_LIKE if available.
fn resolve_ascii_art(distro: &str) -> &'static str {
    // 1. Try direct distro name match
    if let Some(art) = get_builtin_ascii_art(distro) {
        return art;
    }

    #[cfg(target_os = "linux")]
    // 2. No match — try ID_LIKE parent distro
    if let Some(parent) = crate::modules::linux::system::distro::get_id_like() {
        if let Some(art) = get_builtin_ascii_art(&parent) {
            return art;
        }
    }

    // 3. Nothing found — return fallback DEFAULT
    DEFAULT_ASCII
}

const DEFAULT_ASCII: &str = r#"${c2}        #####
${c2}       #######
${c2}       ##${c1}O${c2}#${c1}O${c2}##
${c2}       #${c3}#####${c2}#
${c2}     ##${c1}##${c3}###${c1}##${c2}##
${c2}    #${c1}##########${c2}##
${c2}   #${c1}############${c2}##
${c2}   #${c1}############${c2}###
${c3}  ##${c2}#${c1}###########${c2}##${c3}#
${c3}######${c2}#${c1}#######${c2}#${c3}######
${c3}#######${c2}#${c1}#####${c2}#${c3}#######
${c3}  #####${c2}#######${c3}#####"#;

// ---------------------------------
//        Color Functions
// ---------------------------------

/// Replaces placeholders in a string with ANSI escape codes to colorize
/// the output.
///
/// Placeholders are in the form of `${{key}}`, where `key` is the key in
/// the provided `colors` HashMap. The value associated with the `key` is
/// the ANSI escape code for the color.
pub fn colorize_text(input: String, colors: &HashMap<&str, &str>) -> String {
    let mut result = String::new();

    for line in input.lines() {
        let mut colored = line.to_owned();
        for (key, code) in colors {
            let placeholder = format!("${{{}}}", key);
            colored = colored.replace(&placeholder, code);
        }
        result.push_str(&colored);
        result.push('\n');
    }

    result
}

/// Creates a `HashMap` of ANSI color codes from the given entries.
///
/// Each entry is a tuple containing a key and a corresponding ANSI
/// color code. The function populates the `HashMap` with these entries
/// and also adds bold variants for each color code that starts with
/// `\x1b[0;`. The bold variant key is prefixed with "bold." (e.g.,
/// `bold.c1` for `c1`).
///
/// Additionally, a "reset" key is included in the map, which maps to
/// the ANSI reset code `\x1b[0m`.
///
/// # Arguments
///
/// * `entries` - A slice of tuples where each tuple contains a string
///   key and an ANSI color code.
///
/// # Returns
///
/// * A `HashMap` with the original entries, their bold variants, and a
///   reset entry.
pub fn color_palette(
    entries: &[(&'static str, &'static str)],
) -> HashMap<&'static str, &'static str> {
    let mut map = HashMap::new();
    for (k, v) in entries {
        map.insert(*k, *v);

        // Add bold variant: bold.c1 → \x1b[1;31m if c1 is \x1b[0;31m
        if let Some(code) = v.strip_prefix("\x1b[0;") {
            let bold_code = format!("\x1b[1;{}", code);
            let bold_key = format!("bold.{}", k);
            map.insert(
                Box::leak(bold_key.into_boxed_str()),
                Box::leak(bold_code.into_boxed_str()),
            );
        }
    }

    map.insert("reset", "\x1b[0m");
    map
}

/// Given a slice of color indices or a distro name, generates a HashMap
/// of `cX` keys (where `X` is the index, starting from 1) mapped to the
/// corresponding ANSI foreground color codes.
///
/// The color order is determined by the input slice. The first color is
/// assigned to `c1`, the second to `c2`, and so on. If the input slice is
/// shorter than 16 elements, the remaining colors are filled from the
/// default color palette in the order they appear.
///
/// If the input slice is empty, the function returns an empty HashMap.
///
/// The HashMap also includes bold variants for each color, which can be
/// accessed using the `bold.*` keys. For example, `bold.c1` would be the
/// bold variant of `c1`.
///
/// The `reset` key is also included, which resets the text to the default
/// color.
pub fn get_colors_in_order(color_order: &[u8]) -> HashMap<&'static str, &'static str> {
    // Start with c0 = bold black
    let mut entries: Vec<(&'static str, &'static str)> = vec![("c0", "\x1b[1;30m")];

    let mut used = vec![false; 16]; // support 0–15

    // Fill c1 to cX using given color_order
    for (i, &idx) in color_order.iter().enumerate() {
        if idx < 16 {
            let key: &'static str = Box::leak(format!("c{}", i + 1).into_boxed_str());
            entries.push((key, DEFAULT_ANSI_ALL_COLORS[idx as usize]));
            used[idx as usize] = true;
        }
    }

    // Fill remaining cX from unused colors
    let mut next_index = color_order.len() + 1;
    for (i, &color) in DEFAULT_ANSI_ALL_COLORS.iter().enumerate() {
        if !used[i] {
            let key: &'static str = Box::leak(format!("c{}", next_index).into_boxed_str());
            entries.push((key, color));
            next_index += 1;
        }
    }

    // Generate HashMap with bold.* variants and reset
    let mut map = color_palette(&entries);
    map.insert("reset", "\x1b[0m");
    map
}

/// Given a string of comma-separated color indices or a distro name, returns a
/// HashMap of color codes c0 to cX as found in the given distro's color
/// definition. The color codes are from the ANSI color palette.
pub fn get_custom_colors_order(colors_str_order: &str) -> HashMap<&'static str, &'static str> {
    let custom_color_str_list: Vec<&str> = colors_str_order.split(',').map(str::trim).collect();

    // Try to parse all color indices
    let all_parsed: Option<Vec<u8>> = custom_color_str_list
        .iter()
        .map(|s| s.parse::<u8>().ok())
        .collect();

    let color_list: Vec<u8> = if let Some(list) = all_parsed {
        list
    } else {
        // Fallback: interpret the string as a distro name
        get_builtin_distro_colors(colors_str_order).to_vec()
    };

    get_colors_in_order(&color_list)
}

/// Given a distro name, returns a HashMap of color codes c0 to cX as found
/// in the given distro's color definition. The color codes are from the
/// ANSI color palette. If the distro isn't found, an empty HashMap is returned.
pub fn get_distro_colors(distro: &str) -> HashMap<&'static str, &'static str> {
    let dist_color = get_builtin_distro_colors(distro);

    get_colors_in_order(dist_color)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::EnvLock;

    #[test]
    fn bar_respects_percentage_bounds() {
        assert_eq!(get_bar(0), "[░░░░░░░░░░░░░░]");
        assert_eq!(get_bar(100), "[██████████████]");
        assert!(get_bar(50).contains('█'));
    }

    #[test]
    fn terminal_color_emits_expected_blocks() {
        let visual = get_terminal_color("■");
        assert_eq!(visual.matches('■').count(), 8);
        assert!(visual.contains("\x1b[31m"), "missing ANSI escape: {visual}");
    }

    #[test]
    fn color_palette_includes_bold_and_reset() {
        let map = color_palette(&[("c1", "\x1b[0;31m")]);
        assert_eq!(map.get("c1"), Some(&"\x1b[0;31m"));
        let bold_key = map
            .keys()
            .find(|k| k.starts_with("bold.c1"))
            .expect("bold variant missing");
        assert!(map.get(bold_key).unwrap().starts_with("\x1b[1;"));
        assert_eq!(map.get("reset"), Some(&"\x1b[0m"));
    }

    #[test]
    fn colors_in_order_fills_defaults() {
        let map = get_colors_in_order(&[1, 2, 3]);
        assert_eq!(map.get("c1"), Some(&DEFAULT_ANSI_ALL_COLORS[1]));
        assert_eq!(map.get("c2"), Some(&DEFAULT_ANSI_ALL_COLORS[2]));
        assert_eq!(map.get("c3"), Some(&DEFAULT_ANSI_ALL_COLORS[3]));
        assert!(map.contains_key("c4"));
        assert_eq!(map.get("reset"), Some(&"\x1b[0m"));
    }

    #[test]
    fn image_path_detection_is_extension_based() {
        assert!(is_image_path("/tmp/logo.png"));
        assert!(is_image_path("/tmp/logo.JPG"));
        assert!(is_image_path("/tmp/logo.jpeg"));
        assert!(!is_image_path("/tmp/logo.txt"));
    }

    #[test]
    fn kitty_detection_uses_terminal_hints() {
        let env = EnvLock::acquire(&["TERM", "KITTY_WINDOW_ID"]);

        env.set_var("TERM", "xterm-kitty");
        env.remove_var("KITTY_WINDOW_ID");
        assert!(terminal_supports_inline_images());

        env.set_var("TERM", "xterm-256color");
        env.set_var("KITTY_WINDOW_ID", "1");
        assert!(terminal_supports_inline_images());

        env.set_var("TERM", "xterm-256color");
        env.remove_var("KITTY_WINDOW_ID");
        assert!(!terminal_supports_inline_images());
    }
}