smart_markdown/

stream_renderer.rs

use crate::parser::parse_document;
use crate::renderer::render_element_with_options;
use crate::ThemeMode;

/// Incrementally renders markdown text chunks as they arrive.
///
/// `StreamRenderer` is designed for streaming LLM responses: as the model
/// generates markdown text chunk by chunk, this renderer produces complete,
/// renderable lines as soon as enough input has been buffered to form a
/// complete markdown element (e.g. a paragraph ended by a blank line, a
/// complete table, a closed fenced code block).
///
/// # Examples
///
/// ```rust
/// use smart_markdown::{StreamRenderer, ThemeMode, is_light_terminal};
///
/// let width = terminal_size::terminal_size()
///     .map(|(w, _)| w.0 as usize)
///     .unwrap_or(80);
/// let theme = if is_light_terminal() { ThemeMode::Light } else { ThemeMode::Dark };
/// let mut sr = StreamRenderer::new(width, theme)
///     .with_ascii_table_borders(true)
///     .with_code_theme("base16-ocean.dark");
///
/// // Feed chunks as they arrive from the LLM
/// for line in sr.push("# Hello\n\n") {
///     println!("{line}");
/// }
/// for line in sr.push("this is **bold** text") {
///     println!("{line}");
/// }
///
/// // Flush anything still buffered at the end
/// for line in sr.flush_remaining() {
///     println!("{line}");
/// }
/// ```
pub struct StreamRenderer {
    buffer: String,
    width: usize,
    theme_mode: ThemeMode,
    code_theme: Option<String>,
    ascii_table_borders: bool,
    rendered_count: usize,
}

impl StreamRenderer {
    /// Create a new stream renderer.
    ///
    /// - `width`: terminal width in columns (e.g. from the `terminal_size` crate).
    /// - `theme_mode`: controls syntax highlighting theme for code blocks.
    pub fn new(width: usize, theme_mode: ThemeMode) -> Self {
        StreamRenderer {
            buffer: String::new(),
            width,
            theme_mode,
            code_theme: None,
            ascii_table_borders: false,
            rendered_count: 0,
        }
    }

    /// Set a custom syntax highlighting theme by name.
    ///
    /// See [`crate::highlight::list_themes`] for available theme names.
    pub fn with_code_theme(mut self, theme: &str) -> Self {
        self.code_theme = Some(theme.to_string());
        self
    }

    /// Use ASCII-only table borders (`+`, `-`, `|`) instead of Unicode
    /// box-drawing characters (`┌`, `─`, `│`, etc.).
    ///
    /// Useful for terminals where Unicode box-drawing renders poorly
    /// (e.g. light-background themes without proper color inversion).
    pub fn with_ascii_table_borders(mut self, ascii: bool) -> Self {
        self.ascii_table_borders = ascii;
        self
    }

    /// Push additional text chunks.
    ///
    /// Returns rendered complete lines as they become available.
    /// Incomplete markdown (partial fenced blocks, tables, paragraphs)
    /// is buffered internally.
    pub fn push(&mut self, text: &str) -> Vec<String> {
        self.buffer.push_str(text);
        self.emit_complete()
    }

    /// Flush any remaining buffered content and return the final lines.
    ///
    /// Call this once at the end of the stream to emit any markdown that
    /// hasn't been completed by a blank line or structural close.
    pub fn flush_remaining(&mut self) -> Vec<String> {
        if self.buffer.trim().is_empty() {
            return Vec::new();
        }
        if !self.buffer.ends_with('\n') {
            self.buffer.push('\n');
        }
        let elements = parse_document(&self.buffer);
        let total = elements.len();
        let new_elements: Vec<_> = elements
            .into_iter()
            .skip(self.rendered_count)
            .collect();
        self.rendered_count = total;

        let mut output: Vec<String> = Vec::new();
        for elem in &new_elements {
            output.extend(render_element_with_options(
                elem,
                self.width,
                self.theme_mode,
                self.code_theme.as_deref(),
                self.ascii_table_borders,
            ));
        }
        self.buffer.clear();
        self.rendered_count = 0;
        output
    }

    fn emit_complete(&mut self) -> Vec<String> {
        let (complete, remaining) = split_at_complete_boundary(&self.buffer);
        if complete.is_empty() {
            return Vec::new();
        }

        let elements = parse_document(&complete);
        let total = elements.len();
        let new_elements: Vec<_> = elements
            .into_iter()
            .skip(self.rendered_count)
            .collect();
        self.rendered_count = total;

        let mut output: Vec<String> = Vec::new();
        for elem in &new_elements {
            output.extend(render_element_with_options(
                elem,
                self.width,
                self.theme_mode,
                self.code_theme.as_deref(),
                self.ascii_table_borders,
            ));
        }

        self.buffer = remaining;
        self.rendered_count = 0;
        output
    }
}

/// Split buffer at the last complete markdown element boundary.
/// Returns (complete_prefix, remainder) where complete_prefix ends at a
/// safe boundary (blank line, end of a fenced block, etc.).
fn split_at_complete_boundary(text: &str) -> (String, String) {
    if text.is_empty() {
        return (String::new(), String::new());
    }

    // Find the last double-newline (blank line) boundary — that's always safe.
    // Paragraphs, headings, blockquotes, horizontal rules all end at blank lines.
    if let Some(pos) = text.rfind("\n\n") {
        return (text[..pos].to_string(), trim_leading_newlines(&text[pos + 2..]));
    }

    // Check for completed fenced code block (``` or ~~~).
    let lines: Vec<&str> = text.lines().collect();
    if lines.len() >= 2 {
        let first = lines[0];
        if (first.starts_with("```") || first.starts_with("~~~")) && first.len() >= 3 {
            let fence = &first[..3];
            for i in 1..lines.len() {
                if lines[i].trim().starts_with(fence) && lines[i].trim().len() >= 3
                    && lines[i].trim().chars().take(3).all(|c| c == fence.chars().next().unwrap())
                {
                    let end_pos = text
                        .char_indices()
                        .nth(text.lines().take(i + 1).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
                        .map(|(idx, _)| idx)
                        .unwrap_or(text.len());
                    return (text[..end_pos].to_string(), trim_leading_newlines(&text[end_pos..]));
                }
            }
            // Fenced block started but not closed — buffer it entirely
            return (String::new(), text.to_string());
        }
    }

    // Check for table. A table has: header line, separator line, then rows.
    // It's complete when a non-table, non-empty line appears after rows.
    if let Some(table_end) = find_complete_table_end(&lines) {
        if table_end < lines.len() {
            let end_pos = text
                .char_indices()
                .nth(text.lines().take(table_end).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
                .map(|(idx, _)| idx)
                .unwrap_or(text.len());
            return (text[..end_pos].to_string(), trim_leading_newlines(&text[end_pos..]));
        }
        // All lines are part of the table — don't split
        return (String::new(), text.to_string());
    }

    // Check for definition list — needs the term line + at least one ": " definition line
    if let Some(def_end) = find_complete_definition_list_end(&lines) {
        let end_pos = text
            .char_indices()
            .nth(text.lines().take(def_end).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
            .map(|(idx, _)| idx)
            .unwrap_or(text.len());
        return (text[..end_pos].to_string(), trim_leading_newlines(&text[end_pos..]));
    }

    // Check for HTML block — starts with a tag like <div>, needs closing </div> or blank line
    if let Some(html_end) = find_complete_html_block_end(&lines) {
        let end_pos = text
            .char_indices()
            .nth(text.lines().take(html_end).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
            .map(|(idx, _)| idx)
            .unwrap_or(text.len());
        return (text[..end_pos].to_string(), trim_leading_newlines(&text[end_pos..]));
    }

    // Check for indented code block — followed by non-indented, non-empty line or blank line
    if let Some(code_end) = find_complete_indented_code_end(&lines) {
        let end_pos = text
            .char_indices()
            .nth(text.lines().take(code_end).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
            .map(|(idx, _)| idx)
            .unwrap_or(text.len());
        return (text[..end_pos].to_string(), trim_leading_newlines(&text[end_pos..]));
    }

    // Check for complete lists (ordered, unordered, task) — a list ends when a non-list line appears
    if let Some(list_end) = find_complete_list_end(&lines) {
        let end_pos = text
            .char_indices()
            .nth(text.lines().take(list_end).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
            .map(|(idx, _)| idx)
            .unwrap_or(text.len());
        return (text[..end_pos].to_string(), trim_leading_newlines(&text[end_pos..]));
    }

    // Check for footnote definitions — they can be multiline (continuation lines indented)
    if let Some(fn_end) = find_complete_footnote_end(&lines) {
        let end_pos = text
            .char_indices()
            .nth(text.lines().take(fn_end).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
            .map(|(idx, _)| idx)
            .unwrap_or(text.len());
        return (text[..end_pos].to_string(), trim_leading_newlines(&text[end_pos..]));
    }

    // Single-line elements: headings, horizontal rules, blockquotes (single line), paragraphs
    // If the last line is a heading or HR, emit everything
    if let Some(last) = lines.last() {
        let trimmed = last.trim();
        if trimmed.starts_with('#') && trimmed.len() > 1 && trimmed.as_bytes().get(1) == Some(&b' ') {
            // ATX heading — complete as a single line. Split off heading + preceding lines.
            if lines.len() > 1 {
                let end_pos = text
                    .char_indices()
                    .nth(text.lines().take(lines.len() - 1).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
                    .map(|(idx, _)| idx)
                    .unwrap_or(text.len());
                return (text[..end_pos].to_string(), text[end_pos..].to_string());
            }
            return (text.to_string(), String::new());
        }
        if trimmed == "---" || trimmed == "***" || trimmed == "___" {
            return (text.to_string(), String::new());
        }
        if trimmed.starts_with('>') {
            // Blockquote: emit everything before the blockquote line
            if lines.len() > 1 {
                let end_pos = text
                    .char_indices()
                    .nth(text.lines().take(lines.len() - 1).map(|l| l.len() + 1).sum::<usize>().saturating_sub(1))
                    .map(|(idx, _)| idx)
                    .unwrap_or(text.len());
                return (text[..end_pos].to_string(), text[end_pos..].to_string());
            }
            return (text.to_string(), String::new());
        }
    }

    // If the text ends with a newline, it's a complete paragraph or set of paragraphs
    if text.ends_with('\n') {
        return (text.to_string(), String::new());
    }

    // Scan backwards from the last \n to find a complete element boundary.
    // If the preceding line looks standalone (heading, HR, blockquote), split there.
    if let Some(last_nl) = text.rfind('\n') {
        let prefix = &text[..last_nl];
        let pre_lines: Vec<&str> = prefix.lines().collect();
        if let Some(pre_last) = pre_lines.last() {
            if is_standalone_line(pre_last) {
                return (text[..last_nl + 1].to_string(), text[last_nl + 1..].to_string());
            }
        }
    }

    // Buffer the text — more may arrive that belongs to the same paragraph
    (String::new(), text.to_string())
}

fn is_standalone_line(line: &str) -> bool {
    let line = line.trim();
    if line.starts_with('#') {
        let level = line.chars().take_while(|&c| c == '#').count();
        return level <= 6 && line.len() > level && line.as_bytes().get(level) == Some(&b' ');
    }
    line == "---" || line == "***" || line == "___" || line.starts_with('>')
}

fn trim_leading_newlines(s: &str) -> String {
    s.trim_start_matches('\n').to_string()
}

fn find_complete_table_end(lines: &[&str]) -> Option<usize> {
    if lines.len() < 2 {
        return None;
    }
    let header = lines[0].trim();
    let sep = lines[1].trim();
    if !header.starts_with('|') || !header.ends_with('|')
        || !sep.starts_with('|') || !sep.ends_with('|')
    {
        return None;
    }
    let is_sep = sep
        .chars()
        .filter(|&c| c != ' ' && c != '|' && c != '-' && c != ':')
        .count()
        == 0;
    if !is_sep {
        return None;
    }
    // Count header columns to know expected column count
    let header_cols = header.split('|').filter(|s| !s.is_empty()).count();
    for i in 2..lines.len() {
        let tmp = lines[i].trim();
        if tmp.is_empty() {
            // blank line ends the table
            return Some(i + 1);
        }
        if !tmp.starts_with('|') || !tmp.ends_with('|') {
            // non-table line ends the table
            return Some(i);
        }
        let cols = tmp.split('|').filter(|s| !s.is_empty()).count();
        if cols != header_cols {
            return Some(i);
        }
    }
    // If we've processed all lines and they're all table rows, the table isn't complete yet
    None
}

fn find_complete_definition_list_end(lines: &[&str]) -> Option<usize> {
    if lines.len() < 2 {
        return None;
    }
    let first = lines[0].trim();
    if first.starts_with('#') || first.starts_with('>') || first.starts_with('|')
        || first.starts_with('-') || first.starts_with('*') || first.starts_with('`')
        || first.is_empty()
    {
        return None;
    }
    if !lines[1].trim().starts_with(": ") {
        return None;
    }
    let mut i = 2;
    while i < lines.len() {
        let tmp = lines[i].trim();
        if tmp.starts_with(": ") {
            i += 1;
        } else if tmp.is_empty() {
            return Some(i + 1);
        } else {
            return Some(i);
        }
    }
    None
}

fn find_complete_html_block_end(lines: &[&str]) -> Option<usize> {
    let first = lines[0].trim();
    if !first.starts_with('<') {
        return None;
    }
    let rest = &first[1..];
    let tag_end = rest.find(|c: char| c == '>' || c.is_whitespace())?;
    let tag = &rest[..tag_end];
    let lower = tag.to_lowercase();
    let valid = matches!(
        lower.as_str(),
        "div" | "pre" | "table" | "script" | "style" | "section"
            | "article" | "nav" | "footer" | "header" | "aside" | "main"
            | "blockquote" | "form" | "fieldset" | "details" | "dialog"
            | "figure" | "figcaption" | "dl" | "ol" | "ul" | "h1" | "h2"
            | "h3" | "h4" | "h5" | "h6"
    );
    if !valid {
        return None;
    }
    let close = format!("</{}>", tag);
    for i in 1..lines.len() {
        if lines[i].to_lowercase().contains(&close) {
            return Some(i + 1);
        }
        if lines[i].trim().is_empty() {
            return Some(i + 1);
        }
    }
    None
}

fn find_complete_indented_code_end(lines: &[&str]) -> Option<usize> {
    let first = lines[0];
    if !first.starts_with("    ") && !(first.starts_with('\t') && first.len() > 1) {
        return None;
    }
    for i in 1..lines.len() {
        let l = lines[i];
        if l.starts_with("    ") || (l.starts_with('\t') && l.len() > 1) {
            continue;
        }
        if l.is_empty() {
            continue;
        }
        return Some(i);
    }
    None
}

fn find_complete_list_end(lines: &[&str]) -> Option<usize> {
    let first = lines[0].trim();
    let is_unordered = first.starts_with("* ") || first.starts_with("- ") || first.starts_with("+ ");
    let is_task = first.starts_with("- [ ] ") || first.starts_with("- [x] ") || first.starts_with("- [X] ")
        || first.starts_with("* [ ] ") || first.starts_with("* [x] ") || first.starts_with("* [X] ");
    let is_ordered = first.find(". ").map_or(false, |pos| first[..pos].parse::<u64>().is_ok());

    if !is_unordered && !is_task && !is_ordered {
        return None;
    }

    for i in 1..lines.len() {
        let tmp = lines[i].trim();
        if tmp.is_empty() {
            return Some(i + 1);
        }

        if is_unordered || is_task {
            let still_list = tmp.starts_with("* ") || tmp.starts_with("- ") || tmp.starts_with("+ ")
                || (is_task && (tmp.starts_with("- [ ] ") || tmp.starts_with("- [x] ") || tmp.starts_with("- [X] ")
                    || tmp.starts_with("* [ ] ") || tmp.starts_with("* [x] ") || tmp.starts_with("* [X] ")));
            if !still_list {
                return Some(i);
            }
        }
        if is_ordered {
            if tmp.find(". ").map_or(true, |pos| tmp[..pos].parse::<u64>().is_err()) {
                return Some(i);
            }
        }
    }
    None
}

fn find_complete_footnote_end(lines: &[&str]) -> Option<usize> {
    let first = lines[0].trim();
    if !first.starts_with("[^") {
        return None;
    }
    let close_br = first.find("]:")?;
    if close_br <= 2 {
        return None;
    }
    for i in 1..lines.len() {
        let tmp = lines[i];
        if tmp.trim().is_empty() {
            // blank line ends footnote
            return Some(i + 1);
        }
        if !tmp.starts_with("    ") {
            return Some(i);
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_split_at_blank_line() {
        let (complete, remaining) = split_at_complete_boundary("hello\n\nworld");
        assert_eq!(complete, "hello");
        assert_eq!(remaining, "world");
    }

    #[test]
    fn test_split_no_boundary() {
        let (complete, remaining) = split_at_complete_boundary("hello world");
        assert_eq!(complete, "");
        assert_eq!(remaining, "hello world");
    }

    #[test]
    fn test_split_trailing_newline() {
        let (complete, remaining) = split_at_complete_boundary("hello\n");
        assert_eq!(complete, "hello\n");
        assert_eq!(remaining, "");
    }

    #[test]
    fn test_split_complete_fenced_block() {
        let input = "```rust\nlet x = 1;\n```\nsome text";
        let (complete, remaining) = split_at_complete_boundary(input);
        assert!(complete.contains("```"));
        assert!(complete.contains("```"));
        assert_eq!(remaining, "some text");
    }

    #[test]
    fn test_split_incomplete_fenced_block() {
        let input = "```rust\nlet x = 1;\nstill writing";
        let (complete, remaining) = split_at_complete_boundary(input);
        assert_eq!(complete, "");
        assert_eq!(remaining, input);
    }

    #[test]
    fn test_split_complete_table() {
        let input = "| a | b |\n|---|---|\n| 1 | 2 |\nnext";
        let (complete, remaining) = split_at_complete_boundary(input);
        assert!(complete.contains("| a"));
        assert!(!complete.ends_with('\n'));
        assert_eq!(remaining, "next");
    }

    #[test]
    fn test_split_complete_heading() {
        let (complete, remaining) = split_at_complete_boundary("### Hello\nmore");
        assert_eq!(complete, "### Hello\n");
        assert_eq!(remaining, "more");
    }

    #[test]
    fn test_stream_renderer_paragraph_then_flush() {
        let mut sr = StreamRenderer::new(80, ThemeMode::Dark);
        let lines = sr.push("Hello world.");
        assert!(lines.is_empty(), "unterminated paragraph should buffer");
        let remaining = sr.flush_remaining();
        assert!(!remaining.is_empty());
    }

    #[test]
    fn test_stream_renderer_incremental() {
        let mut sr = StreamRenderer::new(80, ThemeMode::Dark);
        let lines1 = sr.push("First paragraph.");
        assert!(lines1.is_empty() || lines1.iter().any(|l| l.contains("First")));
        let lines2 = sr.push("\n\nSecond paragraph.");
        assert!(!lines2.is_empty());
        let final_lines = sr.flush_remaining();
        assert!(!final_lines.is_empty() || lines2.iter().any(|l| l.contains("Second")));
    }

    #[test]
    fn test_stream_renderer_fenced_block() {
        let mut sr = StreamRenderer::new(80, ThemeMode::Dark);
        let lines1 = sr.push("```rust\nlet x = 1;\n```\n");
        assert!(!lines1.is_empty());
        let remaining = sr.flush_remaining();
        assert!(remaining.is_empty());
    }

    #[test]
    fn test_stream_renderer_table() {
        let mut sr = StreamRenderer::new(80, ThemeMode::Dark);
        let lines = sr.push("| a | b |\n|---|---|\n| 1 | 2 |\n");
        assert!(!lines.is_empty());
        assert!(lines.iter().any(|l| l.contains('│') || l.contains('+')));
    }

    #[test]
    fn test_stream_renderer_ascii_borders() {
        let mut sr = StreamRenderer::new(80, ThemeMode::Dark).with_ascii_table_borders(true);
        let lines = sr.push("| a | b |\n|---|---|\n| 1 | 2 |\n");
        assert!(lines.iter().any(|l| l.contains('+')));
    }

    #[test]
    fn test_stream_renderer_code_theme() {
        let mut sr = StreamRenderer::new(80, ThemeMode::Dark).with_code_theme("base16-ocean.dark");
        let lines = sr.push("```rust\nlet x = 1;\n```\n");
        assert!(!lines.is_empty());
    }
}