mdpreview/src/markdown.rs

use pulldown_cmark::{CodeBlockKind, Event, Options, Tag, html};
use syntect::easy::HighlightLines;
use syntect::highlighting::ThemeSet;
use syntect::html::{IncludeBackground, styled_line_to_highlighted_html};
use syntect::parsing::SyntaxSet;

pub fn markdown_to_html(markdown: &str, ss: &SyntaxSet, ts: &ThemeSet, _file_path: &str) -> String {
    let theme = &ts.themes["base16-ocean.dark"];

    let mut options = Options::empty();
    options.insert(Options::ENABLE_TABLES);
    options.insert(Options::ENABLE_FOOTNOTES);
    options.insert(Options::ENABLE_STRIKETHROUGH);
    options.insert(Options::ENABLE_TASKLISTS);
    options.insert(Options::ENABLE_SMART_PUNCTUATION);

    let parser = pulldown_cmark::Parser::new_ext(markdown, options);

    let mut processed_events: Vec<Event> = Vec::new();
    let mut in_code_block = false;
    let mut current_lang: Option<String> = None;
    let mut current_code = String::new();

    for event in parser {
        match event {
            Event::Start(Tag::CodeBlock(kind)) => {
                in_code_block = true;
                current_code.clear();
                current_lang = if let CodeBlockKind::Fenced(l) = kind {
                    Some(l.to_string())
                } else {
                    None
                };
            }
            Event::End(Tag::CodeBlock(_)) => {
                in_code_block = false;
                let is_mermaid = current_lang.as_deref() == Some("mermaid");

                if is_mermaid {
                    let escaped_code = escape_html(&current_code);
                    let mermaid_html = format!(
                        r#"<div class="code-block-wrapper mermaid-wrapper">
                            <div class="code-header">
                                <span class="code-lang">Mermaid Diagram</span>
                                <button class="copy-btn" onclick="copyCode(this)">Copy</button>
                            </div>
                            <div class="mermaid" style="background: transparent; padding: 20px; text-align: center;">{escaped_code}</div>
                        </div>"#
                    );
                    processed_events.push(Event::Html(mermaid_html.into()));
                } else {
                    let lang_display = current_lang.as_deref().unwrap_or("text");
                    let lang_escaped = escape_html(lang_display);

                    let highlighted_html = if let Some(lang) = &current_lang {
                        if let Some(syntax) = ss.find_syntax_by_token(lang) {
                            let mut h = HighlightLines::new(syntax, theme);
                            let mut result_html = String::new();
                            for line in current_code.lines() {
                                let line_with_newline = format!("{line}\n");
                                match h.highlight_line(&line_with_newline, ss) {
                                    Ok(regions) => {
                                        let html_line = styled_line_to_highlighted_html(
                                            &regions[..],
                                            IncludeBackground::No,
                                        )
                                        .unwrap_or_else(|_| escape_html(&line_with_newline));
                                        result_html.push_str(&html_line);
                                    }
                                    Err(_) => {
                                        result_html.push_str(&escape_html(&line_with_newline));
                                    }
                                }
                            }
                            result_html
                        } else {
                            escape_html(&current_code)
                        }
                    } else {
                        escape_html(&current_code)
                    };

                    let code_container = format!(
                        r#"<div class="code-block-wrapper">
                            <div class="code-header">
                                <span class="code-lang">{lang_escaped}</span>
                                <button class="copy-btn" onclick="copyCode(this)">Copy</button>
                            </div>
                            <pre style="margin: 0; border-radius: 0 0 6px 6px;"><code>{highlighted_html}</code></pre>
                        </div>"#
                    );
                    processed_events.push(Event::Html(code_container.into()));
                }
            }
            Event::Text(text) if in_code_block => {
                current_code.push_str(&text);
            }
            _ => {
                if !in_code_block {
                    processed_events.push(event);
                }
            }
        }
    }

    let mut body_html = String::new();
    html::push_html(&mut body_html, processed_events.into_iter());
    body_html
}

fn escape_html(text: &str) -> String {
    text.replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
        .replace('\'', "&#39;")
}