thoth_cli/
markdown_renderer.rs

1use std::collections::HashMap;
2
3use anyhow::{anyhow, Result};
4use ratatui::{
5    style::{Color, Modifier, Style},
6    text::{Line, Span, Text},
7};
8use syntect::{
9    easy::HighlightLines,
10    highlighting::{Style as SyntectStyle, ThemeSet},
11    parsing::{SyntaxReference, SyntaxSet},
12};
13
14pub struct MarkdownRenderer {
15    syntax_set: SyntaxSet,
16    theme_set: ThemeSet,
17    theme: String,
18    cache: HashMap<String, Text<'static>>,
19}
20
21impl Default for MarkdownRenderer {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl MarkdownRenderer {
28    pub fn new() -> Self {
29        MarkdownRenderer {
30            syntax_set: SyntaxSet::load_defaults_newlines(),
31            theme_set: ThemeSet::load_defaults(),
32            theme: "base16-mocha.dark".to_string(),
33            cache: HashMap::new(),
34        }
35    }
36
37    pub fn render_markdown(
38        &mut self,
39        markdown: String,
40        title: String,
41        width: usize,
42    ) -> Result<Text<'static>> {
43        if let Some(lines) = self.cache.get(&format!("{}{}", &title, &markdown)) {
44            return Ok(lines.clone());
45        }
46
47        let md_syntax = self.syntax_set.find_syntax_by_extension("md").unwrap();
48        let mut lines = Vec::new();
49        let mut in_code_block = false;
50        let mut code_block_lang = String::new();
51        let mut code_block_content = Vec::new();
52        let theme = &self.theme_set.themes[&self.theme];
53        let mut h = HighlightLines::new(md_syntax, theme);
54
55        const HEADER_COLORS: [Color; 6] = [
56            Color::Red,
57            Color::Green,
58            Color::Yellow,
59            Color::Blue,
60            Color::Magenta,
61            Color::Cyan,
62        ];
63
64        // Check if the entire markdown is JSON
65        if (markdown.trim_start().starts_with('{') || markdown.trim_start().starts_with('['))
66            && (markdown.trim_end().ends_with('}') || markdown.trim_end().ends_with(']'))
67        {
68            let json_syntax = self.syntax_set.find_syntax_by_extension("json").unwrap();
69            return Ok(Text::from(self.highlight_code_block(
70                &markdown.lines().map(|x| x.to_string()).collect::<Vec<_>>(),
71                "json",
72                json_syntax,
73                theme,
74                width,
75            )?));
76        }
77
78        let updated_markdown = markdown.clone();
79        let mut markdown_lines = updated_markdown.lines().map(|x| x.to_string()).peekable();
80        while let Some(line) = markdown_lines.next() {
81            if line.starts_with("```") {
82                if in_code_block {
83                    // End of code block
84                    let syntax = self
85                        .syntax_set
86                        .find_syntax_by_token(&code_block_lang)
87                        .unwrap_or(md_syntax);
88                    lines.extend(self.highlight_code_block(
89                        &code_block_content.clone(),
90                        &code_block_lang,
91                        syntax,
92                        theme,
93                        width,
94                    )?);
95                    code_block_content.clear();
96                    in_code_block = false;
97                } else {
98                    // Start of code block
99                    in_code_block = true;
100                    code_block_lang = line.trim_start_matches('`').to_string();
101
102                    // Check if it's a one-line code block
103                    if let Some(next_line) = markdown_lines.peek() {
104                        if next_line.starts_with("```") {
105                            // It's a one-line code block
106                            let syntax = self
107                                .syntax_set
108                                .find_syntax_by_token(&code_block_lang)
109                                .unwrap_or(md_syntax);
110                            lines.extend(self.highlight_code_block(
111                                &["".to_string()],
112                                &code_block_lang,
113                                syntax,
114                                theme,
115                                width,
116                            )?);
117                            in_code_block = false;
118                            markdown_lines.next(); // Skip the closing ```
119                            continue;
120                        }
121                    }
122                }
123            } else if in_code_block {
124                code_block_content.push(line.to_string());
125            } else {
126                let highlighted = h
127                    .highlight_line(&line, &self.syntax_set)
128                    .map_err(|e| anyhow!("Highlight error: {}", e))?;
129                let mut spans: Vec<Span> = highlighted.into_iter().map(into_span).collect();
130
131                // Optimized header handling
132                if let Some(header_level) = line.bytes().position(|b| b != b'#') {
133                    if header_level > 0
134                        && header_level <= 6
135                        && line.as_bytes().get(header_level) == Some(&b' ')
136                    {
137                        let header_color = HEADER_COLORS[header_level.saturating_sub(1)];
138                        spans = vec![Span::styled(
139                            line,
140                            Style::default()
141                                .fg(header_color)
142                                .add_modifier(Modifier::BOLD),
143                        )];
144                    }
145                }
146
147                // Pad regular Markdown lines to full width
148                let line_content: String =
149                    spans.iter().map(|span| span.content.to_string()).collect();
150                let padding_width = width.saturating_sub(line_content.len());
151                if padding_width > 0 {
152                    spans.push(Span::styled(" ".repeat(padding_width), Style::default()));
153                }
154
155                lines.push(Line::from(spans));
156            }
157        }
158
159        let markdown_lines = Text::from(lines);
160        let new_key = &format!("{}{}", &title, &markdown);
161        self.cache.insert(new_key.clone(), markdown_lines.clone());
162        Ok(markdown_lines)
163    }
164
165    fn highlight_code_block(
166        &self,
167        code: &[String],
168        lang: &str,
169        syntax: &SyntaxReference,
170        theme: &syntect::highlighting::Theme,
171        width: usize,
172    ) -> Result<Vec<Line<'static>>> {
173        let mut h = HighlightLines::new(syntax, theme);
174        let mut result = Vec::new();
175
176        let max_line_num = code.len();
177        let line_num_width = max_line_num.to_string().len();
178
179        if lang != "json" {
180            result.push(Line::from(Span::styled(
181                "─".repeat(width),
182                Style::default().fg(Color::White),
183            )));
184        }
185
186        for (line_number, line) in code.iter().enumerate() {
187            let highlighted = h
188                .highlight_line(line, &self.syntax_set)
189                .map_err(|e| anyhow!("Highlight error: {}", e))?;
190
191            let mut spans = if lang == "json" {
192                vec![Span::styled(
193                    format!("{:>width$} ", line_number + 1, width = line_num_width),
194                    Style::default().fg(Color::White),
195                )]
196            } else {
197                vec![Span::styled(
198                    format!("{:>width$} │ ", line_number + 1, width = line_num_width),
199                    Style::default().fg(Color::White),
200                )]
201            };
202            spans.extend(highlighted.into_iter().map(into_span));
203
204            // Pad the line to full width
205            let line_content: String = spans.iter().map(|span| span.content.to_string()).collect();
206            let padding_width = width.saturating_sub(line_content.len());
207            if padding_width > 0 {
208                spans.push(Span::styled(" ".repeat(padding_width), Style::default()));
209            }
210
211            result.push(Line::from(spans));
212        }
213
214        if lang != "json" {
215            result.push(Line::from(Span::styled(
216                "─".repeat(width),
217                Style::default().fg(Color::White),
218            )));
219        }
220
221        Ok(result)
222    }
223}
224
225fn syntect_style_to_ratatui_style(style: SyntectStyle) -> Style {
226    let mut ratatui_style = Style::default().fg(Color::Rgb(
227        style.foreground.r,
228        style.foreground.g,
229        style.foreground.b,
230    ));
231
232    if style
233        .font_style
234        .contains(syntect::highlighting::FontStyle::BOLD)
235    {
236        ratatui_style = ratatui_style.add_modifier(Modifier::BOLD);
237    }
238    if style
239        .font_style
240        .contains(syntect::highlighting::FontStyle::ITALIC)
241    {
242        ratatui_style = ratatui_style.add_modifier(Modifier::ITALIC);
243    }
244    if style
245        .font_style
246        .contains(syntect::highlighting::FontStyle::UNDERLINE)
247    {
248        ratatui_style = ratatui_style.add_modifier(Modifier::UNDERLINED);
249    }
250
251    ratatui_style
252}
253
254fn into_span((style, text): (SyntectStyle, &str)) -> Span<'static> {
255    Span::styled(text.to_string(), syntect_style_to_ratatui_style(style))
256}
257
258#[cfg(test)]
259mod tests {
260    use crate::MIN_TEXTAREA_HEIGHT;
261
262    use super::*;
263
264    #[test]
265    fn test_render_markdown() {
266        let mut renderer = MarkdownRenderer::new();
267        let markdown = "# Header\n\nThis is **bold** and *italic* text.";
268        let rendered = renderer
269            .render_markdown(markdown.to_string(), "".to_string(), 40)
270            .unwrap();
271
272        assert!(rendered.lines.len() >= MIN_TEXTAREA_HEIGHT);
273        assert!(rendered.lines[0]
274            .spans
275            .iter()
276            .any(|span| span.content.contains("Header")));
277        assert!(rendered.lines[2]
278            .spans
279            .iter()
280            .any(|span| span.content.contains("This is")));
281        assert!(rendered.lines[2]
282            .spans
283            .iter()
284            .any(|span| span.content.contains("bold")));
285        assert!(rendered.lines[2]
286            .spans
287            .iter()
288            .any(|span| span.content.contains("italic")));
289    }
290
291    #[test]
292    fn test_render_markdown_with_code_block() {
293        let mut renderer = MarkdownRenderer::new();
294        let markdown = "# Header\n\n```rust\nfn main() {\n    println!(\"Hello, world!\");\n}\n```";
295
296        let rendered = renderer
297            .render_markdown(markdown.to_string(), "".to_string(), 40)
298            .unwrap();
299        assert!(rendered.lines.len() > 5);
300        assert!(rendered.lines[0]
301            .spans
302            .iter()
303            .any(|span| span.content.contains("Header")));
304        assert!(rendered
305            .lines
306            .iter()
307            .any(|line| line.spans.iter().any(|span| span.content.contains("main"))));
308    }
309
310    #[test]
311    fn test_render_json() {
312        let mut renderer = MarkdownRenderer::new();
313        let json = r#"{
314  "name": "John Doe",
315  "age": 30,
316  "city": "New York"
317}"#;
318
319        let rendered = renderer
320            .render_markdown(json.to_string(), "".to_string(), 40)
321            .unwrap();
322
323        assert!(rendered.lines.len() == 5);
324        assert!(rendered.lines[0]
325            .spans
326            .iter()
327            .any(|span| span.content.contains("{")));
328        assert!(rendered.lines[4]
329            .spans
330            .iter()
331            .any(|span| span.content.contains("}")));
332    }
333
334    #[test]
335    fn test_render_markdown_with_one_line_code_block() {
336        let mut renderer = MarkdownRenderer::new();
337        let markdown = "# Header\n\n```rust\n```\n\nText after.".to_string();
338        let rendered = renderer
339            .render_markdown(markdown, "".to_string(), 40)
340            .unwrap();
341
342        assert!(rendered.lines.len() > MIN_TEXTAREA_HEIGHT);
343        assert!(rendered.lines[0]
344            .spans
345            .iter()
346            .any(|span| span.content.contains("Header")));
347        assert!(rendered
348            .lines
349            .iter()
350            .any(|line| line.spans.iter().any(|span| span.content.contains("1 │"))));
351        assert!(rendered
352            .lines
353            .last()
354            .unwrap()
355            .spans
356            .iter()
357            .any(|span| span.content.contains("Text after.")));
358    }
359}