emod-cli/src/debug/log.rs

use std::sync::Mutex;

use regex::{Captures, Regex};

#[cfg(windows)]
use windows_sys::Win32::System::Console::{
    CONSOLE_SCREEN_BUFFER_INFO, FOREGROUND_BLUE, FOREGROUND_GREEN, FOREGROUND_INTENSITY,
    FOREGROUND_RED, GetConsoleScreenBufferInfo, GetStdHandle, STD_OUTPUT_HANDLE,
    SetConsoleTextAttribute,
};

static CONSOLE_LOCK: Mutex<()> = Mutex::new(());

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ConsoleColor {
    Default,
    Green,
    Red,
    Yellow,
    Cyan,
    DarkGray,
}

pub struct LineBuffer {
    buffer: String,
    filter_python: bool,
}

impl LineBuffer {
    pub fn new(filter_python: bool) -> Self {
        Self {
            buffer: String::new(),
            filter_python,
        }
    }

    pub fn append<F>(&mut self, bytes: &[u8], mut process_line: F)
    where
        F: FnMut(String),
    {
        self.buffer.push_str(&String::from_utf8_lossy(bytes));
        while let Some(pos) = self.buffer.find('\n') {
            let mut line = self.buffer[..pos].to_string();
            self.buffer.drain(..=pos);
            if line.ends_with('\r') {
                line.pop();
            }
            if self.filter_python && !line.contains("[Python] ") {
                continue;
            }
            process_line(line);
        }
    }

    pub fn flush<F>(&mut self, mut process_line: F)
    where
        F: FnMut(String),
    {
        if self.buffer.is_empty() {
            return;
        }
        let mut line = std::mem::take(&mut self.buffer);
        if line.ends_with('\r') {
            line.pop();
        }
        if self.filter_python && !line.contains("[Python] ") {
            return;
        }
        process_line(line);
    }
}

pub fn process_stdout_line(line: &str) {
    if line.contains(" [INFO][Engine] ") {
        return;
    }
    let color = if line.contains("[INFO][Developer]") {
        ConsoleColor::DarkGray
    } else if contains_ignore_ascii_case(line, "SUC") {
        ConsoleColor::Green
    } else if contains_ignore_ascii_case(line, "ERROR") {
        ConsoleColor::Red
    } else if contains_ignore_ascii_case(line, "WARN") {
        ConsoleColor::Yellow
    } else if contains_ignore_ascii_case(line, "DEBUG") {
        ConsoleColor::Cyan
    } else {
        ConsoleColor::Default
    };
    print_colored(line, color);
}

pub fn process_stderr_line(line: &str) {
    let line = rewrite_python_traceback_path(line);
    print_colored(&line, ConsoleColor::Red);
}

pub fn rewrite_python_traceback_path(line: &str) -> String {
    let re = Regex::new(r#"File "([A-Za-z0-9_\.]+)", line (\d+)"#).unwrap();
    re.replace_all(line, |caps: &Captures<'_>| {
        format!(
            "File \"{}.py\", line {}",
            caps[1].replace('.', "/"),
            &caps[2]
        )
    })
    .into_owned()
}

pub fn print_colored(message: &str, color: ConsoleColor) {
    let _guard = CONSOLE_LOCK.lock().unwrap();
    #[cfg(windows)]
    unsafe {
        let handle = GetStdHandle(STD_OUTPUT_HANDLE);
        if handle.is_null() {
            println!("{}", message);
            return;
        }
        let mut info: CONSOLE_SCREEN_BUFFER_INFO = std::mem::zeroed();
        let has_info = GetConsoleScreenBufferInfo(handle, &mut info) != 0;
        if color != ConsoleColor::Default {
            SetConsoleTextAttribute(handle, color_attr(color));
        }
        println!("{}", message);
        if color != ConsoleColor::Default && has_info {
            SetConsoleTextAttribute(handle, info.wAttributes);
        }
    }
    #[cfg(not(windows))]
    {
        println!("{}", message);
    }
}

#[cfg(windows)]
fn color_attr(color: ConsoleColor) -> u16 {
    match color {
        ConsoleColor::Default => FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_BLUE,
        ConsoleColor::Green => FOREGROUND_GREEN | FOREGROUND_INTENSITY,
        ConsoleColor::Red => FOREGROUND_RED | FOREGROUND_INTENSITY,
        ConsoleColor::Yellow => FOREGROUND_RED | FOREGROUND_GREEN | FOREGROUND_INTENSITY,
        ConsoleColor::Cyan => FOREGROUND_GREEN | FOREGROUND_BLUE | FOREGROUND_INTENSITY,
        ConsoleColor::DarkGray => FOREGROUND_INTENSITY,
    }
}

fn contains_ignore_ascii_case(haystack: &str, needle: &str) -> bool {
    haystack
        .as_bytes()
        .windows(needle.len())
        .any(|window| window.eq_ignore_ascii_case(needle.as_bytes()))
}

#[cfg(test)]
mod tests {
    use super::{LineBuffer, rewrite_python_traceback_path};

    #[test]
    fn process_buffer_append_handles_crlf_half_lines_and_flush() {
        let mut buffer = LineBuffer::new(true);
        let mut lines = Vec::new();
        buffer.append(b"noise\r\n[Python] first\r\n[Python] sec", |line| {
            lines.push(line)
        });
        assert_eq!(lines, vec!["[Python] first"]);
        buffer.append(b"ond\n", |line| lines.push(line));
        buffer.flush(|line| lines.push(line));
        assert_eq!(lines, vec!["[Python] first", "[Python] second"]);
    }

    #[test]
    fn rewrite_traceback_module_path_to_file_path() {
        assert_eq!(
            rewrite_python_traceback_path(r#"Traceback File "a.b", line 3"#),
            r#"Traceback File "a/b.py", line 3"#
        );
    }
}