live_markdown/

lib.rs

#[cfg(feature = "python")]
use pyo3::prelude::*;

use regex::Regex;
use std::collections::HashSet;
use std::thread;
use std::time::Duration;
use termion::{color, cursor, style};

// Syntax highlighting imports
use syntect::easy::HighlightLines;
use syntect::highlighting::{Style, ThemeSet};
use syntect::parsing::SyntaxSet;
use syntect::util::{LinesWithEndings, as_24_bit_terminal_escaped};

/// Core engine for streaming Markdown to a terminal with live formatting.
/// Handles word wrapping, syntax highlighting for code blocks, and nested block styles.
pub struct MarkdownStreamerCore {
    buffer: String,
    // Internal byte buffer to handle fragmented UTF-8 streams
    byte_cache: Vec<u8>,
    output_buffer: String,
    code_block_depth: usize,
    depth_stack: Vec<String>,
    current_code_lang: String,
    current_code_line: String,
    code_history: String,
    at_line_start: bool,
    active_styles: HashSet<String>,
    word_buffer: String,
    line_pos: u16,
    term_width: u16,
    ps: SyntaxSet,
    ts: ThemeSet,
    /// Delay between processing characters for a "smooth" typewriter effect
    render_delay: Duration,
    /// Minimum buffer size before starting the initial render
    initial_buffer_threshold: usize,
    /// Flag to track if the initial buffer has been cleared
    initial_buffer_passed: bool,
}

impl MarkdownStreamerCore {
    /// Background colors for nested code blocks (ANSI 256-color palette)
    const BGS: &'static [&'static str] = &[
        "\x1b[48;5;235m",
        "\x1b[48;5;237m",
        "\x1b[48;5;239m",
        "\x1b[48;5;241m",
    ];

    pub fn new() -> Self {
        let (width, _) = termion::terminal_size().unwrap_or((80, 24));
        Self {
            buffer: String::new(),
            byte_cache: Vec::new(),
            output_buffer: String::new(),
            code_block_depth: 0,
            depth_stack: Vec::new(),
            current_code_lang: String::new(),
            current_code_line: String::new(),
            code_history: String::new(),
            at_line_start: true,
            active_styles: HashSet::new(),
            word_buffer: String::new(),
            line_pos: 0,
            term_width: width,
            ps: SyntaxSet::load_defaults_newlines(),
            ts: ThemeSet::load_defaults(),
            render_delay: Duration::from_millis(0),
            initial_buffer_threshold: 20,
            initial_buffer_passed: false,
        }
    }

    /// Pulls processed terminal output from the internal buffer.
    pub fn read_available(&mut self) -> String {
        std::mem::take(&mut self.output_buffer)
    }

    fn get_block_bg(&self) -> String {
        if self.code_block_depth == 0 {
            return String::new();
        }
        let idx = (self.code_block_depth - 1).min(Self::BGS.len() - 1);
        Self::BGS[idx].to_string()
    }

    fn apply_indentation(&mut self) {
        if self.code_block_depth == 0 {
            return;
        }
        let move_len = (self.code_block_depth as u16 - 1) * 4;
        if move_len > 0 {
            self.output_buffer
                .push_str(&format!("{}", cursor::Right(move_len)));
            self.line_pos += move_len;
        }
    }

    /// Transforms raw code text into syntax-highlighted ANSI strings.
    fn apply_highlighting(&self, text: &str, lang: &str) -> String {
        let bg = self.get_block_bg();
        if lang == "markdown" || (text.is_empty() && self.code_block_depth > 0) {
            return format!("{}{}\x1b[K", bg, text);
        }

        let syntax = self
            .ps
            .find_syntax_by_token(lang)
            .unwrap_or_else(|| self.ps.find_syntax_plain_text());
        let mut h = HighlightLines::new(syntax, &self.ts.themes["base16-ocean.dark"]);

        // We use code history to maintain highlighting state across incremental chunks
        let mut full_content = self.code_history.clone();
        full_content.push_str(text);

        let mut result = String::new();
        for line in LinesWithEndings::from(&full_content) {
            let ranges: Vec<(Style, &str)> = h.highlight_line(line, &self.ps).unwrap();
            result = as_24_bit_terminal_escaped(&ranges[..], false);
        }

        // Clean up backgrounds and resets to ensure block background consistency
        let clean_bg_regex = Regex::new(r"\x1b\[48;[0-9;]*m").unwrap();
        let stripped = clean_bg_regex.replace_all(&result, "");
        let reset_regex = Regex::new(r"\x1b\[0?m").unwrap();
        let no_resets = reset_regex.replace_all(&stripped, "");
        let esc_regex = Regex::new(r"(\x1b\[[0-9;]*m)").unwrap();
        let fixed_result = esc_regex.replace_all(&no_resets, format!("$1{}", bg));

        format!("{}{}{}\x1b[K", bg, fixed_result, bg)
    }

    /// Handles terminal cursor movement to overwrite lines during live updates.
    fn clear_and_move_up(&mut self, text: String) {
        let re = Regex::new(r"\x1b\[[0-9;]*m").unwrap();
        let stripped = re.replace_all(&text, "");
        let total_len = stripped.chars().count() as u16;

        // Calculate visual line wrap count
        let extra_rows = if total_len > 0 {
            (self.line_pos + total_len - 1) / self.term_width
        } else {
            0
        };

        self.output_buffer.push('\r');
        for _ in 0..extra_rows {
            self.output_buffer.push_str("\x1b[K"); // Clear current line
            self.output_buffer.push_str("\x1b[A"); // Move up
        }

        self.output_buffer.push_str("\x1b[K");
        self.output_buffer.push('\r');

        self.apply_indentation();
    }

    /// Flushes the current word buffer, handling word-wrap logic if the word exceeds terminal width.
    fn flush_word(&mut self, prefix: &str) {
        if self.word_buffer.is_empty() {
            return;
        }
        let re = Regex::new(r"\x1b\[[0-9;]*m").unwrap();
        let visible_word = re.replace_all(&self.word_buffer, "").into_owned();
        let visible_len = visible_word.chars().count() as u16;

        if self.line_pos + visible_len >= self.term_width - 1 {
            let bg = self.get_block_bg();
            self.output_buffer
                .push_str(&format!("{}\x1b[K\n{}", bg, style::Reset));
            self.line_pos = 0;
            self.apply_indentation();
            self.output_buffer.push_str(prefix);
            if self.code_block_depth > 0 {
                self.output_buffer.push_str("\x1b[K");
            }
        }
        self.output_buffer.push_str(&self.word_buffer);
        self.line_pos += visible_len;
        self.word_buffer.clear();
    }

    fn toggle_style(&mut self, style_key: &str) {
        let bg = self.get_block_bg();
        if self.active_styles.contains(style_key) {
            self.active_styles.remove(style_key);
        } else {
            self.active_styles.insert(style_key.to_string());
        }
        self.word_buffer.push_str(&format!("{}", style::Reset));
        self.word_buffer.push_str(&bg);
        for s in &self.active_styles {
            match s.as_str() {
                "bold" => self.word_buffer.push_str(&format!("{}", style::Bold)),
                "italic" => self.word_buffer.push_str(&format!("{}", style::Italic)),
                "code" => {
                    self.word_buffer
                        .push_str(&format!("{}", color::Fg(color::Yellow)));
                    self.word_buffer.push_str("\x1b[48;5;238m");
                }
                _ => {}
            }
        }
    }

    /// Main loop for parsing the incoming Markdown stream.
    pub fn process_buffer(&mut self, final_call: bool) {
        let re_close = Regex::new(r"^[ \t]*```[ \t]*\n?").unwrap();
        let re_open = Regex::new(r"^[ \t]*```([a-zA-Z0-9\-\+#]+)[ \t]*\n?").unwrap();
        let re_hr = Regex::new(r"^[ \t]*(\-{3,}|\*{3,}|\_{3,})[ \t]*\n?").unwrap();

        // Ensure we have enough data for a meaningful initial render
        if !final_call && !self.initial_buffer_passed {
            if self.buffer.chars().count() < self.initial_buffer_threshold {
                return;
            } else {
                self.initial_buffer_passed = true;
            }
        }

        while !self.buffer.is_empty() || (final_call && !self.word_buffer.is_empty()) {
            if final_call && self.buffer.is_empty() && !self.word_buffer.is_empty() {
                let bg = self.get_block_bg();
                self.flush_word(&bg);
                if self.buffer.is_empty() {
                    break;
                }
            }

            // LOOKAHEAD GUARD:
            if !final_call && self.buffer.chars().count() < 15 {
                if self.at_line_start
                    || self.buffer.contains('`')
                    || self.buffer.contains('*')
                    || self.buffer.contains('_')
                {
                    break;
                }
            }

            if !self.render_delay.is_zero() {
                thread::sleep(self.render_delay);
            }

            let bg = self.get_block_bg();

            // Handling inside Code Blocks
            if self.code_block_depth > 0 && self.current_code_lang != "markdown" {
                let close_match = re_close.find(&self.buffer).map(|m| m.end());
                if let Some(end_idx) = close_match {
                    if self.at_line_start {
                        if !self.current_code_line.is_empty() {
                            let line_content = self.current_code_line.clone();
                            let lang = self.current_code_lang.clone();
                            let line_hl = self.apply_highlighting(&line_content, &lang);
                            self.clear_and_move_up(line_content);
                            self.output_buffer.push_str(&line_hl);
                        }
                        self.output_buffer
                            .push_str(&format!("{}{}", style::Reset, "\x1b[K"));
                        self.code_block_depth -= 1;
                        self.current_code_lang = self.depth_stack.pop().unwrap_or_default();
                        self.current_code_line.clear();
                        self.code_history.clear();
                        self.buffer.drain(..end_idx);
                        if self.code_block_depth > 0 {
                            self.output_buffer
                                .push_str(&format!("{}\x1b[K", self.get_block_bg()));
                        }
                        self.at_line_start = true;
                        self.line_pos = 0;
                        continue;
                    }
                }

                let c = match self.buffer.chars().next() {
                    Some(val) => val,
                    None => break,
                };
                self.buffer.drain(..c.len_utf8());

                if c == '\n' {
                    let line_content = self.current_code_line.clone();
                    let lang = self.current_code_lang.clone();
                    let line_hl = self.apply_highlighting(&line_content, &lang);
                    self.clear_and_move_up(line_content);
                    self.output_buffer.push_str(&line_hl);
                    self.output_buffer.push_str(&format!("{}\n", style::Reset));
                    self.code_history.push_str(&self.current_code_line);
                    self.code_history.push('\n');
                    self.current_code_line.clear();
                    self.at_line_start = true;
                    self.line_pos = 0;
                    self.output_buffer
                        .push_str(&format!("{}\x1b[K", self.get_block_bg()));
                } else {
                    let line_content_before = self.current_code_line.clone();
                    self.clear_and_move_up(line_content_before);
                    self.current_code_line.push(c);
                    let line_content_after = self.current_code_line.clone();
                    let lang = self.current_code_lang.clone();
                    let new_hl = self.apply_highlighting(&line_content_after, &lang);
                    self.output_buffer.push_str(&new_hl);
                    self.at_line_start = false;
                }
                continue;
            }

            // Handling Standard Markdown
            if self.at_line_start {
                if let Some(mat) = re_open.captures(&self.buffer.clone()) {
                    let lang = mat.get(1).unwrap().as_str().to_lowercase();
                    self.output_buffer
                        .push_str(&format!("{}{}", style::Reset, "\x1b[K"));
                    let old_lang = self.current_code_lang.clone();
                    self.depth_stack.push(old_lang);
                    self.code_block_depth += 1;
                    self.current_code_lang = lang;
                    self.code_history.clear();
                    self.buffer.drain(..mat.get(0).unwrap().end());
                    self.at_line_start = true;
                    self.output_buffer
                        .push_str(&format!("{}\x1b[K", self.get_block_bg()));
                    continue;
                }

                if let Some(mat) = re_hr.find(&self.buffer.clone()) {
                    let full_len = mat.end();
                    self.apply_indentation();
                    let bar = "─".repeat((self.term_width - self.line_pos) as usize);
                    self.output_buffer.push_str(&format!(
                        "{}{}{}{}\n",
                        bg,
                        color::Fg(color::AnsiValue(244)),
                        bar,
                        style::Reset
                    ));
                    self.buffer.drain(..full_len);
                    self.line_pos = 0;
                    self.at_line_start = true;
                    continue;
                }
            }

            if self.buffer.starts_with("**") || self.buffer.starts_with("__") {
                self.toggle_style("bold");
                self.buffer.drain(..2);
                continue;
            }
            if self.buffer.starts_with('*') || self.buffer.starts_with('_') {
                self.toggle_style("italic");
                self.buffer.drain(..1);
                continue;
            }
            if self.buffer.starts_with('`') {
                self.toggle_style("code");
                self.buffer.drain(..1);
                continue;
            }

            let c = match self.buffer.chars().next() {
                Some(val) => val,
                None => break,
            };
            self.buffer.drain(..c.len_utf8());

            match c {
                '\\' => {
                    if let Some(next) = self.buffer.chars().next() {
                        self.buffer.drain(..next.len_utf8());
                        self.word_buffer.push(next);
                    }
                }
                ' ' | '\n' => {
                    self.flush_word(&bg);
                    if c == '\n' {
                        self.output_buffer
                            .push_str(&format!("{}\x1b[K\n{}", bg, style::Reset));
                        self.line_pos = 0;
                        self.at_line_start = true;
                        let block_bg = self.get_block_bg();
                        if !block_bg.is_empty() {
                            self.output_buffer.push_str(&format!("{}\x1b[K", block_bg));
                        }
                        self.active_styles.clear();
                    } else {
                        self.output_buffer.push_str(&format!("{} ", bg));
                        self.line_pos += 1;
                    }
                }
                _ => {
                    self.word_buffer.push(c);
                    self.at_line_start = false;
                }
            }
        }
    }

    /// Receives raw byte data. Accumulates bytes until a valid UTF-8 string is formed.
    pub fn terminal_stream(&mut self, bytes: &[u8]) {
        self.byte_cache.extend_from_slice(bytes);

        match String::from_utf8(self.byte_cache.clone()) {
            Ok(valid_string) => {
                self.buffer.push_str(&valid_string);
                self.byte_cache.clear();
                self.process_buffer(false);
            }
            Err(e) => {
                let valid_up_to = e.utf8_error().valid_up_to();
                if valid_up_to > 0 {
                    let valid_part =
                        String::from_utf8_lossy(&self.byte_cache[..valid_up_to]).into_owned();
                    self.buffer.push_str(&valid_part);
                    self.byte_cache.drain(..valid_up_to);
                    self.process_buffer(false);
                }
            }
        }
    }

    pub fn finish(&mut self) {
        if !self.byte_cache.is_empty() {
            let leftover = String::from_utf8_lossy(&self.byte_cache).into_owned();
            self.buffer.push_str(&leftover);
            self.byte_cache.clear();
        }
        self.process_buffer(true);
        self.output_buffer
            .push_str(&format!("{}{}", style::Reset, "\x1b[K"));
    }
}

/// PyO3 Wrapper for exposing the streamer to Python
#[cfg_attr(feature = "python", pyclass)]
pub struct MarkdownStreamer {
    core: MarkdownStreamerCore,
}

#[cfg_attr(feature = "python", pymethods)]
impl MarkdownStreamer {
    #[cfg(feature = "python")]
    #[new]
    pub fn new() -> Self {
        Self {
            core: MarkdownStreamerCore::new(),
        }
    }
    pub fn read_available(&mut self) -> String {
        self.core.read_available()
    }
    pub fn terminal_stream(&mut self, data: &[u8]) {
        self.core.terminal_stream(data);
    }
    pub fn finish(&mut self) {
        self.core.finish();
    }
}

#[cfg(feature = "python")]
#[pymodule]
fn my_rust(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<MarkdownStreamer>()?;
    Ok(())
}