Baseline project

2025-10-07 10:45:42 -06:00
parent 88288a98b6
commit 650dbd92e0
54 changed files with 18982 additions and 0 deletions
@@ -0,0 +1,393 @@
+use crate::utils::decode_bin;
+
+use ansi_colours::AsRGB;
+use anyhow::{anyhow, Context, Result};
+use crossterm::style::{Color, Stylize};
+use crossterm::terminal;
+use std::collections::HashMap;
+use std::sync::LazyLock;
+use syntect::highlighting::{Color as SyntectColor, FontStyle, Style, Theme};
+use syntect::parsing::SyntaxSet;
+use syntect::{easy::HighlightLines, parsing::SyntaxReference};
+
+/// Comes from <https://github.com/sharkdp/bat/raw/5e77ca37e89c873e4490b42ff556370dc5c6ba4f/assets/syntaxes.bin>
+const SYNTAXES: &[u8] = include_bytes!("../../assets/syntaxes.bin");
+
+static LANG_MAPS: LazyLock<HashMap<String, String>> = LazyLock::new(|| {
+    let mut m = HashMap::new();
+    m.insert("csharp".into(), "C#".into());
+    m.insert("php".into(), "PHP Source".into());
+    m
+});
+
+pub struct MarkdownRender {
+    options: RenderOptions,
+    syntax_set: SyntaxSet,
+    code_color: Option<Color>,
+    md_syntax: SyntaxReference,
+    code_syntax: Option<SyntaxReference>,
+    prev_line_type: LineType,
+    wrap_width: Option<u16>,
+}
+
+impl MarkdownRender {
+    pub fn init(options: RenderOptions) -> Result<Self> {
+        let syntax_set: SyntaxSet =
+            decode_bin(SYNTAXES).with_context(|| "MarkdownRender: invalid syntaxes binary")?;
+
+        let code_color = options
+            .theme
+            .as_ref()
+            .map(|theme| get_code_color(theme, options.truecolor));
+        let md_syntax = syntax_set.find_syntax_by_extension("md").unwrap().clone();
+        let line_type = LineType::Normal;
+        let wrap_width = match options.wrap.as_deref() {
+            None => None,
+            Some(value) => match terminal::size() {
+                Ok((columns, _)) => {
+                    if value == "auto" {
+                        Some(columns)
+                    } else {
+                        let value = value
+                            .parse::<u16>()
+                            .map_err(|_| anyhow!("Invalid wrap value"))?;
+                        Some(columns.min(value))
+                    }
+                }
+                Err(_) => None,
+            },
+        };
+        Ok(Self {
+            syntax_set,
+            code_color,
+            md_syntax,
+            code_syntax: None,
+            prev_line_type: line_type,
+            wrap_width,
+            options,
+        })
+    }
+
+    pub fn render(&mut self, text: &str) -> String {
+        text.split('\n')
+            .map(|line| self.render_line_mut(line))
+            .collect::<Vec<String>>()
+            .join("\n")
+    }
+
+    pub fn render_line(&self, line: &str) -> String {
+        let (_, code_syntax, is_code) = self.check_line(line);
+        if is_code {
+            self.highlight_code_line(line, &code_syntax)
+        } else {
+            self.highlight_line(line, &self.md_syntax, false)
+        }
+    }
+
+    fn render_line_mut(&mut self, line: &str) -> String {
+        let (line_type, code_syntax, is_code) = self.check_line(line);
+        let output = if is_code {
+            self.highlight_code_line(line, &code_syntax)
+        } else {
+            self.highlight_line(line, &self.md_syntax, false)
+        };
+        self.prev_line_type = line_type;
+        self.code_syntax = code_syntax;
+        output
+    }
+
+    fn check_line(&self, line: &str) -> (LineType, Option<SyntaxReference>, bool) {
+        let mut line_type = self.prev_line_type;
+        let mut code_syntax = self.code_syntax.clone();
+        let mut is_code = false;
+        if let Some(lang) = detect_code_block(line) {
+            match line_type {
+                LineType::Normal | LineType::CodeEnd => {
+                    line_type = LineType::CodeBegin;
+                    code_syntax = if lang.is_empty() {
+                        None
+                    } else {
+                        self.find_syntax(&lang).cloned()
+                    };
+                }
+                LineType::CodeBegin | LineType::CodeInner => {
+                    line_type = LineType::CodeEnd;
+                    code_syntax = None;
+                }
+            }
+        } else {
+            match line_type {
+                LineType::Normal => {}
+                LineType::CodeEnd => {
+                    line_type = LineType::Normal;
+                }
+                LineType::CodeBegin => {
+                    if code_syntax.is_none() {
+                        if let Some(syntax) = self.syntax_set.find_syntax_by_first_line(line) {
+                            code_syntax = Some(syntax.clone());
+                        }
+                    }
+                    line_type = LineType::CodeInner;
+                    is_code = true;
+                }
+                LineType::CodeInner => {
+                    is_code = true;
+                }
+            }
+        }
+        (line_type, code_syntax, is_code)
+    }
+
+    fn highlight_line(&self, line: &str, syntax: &SyntaxReference, is_code: bool) -> String {
+        let ws: String = line.chars().take_while(|c| c.is_whitespace()).collect();
+        let trimmed_line: &str = &line[ws.len()..];
+        let mut line_highlighted = None;
+        if let Some(theme) = &self.options.theme {
+            let mut highlighter = HighlightLines::new(syntax, theme);
+            if let Ok(ranges) = highlighter.highlight_line(trimmed_line, &self.syntax_set) {
+                line_highlighted = Some(format!(
+                    "{ws}{}",
+                    as_terminal_escaped(&ranges, self.options.truecolor)
+                ))
+            }
+        }
+        let line = line_highlighted.unwrap_or_else(|| line.into());
+        self.wrap_line(line, is_code)
+    }
+
+    fn highlight_code_line(&self, line: &str, code_syntax: &Option<SyntaxReference>) -> String {
+        if let Some(syntax) = code_syntax {
+            self.highlight_line(line, syntax, true)
+        } else {
+            let line = match self.code_color {
+                Some(color) => line.with(color).to_string(),
+                None => line.to_string(),
+            };
+            self.wrap_line(line, true)
+        }
+    }
+
+    fn wrap_line(&self, line: String, is_code: bool) -> String {
+        if let Some(width) = self.wrap_width {
+            if is_code && !self.options.wrap_code {
+                return line;
+            }
+            wrap(&line, width as usize)
+        } else {
+            line
+        }
+    }
+
+    fn find_syntax(&self, lang: &str) -> Option<&SyntaxReference> {
+        if let Some(new_lang) = LANG_MAPS.get(&lang.to_ascii_lowercase()) {
+            self.syntax_set.find_syntax_by_name(new_lang)
+        } else {
+            self.syntax_set
+                .find_syntax_by_token(lang)
+                .or_else(|| self.syntax_set.find_syntax_by_extension(lang))
+        }
+    }
+}
+
+fn wrap(text: &str, width: usize) -> String {
+    let indent: usize = text.chars().take_while(|c| *c == ' ').count();
+    let wrap_options = textwrap::Options::new(width)
+        .wrap_algorithm(textwrap::WrapAlgorithm::FirstFit)
+        .initial_indent(&text[0..indent]);
+    textwrap::wrap(&text[indent..], wrap_options).join("\n")
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct RenderOptions {
+    pub theme: Option<Theme>,
+    pub wrap: Option<String>,
+    pub wrap_code: bool,
+    pub truecolor: bool,
+}
+
+impl RenderOptions {
+    pub(crate) fn new(
+        theme: Option<Theme>,
+        wrap: Option<String>,
+        wrap_code: bool,
+        truecolor: bool,
+    ) -> Self {
+        Self {
+            theme,
+            wrap,
+            wrap_code,
+            truecolor,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LineType {
+    Normal,
+    CodeBegin,
+    CodeInner,
+    CodeEnd,
+}
+
+fn as_terminal_escaped(ranges: &[(Style, &str)], truecolor: bool) -> String {
+    let mut output = String::new();
+    for (style, text) in ranges {
+        let fg = blend_fg_color(style.foreground, style.background);
+        let mut text = text.with(convert_color(fg, truecolor));
+        if style.font_style.contains(FontStyle::BOLD) {
+            text = text.bold();
+        }
+        if style.font_style.contains(FontStyle::UNDERLINE) {
+            text = text.underlined();
+        }
+        output.push_str(&text.to_string());
+    }
+    output
+}
+
+fn convert_color(c: SyntectColor, truecolor: bool) -> Color {
+    if truecolor {
+        Color::Rgb {
+            r: c.r,
+            g: c.g,
+            b: c.b,
+        }
+    } else {
+        let value = (c.r, c.g, c.b).to_ansi256();
+        // lower contrast
+        let value = match value {
+            7 | 15 | 231 | 252..=255 => 252,
+            _ => value,
+        };
+        Color::AnsiValue(value)
+    }
+}
+
+fn blend_fg_color(fg: SyntectColor, bg: SyntectColor) -> SyntectColor {
+    if fg.a == 0xff {
+        return fg;
+    }
+    let ratio = u32::from(fg.a);
+    let r = (u32::from(fg.r) * ratio + u32::from(bg.r) * (255 - ratio)) / 255;
+    let g = (u32::from(fg.g) * ratio + u32::from(bg.g) * (255 - ratio)) / 255;
+    let b = (u32::from(fg.b) * ratio + u32::from(bg.b) * (255 - ratio)) / 255;
+    SyntectColor {
+        r: u8::try_from(r).unwrap_or(u8::MAX),
+        g: u8::try_from(g).unwrap_or(u8::MAX),
+        b: u8::try_from(b).unwrap_or(u8::MAX),
+        a: 255,
+    }
+}
+
+fn detect_code_block(line: &str) -> Option<String> {
+    let line = line.trim_start();
+    if !line.starts_with("```") {
+        return None;
+    }
+    let lang = line
+        .chars()
+        .skip(3)
+        .take_while(|v| !v.is_whitespace())
+        .collect();
+    Some(lang)
+}
+
+fn get_code_color(theme: &Theme, truecolor: bool) -> Color {
+    let scope = theme.scopes.iter().find(|v| {
+        v.scope
+            .selectors
+            .iter()
+            .any(|v| v.path.scopes.iter().any(|v| v.to_string() == "string"))
+    });
+    scope
+        .and_then(|v| v.style.foreground)
+        .map_or_else(|| Color::Yellow, |c| convert_color(c, truecolor))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const TEXT: &str = r#"
+To unzip a file in Rust, you can use the `zip` crate. Here's an example code that shows how to unzip a file:
+
+```rust
+use std::fs::File;
+
+fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box<dyn std::error::Error>> {
+    todo!()
+}
+```
+"#;
+    const TEXT_NO_WRAP_CODE: &str = r#"
+To unzip a file in Rust, you can use the `zip` crate. Here's an example code
+that shows how to unzip a file:
+
+```rust
+use std::fs::File;
+
+fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box<dyn std::error::Error>> {
+    todo!()
+}
+```
+"#;
+
+    const TEXT_WRAP_ALL: &str = r#"
+To unzip a file in Rust, you can use the `zip` crate. Here's an example code
+that shows how to unzip a file:
+
+```rust
+use std::fs::File;
+
+fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box<dyn
+std::error::Error>> {
+    todo!()
+}
+```
+"#;
+
+    #[test]
+    fn test_render() {
+        let options = RenderOptions::default();
+        let render = MarkdownRender::init(options).unwrap();
+        assert!(render.find_syntax("csharp").is_some());
+    }
+
+    #[test]
+    fn no_theme() {
+        let options = RenderOptions::default();
+        let mut render = MarkdownRender::init(options).unwrap();
+        let output = render.render(TEXT);
+        assert_eq!(TEXT, output);
+    }
+
+    #[test]
+    fn no_wrap_code() {
+        let options = RenderOptions::default();
+        let mut render = MarkdownRender::init(options).unwrap();
+        render.wrap_width = Some(80);
+        let output = render.render(TEXT);
+        assert_eq!(TEXT_NO_WRAP_CODE, output);
+    }
+
+    #[test]
+    fn wrap_all() {
+        let options = RenderOptions {
+            wrap_code: true,
+            ..Default::default()
+        };
+        let mut render = MarkdownRender::init(options).unwrap();
+        render.wrap_width = Some(80);
+        let output = render.render(TEXT);
+        assert_eq!(TEXT_WRAP_ALL, output);
+    }
+
+    #[test]
+    fn test_detect_code_block() {
+        assert_eq!(detect_code_block("```rust"), Some("rust".into()));
+        assert_eq!(detect_code_block("```c++"), Some("c++".into()));
+        assert_eq!(detect_code_block("  ```rust"), Some("rust".into()));
+        assert_eq!(detect_code_block("```"), Some("".into()));
+        assert_eq!(detect_code_block("``rust"), None);
+    }
+}
@@ -0,0 +1,30 @@
+mod markdown;
+mod stream;
+
+pub use self::markdown::{MarkdownRender, RenderOptions};
+use self::stream::{markdown_stream, raw_stream};
+
+use crate::utils::{error_text, pretty_error, AbortSignal, IS_STDOUT_TERMINAL};
+use crate::{client::SseEvent, config::GlobalConfig};
+
+use anyhow::Result;
+use tokio::sync::mpsc::UnboundedReceiver;
+
+pub async fn render_stream(
+    rx: UnboundedReceiver<SseEvent>,
+    config: &GlobalConfig,
+    abort_signal: AbortSignal,
+) -> Result<()> {
+    let ret = if *IS_STDOUT_TERMINAL && config.read().highlight {
+        let render_options = config.read().render_options()?;
+        let mut render = MarkdownRender::init(render_options)?;
+        markdown_stream(rx, &mut render, &abort_signal).await
+    } else {
+        raw_stream(rx, &abort_signal).await
+    };
+    ret.map_err(|err| err.context("Failed to reader stream"))
+}
+
+pub fn render_error(err: anyhow::Error) {
+    eprintln!("{}", error_text(&pretty_error(&err)));
+}
@@ -0,0 +1,217 @@
+use super::{MarkdownRender, SseEvent};
+
+use crate::utils::{poll_abort_signal, spawn_spinner, AbortSignal};
+
+use anyhow::Result;
+use crossterm::{
+    cursor, queue, style,
+    terminal::{self, disable_raw_mode, enable_raw_mode},
+};
+use std::{
+    io::{stdout, Stdout, Write},
+    time::Duration,
+};
+use textwrap::core::display_width;
+use tokio::sync::mpsc::UnboundedReceiver;
+
+pub async fn markdown_stream(
+    rx: UnboundedReceiver<SseEvent>,
+    render: &mut MarkdownRender,
+    abort_signal: &AbortSignal,
+) -> Result<()> {
+    enable_raw_mode()?;
+    let mut stdout = stdout();
+
+    let ret = markdown_stream_inner(rx, render, abort_signal, &mut stdout).await;
+
+    disable_raw_mode()?;
+
+    if ret.is_err() {
+        println!();
+    }
+    ret
+}
+
+pub async fn raw_stream(
+    mut rx: UnboundedReceiver<SseEvent>,
+    abort_signal: &AbortSignal,
+) -> Result<()> {
+    let mut spinner = Some(spawn_spinner("Generating"));
+
+    loop {
+        if abort_signal.aborted() {
+            break;
+        }
+        if let Some(evt) = rx.recv().await {
+            if let Some(spinner) = spinner.take() {
+                spinner.stop();
+            }
+
+            match evt {
+                SseEvent::Text(text) => {
+                    print!("{text}");
+                    stdout().flush()?;
+                }
+                SseEvent::Done => {
+                    break;
+                }
+            }
+        }
+    }
+    if let Some(spinner) = spinner.take() {
+        spinner.stop();
+    }
+    Ok(())
+}
+
+async fn markdown_stream_inner(
+    mut rx: UnboundedReceiver<SseEvent>,
+    render: &mut MarkdownRender,
+    abort_signal: &AbortSignal,
+    writer: &mut Stdout,
+) -> Result<()> {
+    let mut buffer = String::new();
+    let mut buffer_rows = 1;
+
+    let columns = terminal::size()?.0;
+
+    let mut spinner = Some(spawn_spinner("Generating"));
+
+    'outer: loop {
+        if abort_signal.aborted() {
+            break;
+        }
+        for reply_event in gather_events(&mut rx).await {
+            if let Some(spinner) = spinner.take() {
+                spinner.stop();
+            }
+
+            match reply_event {
+                SseEvent::Text(mut text) => {
+                    // tab width hacking
+                    text = text.replace('\t', "    ");
+
+                    let mut attempts = 0;
+                    let (col, mut row) = loop {
+                        match cursor::position() {
+                            Ok(pos) => break pos,
+                            Err(_) if attempts < 3 => attempts += 1,
+                            Err(e) => return Err(e.into()),
+                        }
+                    };
+
+                    // Fix unexpected duplicate lines on kitty
+                    if col == 0 && row > 0 && display_width(&buffer) == columns as usize {
+                        row -= 1;
+                    }
+
+                    if row + 1 >= buffer_rows {
+                        queue!(writer, cursor::MoveTo(0, row + 1 - buffer_rows),)?;
+                    } else {
+                        let scroll_rows = buffer_rows - row - 1;
+                        queue!(
+                            writer,
+                            terminal::ScrollUp(scroll_rows),
+                            cursor::MoveTo(0, 0),
+                        )?;
+                    }
+
+                    // No guarantee that text returned by render will not be re-layouted, so it is better to clear it.
+                    queue!(writer, terminal::Clear(terminal::ClearType::FromCursorDown))?;
+
+                    if text.contains('\n') {
+                        let text = format!("{buffer}{text}");
+                        let (head, tail) = split_line_tail(&text);
+                        let output = render.render(head);
+                        print_block(writer, &output, columns)?;
+                        buffer = tail.to_string();
+                    } else {
+                        buffer = format!("{buffer}{text}");
+                    }
+
+                    let output = render.render_line(&buffer);
+                    if output.contains('\n') {
+                        let (head, tail) = split_line_tail(&output);
+                        buffer_rows = print_block(writer, head, columns)?;
+                        queue!(writer, style::Print(&tail),)?;
+
+                        // No guarantee the buffer width of the buffer will not exceed the number of columns.
+                        // So we calculate the number of rows needed, rather than setting it directly to 1.
+                        buffer_rows += need_rows(tail, columns);
+                    } else {
+                        queue!(writer, style::Print(&output))?;
+                        buffer_rows = need_rows(&output, columns);
+                    }
+
+                    writer.flush()?;
+                }
+                SseEvent::Done => {
+                    break 'outer;
+                }
+            }
+        }
+
+        if poll_abort_signal(abort_signal)? {
+            break;
+        }
+    }
+
+    if let Some(spinner) = spinner.take() {
+        spinner.stop();
+    }
+    Ok(())
+}
+
+async fn gather_events(rx: &mut UnboundedReceiver<SseEvent>) -> Vec<SseEvent> {
+    let mut texts = vec![];
+    let mut done = false;
+    tokio::select! {
+        _ = async {
+            while let Some(reply_event) = rx.recv().await {
+                match reply_event {
+                    SseEvent::Text(v) => texts.push(v),
+                    SseEvent::Done => {
+                        done = true;
+                        break;
+                    }
+                }
+            }
+        } => {}
+        _ = tokio::time::sleep(Duration::from_millis(50)) => {}
+    }
+    let mut events = vec![];
+    if !texts.is_empty() {
+        events.push(SseEvent::Text(texts.join("")))
+    }
+    if done {
+        events.push(SseEvent::Done)
+    }
+    events
+}
+
+fn print_block(writer: &mut Stdout, text: &str, columns: u16) -> Result<u16> {
+    let mut num = 0;
+    for line in text.split('\n') {
+        queue!(
+            writer,
+            style::Print(line),
+            style::Print("\n"),
+            cursor::MoveLeft(columns),
+        )?;
+        num += 1;
+    }
+    Ok(num)
+}
+
+fn split_line_tail(text: &str) -> (&str, &str) {
+    if let Some((head, tail)) = text.rsplit_once('\n') {
+        (head, tail)
+    } else {
+        ("", text)
+    }
+}
+
+fn need_rows(text: &str, columns: u16) -> u16 {
+    let buffer_width = display_width(text).max(1) as u16;
+    buffer_width.div_ceil(columns)
+}