streaming_buffer.rs

//! Streaming markdown buffer for safe incremental rendering.
//!
//! This module provides a buffer that accumulates streaming markdown chunks
//! and determines safe points to flush content for rendering. It tracks
//! open markdown constructs (code blocks, bold, links, etc.) to ensure
//! we only output complete, well-formed markdown.
//!
//! # Example
//!
//! ```
//! use goose_cli::session::streaming_buffer::MarkdownBuffer;
//!
//! let mut buf = MarkdownBuffer::new();
//!
//! // Partial bold - buffers until closed
//! assert_eq!(buf.push("Hello **wor"), Some("Hello ".to_string()));
//! assert_eq!(buf.push("ld**!"), Some("**world**!".to_string()));
//!
//! // At end of stream, flush remaining content
//! let remaining = buf.flush();
//! ```

use regex::Regex;
use std::io::Write;
use std::sync::LazyLock;

const MAX_CODE_BLOCK_LINES: usize = 50;
const TRUNCATED_SHOW_LINES: usize = 20;

fn truncate_code_blocks(content: &str) -> String {
    let (open_pos, fence) = match (content.find("```"), content.find("~~~")) {
        (Some(a), Some(b)) if a <= b => (a, "```"),
        (Some(a), None) => (a, "```"),
        (None, Some(b)) => (b, "~~~"),
        (Some(_), Some(b)) => (b, "~~~"),
        (None, None) => return content.to_string(),
    };

    let Some(after_open) = content.get(open_pos + 3..) else {
        return content.to_string();
    };
    let Some(newline_pos) = after_open.find('\n') else {
        return content.to_string();
    };
    let code_start = open_pos + 3 + newline_pos + 1;

    let Some(code_region) = content.get(code_start..) else {
        return content.to_string();
    };
    let close_pattern = format!("\n{}", fence);
    let Some(close_offset) = code_region.find(&close_pattern) else {
        return content.to_string();
    };

    let Some(code_content) = code_region.get(..close_offset) else {
        return content.to_string();
    };
    let lines: Vec<&str> = code_content.lines().collect();

    if lines.len() <= MAX_CODE_BLOCK_LINES {
        return content.to_string();
    }

    let truncated: String = lines
        .iter()
        .take(TRUNCATED_SHOW_LINES)
        .copied()
        .collect::<Vec<_>>()
        .join("\n");
    let remaining = lines.len() - TRUNCATED_SHOW_LINES;

    let file_msg = save_to_temp_file(code_content)
        .map(|p| format!(" → {}", p))
        .unwrap_or_default();

    let close_pos = code_start + close_offset + 1; // +1 to include the \n
    let prefix = content.get(..code_start).unwrap_or("");
    let suffix = content.get(close_pos..).unwrap_or("");
    format!(
        "{}{}\n... ({} more lines{})\n{}",
        prefix, truncated, remaining, file_msg, suffix
    )
}

fn save_to_temp_file(content: &str) -> Option<String> {
    let mut file = tempfile::Builder::new()
        .prefix("goose-")
        .suffix(".txt")
        .tempfile()
        .ok()?;

    file.write_all(content.as_bytes()).ok()?;

    // Keep the file (don't delete on drop) and get the path
    let (_, path) = file.keep().ok()?;
    Some(path.display().to_string())
}

/// Regex that tokenizes markdown inline elements.
/// Order matters: longer/more-specific patterns first.
static INLINE_TOKEN_RE: LazyLock<Regex> = LazyLock::new(|| {
    Regex::new(concat!(
        r"(",
        r"\\.",                 // Escaped char (highest priority)
        r"|`+",                 // Inline code (variable length backticks)
        r"|\*\*\*",             // Bold+italic
        r"|\*\*",               // Bold
        r"|\*",                 // Italic
        r"|___",                // Bold+italic (underscore)
        r"|__",                 // Bold (underscore)
        r"|_",                  // Italic (underscore)
        r"|~~",                 // Strikethrough
        r"|\!\[",               // Image start
        r"|\]\(",               // Link URL start
        r"|\[",                 // Link text start
        r"|\]",                 // Bracket close (without following paren)
        r"|\)",                 // Link URL end
        r"|[^\\\*_`~\[\]!()]+", // Plain text (no special chars)
        r"|.",                  // Any other single char
        r")"
    ))
    .unwrap()
});

/// A streaming markdown buffer that tracks open constructs.
///
/// Accumulates chunks and returns content that is safe to render,
/// holding back any incomplete markdown constructs. Large code blocks
/// are automatically truncated with full content saved to a temp file.
#[derive(Default)]
pub struct MarkdownBuffer {
    buffer: String,
}

/// Tracks the current parsing state for markdown constructs.
#[derive(Default, Debug, Clone, PartialEq)]
struct ParseState {
    in_code_block: bool,
    code_fence_char: char,
    code_fence_len: usize,
    in_table: bool,
    pending_heading: bool,
    in_inline_code: bool,
    inline_code_len: usize,
    in_bold: bool,
    in_italic: bool,
    in_strikethrough: bool,
    in_link_text: bool,
    in_link_url: bool,
    in_image_alt: bool,
}

impl ParseState {
    /// Returns true if no markdown constructs are currently open.
    fn is_clean(&self) -> bool {
        !self.in_code_block
            && !self.in_table
            && !self.pending_heading
            && !self.in_inline_code
            && !self.in_bold
            && !self.in_italic
            && !self.in_strikethrough
            && !self.in_link_text
            && !self.in_link_url
            && !self.in_image_alt
    }
}

// SAFETY: All string slicing in this impl is safe because:
// - We only slice at positions derived from ASCII characters (newlines, #, |, etc.)
// - The regex tokenizer operates on valid UTF-8 and returns byte positions at char boundaries
// - Code fence detection uses chars().take_while() which respects UTF-8
#[allow(clippy::string_slice)]
impl MarkdownBuffer {
    /// Create a new empty buffer.
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a chunk of markdown text to the buffer.
    ///
    /// Returns any content that is safe to render, or None if the buffer
    /// contains only incomplete constructs. Large code blocks are automatically
    /// truncated with full content saved to a temp file.
    pub fn push(&mut self, chunk: &str) -> Option<String> {
        self.buffer.push_str(chunk);
        let safe_end = self.find_safe_end();

        if safe_end > 0 {
            // SAFETY: safe_end is always at a valid UTF-8 char boundary because:
            // - We only set it after processing complete regex tokens (which match
            //   valid UTF-8 sequences) or at newline positions (ASCII, single byte)
            // - The regex tokenizer operates on &str which guarantees UTF-8
            let to_render = self.buffer[..safe_end].to_string();
            self.buffer = self.buffer[safe_end..].to_string();
            Some(truncate_code_blocks(&to_render))
        } else {
            None
        }
    }

    /// Flush any remaining content from the buffer.
    ///
    /// Call this at the end of a stream to get any buffered content,
    /// even if markdown constructs are unclosed.
    pub fn flush(&mut self) -> String {
        std::mem::take(&mut self.buffer)
    }

    /// Find the last byte position where the parse state is "clean".
    fn find_safe_end(&self) -> usize {
        let mut state = ParseState::default();
        let mut last_safe: usize = 0;
        let bytes = self.buffer.as_bytes();
        let len = bytes.len();
        let mut pos: usize = 0;

        while pos < len {
            let at_line_start = pos == 0 || bytes[pos - 1] == b'\n';

            if at_line_start {
                if let Some(new_pos) = self.process_line_start(&mut state, pos) {
                    pos = new_pos;
                    if state.is_clean() {
                        last_safe = pos;
                    }
                    continue;
                }
            }

            if state.in_code_block {
                while pos < len && bytes[pos] != b'\n' {
                    pos += 1;
                }
                if pos < len {
                    pos += 1;
                }
                continue;
            }

            let remaining = &self.buffer[pos..];
            let line_end = remaining.find('\n').map(|i| pos + i + 1).unwrap_or(len);
            let line_content = &self.buffer[pos..line_end];

            for cap in INLINE_TOKEN_RE.find_iter(line_content) {
                let token = cap.as_str();
                let token_end = pos + cap.end();

                self.process_inline_token(&mut state, token);

                if state.is_clean() {
                    last_safe = token_end;
                }
            }

            if line_end <= len && line_end > pos && bytes[line_end - 1] == b'\n' {
                state.pending_heading = false;
                if state.is_clean() {
                    last_safe = line_end;
                }
            }

            pos = line_end;
        }

        last_safe
    }

    /// Process block-level constructs at the start of a line.
    ///
    /// Returns the new position after processing, or None if no block construct found.
    fn process_line_start(&self, state: &mut ParseState, pos: usize) -> Option<usize> {
        let remaining = &self.buffer[pos..];

        if state.pending_heading {
            state.pending_heading = false;
        }

        if let Some(fence_result) = self.check_code_fence(remaining, state) {
            return Some(pos + fence_result);
        }

        if state.in_code_block {
            return None;
        }

        if remaining.starts_with('#') {
            let hashes = remaining.chars().take_while(|&c| c == '#').count();
            if hashes <= 6 {
                let after_hashes = &remaining[hashes..];
                if after_hashes.is_empty()
                    || after_hashes.starts_with(' ')
                    || after_hashes.starts_with('\n')
                {
                    state.pending_heading = true;
                    return None;
                }
            }
        }

        if remaining.starts_with('|') {
            state.in_table = true;
            return None;
        }

        if (remaining.starts_with('\n') || remaining.is_empty()) && state.in_table {
            state.in_table = false;
            return Some(pos + 1);
        }

        if state.in_table && !remaining.starts_with('|') {
            state.in_table = false;
        }

        None
    }

    /// Check for a code fence and update state accordingly.
    ///
    /// Returns the position after the fence line if found, None otherwise.
    fn check_code_fence(&self, line: &str, state: &mut ParseState) -> Option<usize> {
        let trimmed = line.trim_start();

        let fence_char = trimmed.chars().next()?;
        if fence_char != '`' && fence_char != '~' {
            return None;
        }

        let fence_len = trimmed.chars().take_while(|&c| c == fence_char).count();
        if fence_len < 3 {
            return None;
        }

        let after_fence = &trimmed[fence_len..];

        if state.in_code_block {
            if fence_char == state.code_fence_char
                && fence_len >= state.code_fence_len
                && (after_fence.is_empty()
                    || after_fence.starts_with('\n')
                    || after_fence.trim().is_empty())
            {
                state.in_code_block = false;
                state.code_fence_char = '\0';
                state.code_fence_len = 0;

                if let Some(newline_pos) = line.find('\n') {
                    return Some(newline_pos + 1);
                } else {
                    return Some(line.len());
                }
            }
        } else {
            state.in_code_block = true;
            state.code_fence_char = fence_char;
            state.code_fence_len = fence_len;

            if let Some(newline_pos) = line.find('\n') {
                return Some(newline_pos + 1);
            } else {
                return Some(line.len());
            }
        }

        None
    }

    /// Process an inline token and update state.
    fn process_inline_token(&self, state: &mut ParseState, token: &str) {
        if token.starts_with('\\') && token.len() == 2 {
            return;
        }

        if token.starts_with('`') {
            let tick_count = token.len();
            if state.in_inline_code {
                if tick_count == state.inline_code_len {
                    state.in_inline_code = false;
                    state.inline_code_len = 0;
                }
            } else {
                state.in_inline_code = true;
                state.inline_code_len = tick_count;
            }
            return;
        }

        if state.in_inline_code {
            return;
        }

        match token {
            "***" | "___" => {
                if state.in_bold && state.in_italic {
                    state.in_bold = false;
                    state.in_italic = false;
                } else if state.in_bold {
                    state.in_italic = !state.in_italic;
                } else if state.in_italic {
                    state.in_bold = !state.in_bold;
                } else {
                    state.in_bold = true;
                    state.in_italic = true;
                }
            }
            "**" | "__" => {
                state.in_bold = !state.in_bold;
            }
            "*" | "_" => {
                state.in_italic = !state.in_italic;
            }
            "~~" => {
                state.in_strikethrough = !state.in_strikethrough;
            }
            "![" => {
                state.in_image_alt = true;
            }
            "[" => {
                if !state.in_link_text && !state.in_image_alt {
                    state.in_link_text = true;
                }
            }
            "](" => {
                if state.in_link_text {
                    state.in_link_text = false;
                    state.in_link_url = true;
                } else if state.in_image_alt {
                    state.in_image_alt = false;
                    state.in_link_url = true;
                }
            }
            "]" => {}
            ")" => {
                if state.in_link_url {
                    state.in_link_url = false;
                }
            }
            _ => {}
        }
    }
}

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

    /// Process chunks through the buffer and return all outputs (skipping None, including flush)
    fn stream(chunks: &[&str]) -> Vec<String> {
        let mut buf = MarkdownBuffer::new();
        let mut results: Vec<String> = chunks.iter().filter_map(|chunk| buf.push(chunk)).collect();
        let remaining = buf.flush();
        if !remaining.is_empty() {
            results.push(remaining);
        }
        results
    }

    // ===========================================
    // Realistic LLM streaming scenarios
    // ===========================================

    #[test_case(
        &["I'll", " help", " you", " with", " that", "!"],
        &["I'll", " help", " you", " with", " that", "!"]
        ; "simple sentence streams through immediately without markdown"
    )]
    #[test_case(
        &["Here's the **important", "** part."],
        &["Here's the ", "**important** part."]
        ; "bold split mid-word"
    )]
    #[test_case(
        &["Use the `println!", "` macro."],
        &["Use the ", "`println!` macro."]
        ; "inline code split"
    )]
    #[test_case(
        &["Check [the docs](https://doc", "s.rs) for more."],
        &["Check ", "[the docs](https://docs.rs) for more."]
        ; "link url split"
    )]
    fn test_inline_streaming(chunks: &[&str], expected: &[&str]) {
        assert_eq!(stream(chunks), expected);
    }

    // ===========================================
    // Code blocks (most important for bat rendering)
    // ===========================================

    #[test_case(
        &["```rust\n", "fn main() {\n", "    println!(\"hello\");\n", "}\n", "```\n"],
        &["```rust\nfn main() {\n    println!(\"hello\");\n}\n```\n"]
        ; "rust code block streamed line by line"
    )]
    #[test_case(
        &["Here's an exa", "mple:\n\n```python\nprint(\"``", "`nested```\")\n```\n\nNice!"],
        &["Here's an exa", "mple:\n", "\n```python\nprint(\"```nested```\")\n```\n\nNice!"]
        ; "code block with backticks in string literal"
    )]
    #[test_case(
        &["````md\n", "```\ninner\n```\n", "````\n"],
        &["````md\n```\ninner\n```\n````\n"]
        ; "nested code fence with longer outer fence"
    )]
    #[test_case(
        &["~~~bash\n", "echo 'hello'\n", "~", "~~\n"],
        &["~~~bash\necho 'hello'\n~~~\n"]
        ; "tilde code fence"
    )]
    #[test_case(
        &["```\ncode"],
        &["```\ncode"]
        ; "unclosed code block flushes at end"
    )]
    fn test_code_blocks(chunks: &[&str], expected: &[&str]) {
        assert_eq!(stream(chunks), expected);
    }

    // ===========================================
    // Headings
    // ===========================================

    #[test_case(
        &["# Getting St", "arted\n\nFirst, install..."],
        &["# Getting Started\n\nFirst, install..."]
        ; "heading split mid-word"
    )]
    #[test_case(
        &["## API Reference\n\n###", " Methods\n\n"],
        &["## API Reference\n\n", "### Methods\n\n"]
        ; "multiple headings in one chunk"
    )]
    fn test_headings(chunks: &[&str], expected: &[&str]) {
        assert_eq!(stream(chunks), expected);
    }

    // ===========================================
    // Tables
    // ===========================================

    #[test_case(
        &["| Name | Value |\n", "|------|-------|\n", "| foo  | 42    |\n", "\nMore text"],
        &["| Name | Value |\n|------|-------|\n| foo  | 42    |\n\nMore text"]
        ; "table streamed row by row"
    )]
    #[test_case(
        &["| A | B |\n|---|---|\n| 1 | 2 |\n\n"],
        &["| A | B |\n|---|---|\n| 1 | 2 |\n\n"]
        ; "table followed by blank line"
    )]
    fn test_tables(chunks: &[&str], expected: &[&str]) {
        assert_eq!(stream(chunks), expected);
    }

    // ===========================================
    // Mixed formatting (realistic assistant responses)
    // ===========================================

    #[test_case(
        &[
            "Here's how to do it:\n\n",
            "1. First, run `cargo", " build`\n",
            "2. Then check the **out", "put**\n\n",
            "```rust\n",
            "fn main() {}\n",
            "```\n"
        ],
        &[
            "Here's how to do it:\n\n",
            "1. First, run ",
            "`cargo build`\n",
            "2. Then check the ",
            "**output**\n\n",
            "```rust\nfn main() {}\n```\n"
        ]
        ; "typical assistant response with list code and formatting"
    )]
    #[test_case(
        &[
            "See the [**Rust Book**](https://doc.rust-l",
            "ang.org/book/) for more info.\n\n",
            "Key points:\n- Use `Result` for errors\n- Prefer `Option` over null"
        ],
        &[
            "See the ",
            "[**Rust Book**](https://doc.rust-lang.org/book/) for more info.\n\n",
            "Key points:\n- Use `Result` for errors\n- Prefer `Option` over null"
        ]
        ; "link with nested bold and list"
    )]
    #[test_case(
        &[
            "![screenshot](./img/sc",
            "reen.png)\n\nAs shown above..."
        ],
        &[
            "![screenshot](./img/screen.png)\n\nAs shown above..."
        ]
        ; "image with split url"
    )]
    fn test_mixed_content(chunks: &[&str], expected: &[&str]) {
        assert_eq!(stream(chunks), expected);
    }

    // ===========================================
    // Edge cases and escapes
    // ===========================================

    #[test_case(
        &["Use \\* for bullet points, not \\`code\\`"],
        &["Use \\* for bullet points, not \\`code\\`"]
        ; "escaped markdown characters"
    )]
    #[test_case(
        &["Price: $100 * 2 = $200"],
        &["Price: $100 ", "* 2 = $200"]
        ; "asterisk in math context treated as italic marker"
    )]
    #[test_case(
        &[""],
        &[] as &[&str]
        ; "empty input produces no output"
    )]
    #[test_case(
        &["Hello 世界! Here's some **太字** text."],
        &["Hello 世界! Here's some **太字** text."]
        ; "unicode content"
    )]
    #[test_case(
        &["**bold *and italic* together**"],
        &["**bold *and italic* together**"]
        ; "nested bold and italic"
    )]
    #[test_case(
        &["***bold italic***"],
        &["***bold italic***"]
        ; "combined bold italic marker"
    )]
    #[test_case(
        &["~~stri", "ke~~ and **bo", "ld**"],
        &["~~strike~~ and ", "**bold**"]
        ; "strikethrough and bold split"
    )]
    fn test_edge_cases(chunks: &[&str], expected: &[&str]) {
        assert_eq!(stream(chunks), expected);
    }

    // ===========================================
    // Incomplete constructs at stream end
    // ===========================================

    #[test_case(
        &["This is **incomplete bold"],
        &["This is ", "**incomplete bold"]
        ; "unclosed bold flushes"
    )]
    #[test_case(
        &["Check [broken link](http://"],
        &["Check ", "[broken link](http://"]
        ; "unclosed link flushes"
    )]
    #[test_case(
        &["Start of `code"],
        &["Start of ", "`code"]
        ; "unclosed inline code flushes"
    )]
    fn test_incomplete_constructs(chunks: &[&str], expected: &[&str]) {
        assert_eq!(stream(chunks), expected);
    }

    #[test]
    fn test_multiple_code_blocks_mixed() {
        // One small, one large
        let small_block = "```python\nprint('hello')\n```\n";
        let large_lines: Vec<String> = (1..=60).map(|i| format!("line {}", i)).collect();
        let large_block = format!("```rust\n{}\n```\n", large_lines.join("\n"));
        let input = format!(
            "Some text\n\n{}\n\nMore text\n\n{}",
            small_block, large_block
        );

        let result = truncate_code_blocks(&input);

        // Small block should be unchanged
        assert!(result.contains("print('hello')"));
        // Large block should be truncated
        assert!(result.contains("... (40 more lines"));
    }
}