use crate::utils;
use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(windows)]
use windows_sys::Win32::System::Console::{
GetConsoleMode, GetStdHandle, SetConsoleMode, ENABLE_VIRTUAL_TERMINAL_PROCESSING,
STD_OUTPUT_HANDLE,
};
const COLOURS: [&str; 8] = [
"black", "red", "green", "yellow", "blue", "magenta", "cyan", "white",
];
const COLOUR_ATTRIBUTES: [&str; 8] = [
"bold",
"dim",
"italic",
"underline",
"blink",
"reversed",
"hidden",
"strikethrough",
];
const COLOUR_RESET: &str = "\x1b[0m";
static SHOULD_COLORIZE: AtomicBool = AtomicBool::new(false);
pub fn init(always: bool) {
#[cfg(windows)]
unsafe {
let handle = GetStdHandle(STD_OUTPUT_HANDLE);
let mut original_mode = 0;
GetConsoleMode(handle, &mut original_mode);
let enabled = original_mode & ENABLE_VIRTUAL_TERMINAL_PROCESSING
== ENABLE_VIRTUAL_TERMINAL_PROCESSING;
if !enabled {
SetConsoleMode(handle, ENABLE_VIRTUAL_TERMINAL_PROCESSING | original_mode);
}
}
SHOULD_COLORIZE.store(always, Ordering::SeqCst);
}
fn hex_to_rgb(code: &str) -> Option<(u8, u8, u8)> {
if code.len() == 7 {
Some((
u8::from_str_radix(&code[1..3], 16).ok()?,
u8::from_str_radix(&code[3..5], 16).ok()?,
u8::from_str_radix(&code[5..7], 16).ok()?,
))
} else {
None
}
}
fn parse_ansi(code: &str) -> Option<u8> {
code.get(5..(code.len() - 1))?.parse::<u8>().ok()
}
fn parse_rgb(code: &str) -> Option<(u8, u8, u8)> {
let mut values = code
.get(4..(code.len() - 1))?
.split(',')
.filter_map(|x| x.trim().parse::<u8>().ok());
Some((values.next()?, values.next()?, values.next()?))
}
pub fn colour(code: &str) -> Option<String> {
let mut code = code.to_lowercase();
let mut bg = None;
if let Some(index) = code.find("on #") {
let end = 3 + 7;
let (r, g, b) = hex_to_rgb(code.get((index + 3)..(index + end))?)?;
code.replace_range(index..(index + end), "");
bg = Some(format!("48;2;{};{};{}", r, g, b));
} else if let Some(index) = code.find("on rgb(") {
let end = 3 + code.get((index + 3)..)?.find(')')? + 1;
let (r, g, b) = parse_rgb(code.get((index + 3)..(index + end))?)?;
code.replace_range(index..(index + end), "");
bg = Some(format!("48;2;{};{};{}", r, g, b));
} else if let Some(index) = code.find("on ansi(") {
let end = 4 + code.get((index + 4)..)?.find(')')? + 1;
let number = parse_ansi(code.get((index + 3)..(index + end))?)?;
code.replace_range(index..(index + end), "");
bg = Some(format!("48;5;{}", number));
} else {
let mut number = 100_u8;
for colour in COLOURS {
let bright_bg_colour = "on bright ".to_owned() + colour;
if let Some(index) = code.find(&bright_bg_colour) {
code.replace_range(index..(index + bright_bg_colour.len()), "");
bg = Some(number.to_string());
break;
}
number += 1;
}
if bg.is_none() {
number = 40;
for colour in COLOURS {
let bg_colour = "on ".to_owned() + colour;
if let Some(index) = code.find(&bg_colour) {
code.replace_range(index..(index + bg_colour.len()), "");
bg = Some(number.to_string());
break;
}
number += 1;
}
}
}
let mut fg = None;
if let Some(index) = code.find('#') {
let end = 7;
let (r, g, b) = hex_to_rgb(code.get(index..(index + end))?)?;
code.replace_range(index..(index + end), "");
fg = Some(format!("38;2;{};{};{}", r, g, b));
} else if let Some(index) = code.find("rgb(") {
let end = code.get(index..)?.find(')')? + 1;
let (r, g, b) = parse_rgb(code.get(index..(index + end))?)?;
code.replace_range(index..(index + end), "");
fg = Some(format!("38;2;{};{};{}", r, g, b));
} else if let Some(index) = code.find("ansi(") {
let end = code.get(index..)?.find(')')? + 1;
let number = parse_ansi(code.get(index..(index + end))?)?;
code.replace_range(index..(index + end), "");
fg = Some(format!("38;5;{}", number));
} else {
let mut number = 90_u8;
for colour in COLOURS {
let bright_fg_colour = "bright ".to_owned() + colour;
if let Some(index) = code.find(&bright_fg_colour) {
code.replace_range(index..(index + bright_fg_colour.len()), "");
fg = Some(number.to_string());
break;
}
number += 1;
}
if fg.is_none() {
number = 30;
for fg_colour in COLOURS {
if let Some(index) = code.find(fg_colour) {
code.replace_range(index..(index + fg_colour.len()), "");
fg = Some(number.to_string());
break;
}
number += 1;
}
}
}
let mut attributes = String::new();
let mut number = 1_u8;
for attribute in COLOUR_ATTRIBUTES {
if let Some(index) = code.find(attribute) {
code.replace_range(index..(index + attribute.len()), "");
if !attributes.is_empty() {
attributes.push(';')
}
attributes += &number.to_string();
}
number += 1;
}
let attributes = if attributes.is_empty() {
None
} else {
Some(attributes)
};
let escape_code = "\x1b[".to_owned()
+ &[fg, bg, attributes]
.into_iter()
.flatten()
.collect::<Vec<String>>()
.join(";")
+ "m";
if escape_code == "\x1b[m" {
None
} else {
Some(escape_code)
}
}
pub trait Colorizer {
fn colorize(&self, code: &str) -> String;
fn trim_ansi(&self) -> String;
fn len_ansi(&self) -> usize;
}
impl Colorizer for str {
fn colorize(&self, code: &str) -> String {
if !SHOULD_COLORIZE.load(Ordering::Acquire) {
return self.to_owned();
}
let escape_code = colour(code);
if let Some(escape_code) = escape_code {
escape_code + self + COLOUR_RESET
} else {
self.to_owned()
}
}
fn trim_ansi(&self) -> String {
let mut text = self.to_owned();
while let Some(start) = text.find("\x1b[") {
if let Some(end) = text[start..].find('m') {
text.replace_range(start..(start + end + 1), "");
}
}
text
}
fn len_ansi(&self) -> usize {
utils::len(&self.trim_ansi())
}
}