feat: autodetect light theme on unix-like os (#168)

4 files changed, 280 insertions, 7 deletions

diff --git a/src/client/mod.rs b/src/client/mod.rs
index 37440f1..b2ceb84 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -11,7 +11,6 @@ use async_trait::async_trait;
 use reqwest::{ClientBuilder, Proxy};
 use serde::Deserialize;
 use std::{env, time::Duration};
-use tokio::runtime::Runtime;
 use tokio::time::sleep;
 
 use crate::{
@@ -64,7 +63,7 @@ pub trait Client {
     fn get_config(&self) -> &SharedConfig;
 
     fn send_message(&self, content: &str) -> Result<String> {
-        init_runtime()?.block_on(async {
+        init_tokio_runtime()?.block_on(async {
             if self.get_config().read().dry_run {
                 return Ok(self.get_config().read().echo_messages(content));
             }
@@ -88,7 +87,7 @@ pub trait Client {
             }
         }
         let abort = handler.get_abort();
-        init_runtime()?.block_on(async {
+        init_tokio_runtime()?.block_on(async {
             tokio::select! {
                 ret = async {
                     if self.get_config().read().dry_run {
@@ -164,7 +163,7 @@ pub fn list_models(config: &Config) -> Vec<ModelInfo> {
         .collect()
 }
 
-pub fn init_runtime() -> Result<Runtime> {
+pub fn init_tokio_runtime() -> Result<tokio::runtime::Runtime> {
     tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 64a3299..922fb0a 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -129,6 +129,7 @@ impl Config {
         config.merge_env_vars();
         config.load_roles()?;
         config.ensure_sessions_dir()?;
+        config.check_term_theme()?;
 
         Ok(config)
     }
@@ -546,9 +547,6 @@ impl Config {
     }
 
     fn merge_env_vars(&mut self) {
-        if let Ok(value) = env::var(get_env_name("light_theme")) {
-            set_bool(&mut self.light_theme, &value);
-        }
         if let Ok(value) = env::var("NO_COLOR") {
             let mut no_color = false;
             set_bool(&mut no_color, &value);
@@ -568,6 +566,23 @@ impl Config {
         Ok(())
     }
 
+    fn check_term_theme(&mut self) -> Result<()> {
+        if self.light_theme {
+            return Ok(());
+        }
+        if let Ok(value) = env::var(get_env_name("light_theme")) {
+            set_bool(&mut self.light_theme, &value);
+            return Ok(());
+        }
+        #[cfg(not(target_os = "windows"))]
+        if let Ok(crate::utils::termbg::Theme::Light) =
+            crate::utils::termbg::theme(std::time::Duration::from_millis(200))
+        {
+            self.light_theme = true;
+        }
+        Ok(())
+    }
+
     fn compat_old_config(&mut self, config_path: &PathBuf) -> Result<()> {
         let content = read_to_string(config_path)?;
         let value: serde_json::Value = serde_yaml::from_str(&content)?;
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 4ed0d5a..fa26dcf 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -1,3 +1,5 @@
+#[cfg(not(target_os = "windows"))]
+pub(crate) mod termbg;
 mod tiktoken;
 
 use self::tiktoken::cl100k_base;
diff --git a/src/utils/termbg.rs b/src/utils/termbg.rs
new file mode 100644
index 0000000..2322227
--- /dev/null
+++ b/src/utils/termbg.rs
@@ -0,0 +1,257 @@
+//! Terminal background color detection
+/// Fork from https://github.com/dalance/termbg/blob/v0.4.3/src/lib.rs
+use anyhow::{anyhow, Error, Result};
+use crossterm::terminal;
+use std::env;
+use std::io::{self, Read, Write};
+use std::time::Duration;
+
+/// Terminal
+#[derive(Copy, Clone, Debug, PartialEq)]
+pub enum Terminal {
+    Screen,
+    Tmux,
+    XtermCompatible,
+    VSCode,
+    Emacs,
+}
+
+/// 16bit RGB color
+#[derive(Copy, Clone, Debug)]
+pub struct Rgb {
+    pub r: u16,
+    pub g: u16,
+    pub b: u16,
+}
+
+/// Background theme
+#[derive(Copy, Clone, Debug)]
+pub enum Theme {
+    Light,
+    Dark,
+}
+
+/// get detected termnial
+pub fn terminal() -> Terminal {
+    if let Ok(term_program) = env::var("TERM_PROGRAM") {
+        if term_program == "vscode" {
+            return Terminal::VSCode;
+        }
+    }
+
+    if env::var("INSIDE_EMACS").is_ok() {
+        return Terminal::Emacs;
+    }
+
+    if env::var("TMUX").is_ok() {
+        Terminal::Tmux
+    } else {
+        let is_screen = if let Ok(term) = env::var("TERM") {
+            term.starts_with("screen")
+        } else {
+            false
+        };
+        if is_screen {
+            Terminal::Screen
+        } else {
+            Terminal::XtermCompatible
+        }
+    }
+}
+
+/// get background color by `RGB`
+pub fn rgb(timeout: Duration) -> Result<Rgb> {
+    let term = terminal();
+    let rgb = match term {
+        Terminal::VSCode => Err(unsupported_err()),
+        Terminal::Emacs => Err(unsupported_err()),
+        _ => from_xterm(term, timeout),
+    };
+    check_rgb(rgb)
+}
+
+/// get background color by `Theme`
+pub fn theme(timeout: Duration) -> Result<Theme> {
+    let rgb = rgb(timeout)?;
+
+    // ITU-R BT.601
+    let y = rgb.r as f64 * 0.299 + rgb.g as f64 * 0.587 + rgb.b as f64 * 0.114;
+
+    if y > 32768.0 {
+        Ok(Theme::Light)
+    } else {
+        Ok(Theme::Dark)
+    }
+}
+
+fn from_xterm(term: Terminal, timeout: Duration) -> Result<Rgb> {
+    // Query by XTerm control sequence
+    let query = if term == Terminal::Tmux {
+        "\x1bPtmux;\x1b\x1b]11;?\x07\x1b\\\x03"
+    } else if term == Terminal::Screen {
+        "\x1bP\x1b]11;?\x07\x1b\\\x03"
+    } else {
+        "\x1b]11;?\x1b\\"
+    };
+
+    let mut stderr = io::stderr();
+    terminal::enable_raw_mode()?;
+    write!(stderr, "{}", query)?;
+    stderr.flush()?;
+
+    let mut stdin = io::stdin();
+
+    let (tx, rx) = std::sync::mpsc::channel();
+
+    std::thread::spawn(move || {
+        let mut buffer = Vec::new();
+        let mut buf = [0; 1];
+        let mut start = false;
+        loop {
+            let _ = stdin.read_exact(&mut buf);
+            // response terminated by BEL(0x7)
+            if start && (buf[0] == 0x7) {
+                break;
+            }
+            // response terminated by ST(0x1b 0x5c)
+            if start && (buf[0] == 0x1b) {
+                // consume last 0x5c
+                let _ = stdin.read_exact(&mut buf);
+                debug_assert_eq!(buf[0], 0x5c);
+                break;
+            }
+            if start {
+                buffer.push(buf[0]);
+            }
+            if buf[0] == b':' {
+                start = true;
+            }
+        }
+        // Ignore send error because timeout may be occured
+        let _ = tx.send(buffer);
+    });
+
+    let buffer = rx.recv_timeout(timeout);
+
+    terminal::disable_raw_mode()?;
+
+    let buffer = buffer?;
+
+    let s = String::from_utf8_lossy(&buffer);
+    let (r, g, b) = decode_x11_color(&s)?;
+    Ok(Rgb { r, g, b })
+}
+
+fn decode_x11_color(s: &str) -> Result<(u16, u16, u16)> {
+    fn decode_hex(s: &str) -> Result<u16> {
+        let len = s.len() as u32;
+        let mut ret = u16::from_str_radix(s, 16).map_err(|_| parse_err(s))?;
+        ret <<= (4 - len) * 4;
+        Ok(ret)
+    }
+
+    let rgb: Vec<_> = s.split('/').collect();
+
+    let r = rgb.first().ok_or_else(|| parse_err(s))?;
+    let g = rgb.get(1).ok_or_else(|| parse_err(s))?;
+    let b = rgb.get(2).ok_or_else(|| parse_err(s))?;
+    let r = decode_hex(r)?;
+    let g = decode_hex(g)?;
+    let b = decode_hex(b)?;
+
+    Ok((r, g, b))
+}
+
+fn check_rgb(prev: Result<Rgb>) -> Result<Rgb> {
+    if prev.is_ok() {
+        return prev;
+    }
+    let fallback = from_env_colorfgbg();
+    if fallback.is_ok() {
+        return fallback;
+    }
+    prev
+}
+
+fn from_env_colorfgbg() -> Result<Rgb> {
+    let var = env::var("COLORFGBG").map_err(|_| unsupported_err())?;
+    let fgbg: Vec<_> = var.split(';').collect();
+    let bg = fgbg.get(1).ok_or(unsupported_err())?;
+    let bg = bg.parse::<u8>().map_err(|_| parse_err(&var))?;
+
+    // rxvt default color table
+    let (r, g, b) = match bg {
+        // black
+        0 => (0, 0, 0),
+        // red
+        1 => (205, 0, 0),
+        // green
+        2 => (0, 205, 0),
+        // yellow
+        3 => (205, 205, 0),
+        // blue
+        4 => (0, 0, 238),
+        // magenta
+        5 => (205, 0, 205),
+
+        // cyan
+        6 => (0, 205, 205),
+        // white
+        7 => (229, 229, 229),
+        // bright black
+        8 => (127, 127, 127),
+        // bright red
+        9 => (255, 0, 0),
+        // bright green
+        10 => (0, 255, 0),
+        // bright yellow
+        11 => (255, 255, 0),
+        // bright blue
+        12 => (92, 92, 255),
+        // bright magenta
+        13 => (255, 0, 255),
+        // bright cyan
+        14 => (0, 255, 255),
+
+        // bright white
+        15 => (255, 255, 255),
+        _ => (0, 0, 0),
+    };
+
+    Ok(Rgb {
+        r: r * 256,
+        g: g * 256,
+        b: b * 256,
+    })
+}
+
+fn unsupported_err() -> Error {
+    anyhow!("Unsupported terminal")
+}
+
+fn parse_err(value: &str) -> Error {
+    anyhow!("Failed to parse {value}")
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_decode_x11_color() {
+        let s = "0000/0000/0000";
+        assert_eq!((0, 0, 0), decode_x11_color(s).unwrap());
+
+        let s = "1111/2222/3333";
+        assert_eq!((0x1111, 0x2222, 0x3333), decode_x11_color(s).unwrap());
+
+        let s = "111/222/333";
+        assert_eq!((0x1110, 0x2220, 0x3330), decode_x11_color(s).unwrap());
+
+        let s = "11/22/33";
+        assert_eq!((0x1100, 0x2200, 0x3300), decode_x11_color(s).unwrap());
+
+        let s = "1/2/3";
+        assert_eq!((0x1000, 0x2000, 0x3000), decode_x11_color(s).unwrap());
+    }
+}