cnxt 0.1.6

//! A couple of functions to set whether to colorize.
//!
//! You can set the color level to use for the terminal with [`set_should_colorize()`].
//!
//! And get the current color level with [`get_should_colorize()`].
//!
//! # Example
//! ```rust
//! use cnxt::control::{set_should_colorize, ShouldColorize};
//!
//! // Enable colorization with the detected color level
//! set_should_colorize(ShouldColorize::Yes);
//!
//! // Disable colorization
//! set_should_colorize(ShouldColorize::No);
//!
//! // Enable colorization but force 16 ANSI colors
//! set_should_colorize(ShouldColorize::YesWithAnsi16);
//!
//! // Automatically enable or disable colorization based on environment variables and terminal status
//! set_should_colorize(ShouldColorize::from_env());
//! ```

use std::{
    env, io,
    io::IsTerminal,
    sync::{
        LazyLock,
        atomic::{AtomicU8, Ordering},
    },
};

/// The detected color level for the current terminal.
///
/// This is lazily initialized the first time it's accessed and detects
/// the available color level based on the environment and terminal capabilities.
///
/// # Examples
///
/// ```rust
/// use cnxt::control::COLOR_LEVEL_DETECTED;
///
/// // Get the detected color level
/// let detected_level = *COLOR_LEVEL_DETECTED;
/// ```
pub static COLOR_LEVEL_DETECTED: LazyLock<ColorLevel> =
    LazyLock::new(ColorLevel::detect);

/// The global setting for whether and how to colorize output.
///
/// When the default feature `terminal-detection` is disabled, this defaults to `YesWithTrueColor`.
///
/// This atomic value stores the current colorization setting as specified by
/// [`ShouldColorize`]. It's initialized from the environment using
/// [`ShouldColorize::from_env()`] but can be changed at runtime with
/// [`set_should_colorize()`].
pub static SHOULD_COLORIZE: LazyLock<AtomicU8> = LazyLock::new(|| {
    AtomicU8::new(
        #[cfg(feature = "terminal-detection")]
        {
            ShouldColorize::from_env() as u8
        },
        #[cfg(not(feature = "terminal-detection"))]
        {
            ShouldColorize::YesWithTrueColor as u8
        },
    )
});

/// Sets a flag to the console to use a virtual terminal environment.
///
/// This is primarily used for Windows 10 environments which will not correctly colorize
/// the outputs based on ANSI escape codes.
#[allow(clippy::result_unit_err)]
#[cfg(windows)]
pub fn set_virtual_terminal(use_virtual: bool) {
    use windows_sys::Win32::System::Console::{
        ENABLE_VIRTUAL_TERMINAL_PROCESSING, GetConsoleMode, GetStdHandle,
        STD_OUTPUT_HANDLE, SetConsoleMode,
    };

    unsafe {
        let handle = GetStdHandle(STD_OUTPUT_HANDLE);
        if handle.is_null() {
            return;
        }

        let mut mode = 0;
        if GetConsoleMode(handle, &mut mode) == 0 {
            return;
        }

        let new_mode = if use_virtual {
            mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING
        } else {
            mode & !ENABLE_VIRTUAL_TERMINAL_PROCESSING
        };

        if new_mode != mode {
            SetConsoleMode(handle, new_mode);
        }
    }
}

/// Sets the color level to use for the terminal.
///
/// Default value is generated by [`ShouldColorize::from_env()`].
pub fn set_should_colorize(should_colorize: ShouldColorize) {
    SHOULD_COLORIZE.store(should_colorize as u8, Ordering::Relaxed);
}

/// Gets the current color level to use for the terminal.
pub fn get_should_colorize() -> ShouldColorize {
    SHOULD_COLORIZE.load(Ordering::Relaxed).into()
}

pub fn get_current_color_level() -> ColorLevel {
    match get_should_colorize() {
        ShouldColorize::No => ColorLevel::None,
        ShouldColorize::Yes => *COLOR_LEVEL_DETECTED,
        level => level.into(),
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
/// The color level to use for the terminal.
/// Determines the color depth to use for the terminal.
pub enum ColorLevel {
    /// Use no colors.
    None,
    /// Use 16 ANSI colors.
    Ansi16,
    /// Use 256 ANSI colors.
    Ansi256,
    /// Use True Color.
    TrueColor,
}

impl ColorLevel {
    fn detect() -> Self {
        // Check for 24-bit color support via COLORTERM
        if env::var("COLORTERM")
            .is_ok_and(|v| matches!(v.as_str(), "truecolor" | "24bit"))
        {
            return Self::TrueColor;
        }

        // Detect Windows Terminal in Windows and WSL
        if env::var_os("WT_SESSION").is_some() {
            return Self::TrueColor;
        }

        // Detect CI environments, which typically have limited color support
        if env::var_os("CI").is_some() {
            return Self::Ansi256;
        }

        // Windows version-specific checks
        #[cfg(target_os = "windows")]
        {
            use windows_version::OsVersion;
            let version = OsVersion::current();

            // Windows 10 build 14931+ supports TrueColor
            if version >= OsVersion::new(10, 0, 0, 14931) {
                return Self::TrueColor;
            }

            // Windows 10 build 10586+ supports 256 colors
            if version >= OsVersion::new(10, 0, 0, 10586) {
                return Self::Ansi256;
            }
        }

        // Check TERM for 256-color indication
        if let Some(term) =
            env::var_os("TERM").and_then(|term| term.into_string().ok())
        {
            if term.ends_with("-256color") || term.ends_with("256") {
                return Self::Ansi256;
            }
        }

        // Fallback to basic ANSI colors
        Self::Ansi16
    }
}

/// Whether and how to colorize the output.
#[repr(u8)]
#[derive(Clone, Copy, Debug)]
pub enum ShouldColorize {
    /// Do not colorize the output.
    No,
    /// Colorize the output with the automaticly detected color level.
    Yes,
    /// Force colorization with 16 ANSI colors.
    YesWithAnsi16,
    /// Force colorization with 256 ANSI colors.
    YesWithAnsi256,
    /// Force colorization with True Color.
    YesWithTrueColor,
}

impl ShouldColorize {
    /// Determines if colorization should be applied based on environment variables and terminal status.
    ///
    /// Priority order for environment variables:
    /// 1. `CLICOLOR_FORCE` (force enable colorization)
    /// 2. `NO_COLOR` (force disable colorization)
    /// 3. `CLICOLOR` (enable colorization if set, depending on tty)
    /// 4. If none of the above, use the terminal status (enabled if stdout is a tty)
    #[must_use]
    pub fn from_env() -> Self {
        if env::var("CLICOLOR_FORCE").is_ok_and(|v| v != "0") {
            return Self::Yes;
        }

        if env::var("NO_COLOR").is_ok() {
            return Self::No;
        }

        if env::var("CLICOLOR").is_ok_and(|v| v != "0") {
            return Self::Yes;
        }

        if io::stdout().is_terminal() {
            Self::Yes
        } else {
            Self::No
        }
    }
}

impl From<u8> for ShouldColorize {
    fn from(value: u8) -> Self {
        match value {
            0 => Self::No,
            2 => Self::YesWithAnsi16,
            3 => Self::YesWithAnsi256,
            4 => Self::YesWithTrueColor,
            _ => Self::Yes, // unreachable, but default to Yes to avoid panics
        }
    }
}

impl From<ShouldColorize> for ColorLevel {
    fn from(value: ShouldColorize) -> Self {
        match value {
            ShouldColorize::YesWithAnsi16 => Self::Ansi16,
            ShouldColorize::YesWithAnsi256 => Self::Ansi256,
            ShouldColorize::YesWithTrueColor => Self::TrueColor,
            _ => Self::None,
        }
    }
}