mdpreview/src/markdown.rs

use pulldown_cmark::{
    Alignment, BlockQuoteKind, CodeBlockKind, CowStr, Event, LinkType, Options, Parser, Tag,
    TagEnd, TextMergeStream,
};

use std::collections::HashMap;
use std::fmt::Write;

// ========== Класс состояния таблицы ==========
#[derive(Debug, Clone, PartialEq)]
pub enum AlignmentState {
    Left,
    Center,
    Right,
    None,
}

impl From<Alignment> for AlignmentState {
    fn from(alignment: Alignment) -> Self {
        match alignment {
            Alignment::Left => AlignmentState::Left,
            Alignment::Center => AlignmentState::Center,
            Alignment::Right => AlignmentState::Right,
            Alignment::None => AlignmentState::None,
        }
    }
}

impl AlignmentState {
    pub fn to_css_class(&self) -> &'static str {
        match self {
            AlignmentState::Left => "align-left",
            AlignmentState::Center => "align-center",
            AlignmentState::Right => "align-right",
            AlignmentState::None => "",
        }
    }
}

#[derive(Debug, Clone, Default)]
pub struct TableState {
    pub alignments: Vec<AlignmentState>,
    pub column_count: usize,
    pub is_inside_table: bool,
    pub is_inside_head: bool,
    pub is_inside_body: bool,
    pub current_row: usize,
    pub current_column: usize,
    pub cell_data: HashMap<(usize, usize), String>,
    pub has_header: bool,
}

impl TableState {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn start_table(&mut self, alignments: Vec<Alignment>) {
        self.alignments = alignments.into_iter().map(AlignmentState::from).collect();
        self.column_count = self.alignments.len();
        self.is_inside_table = true;
        self.is_inside_head = false;
        self.is_inside_body = false;
        self.current_row = 0;
        self.current_column = 0;
        self.cell_data.clear();
        self.has_header = false;
    }

    pub fn end_table(&mut self) {
        self.is_inside_table = false;
        self.is_inside_head = false;
        self.is_inside_body = false;
    }

    pub fn start_head(&mut self) {
        if self.is_inside_table {
            self.is_inside_head = true;
            self.is_inside_body = false;
            self.current_row = 0;
            self.current_column = 0;
            self.has_header = true;
        }
    }

    pub fn end_head(&mut self) {
        self.is_inside_head = false;
    }

    pub fn start_row(&mut self) {
        if self.is_inside_table {
            if self.is_inside_body {
                self.current_row += 1;
            }
            self.current_column = 0;
        }
    }

    pub fn end_row(&mut self) {
        if self.current_column != self.column_count && self.is_inside_table && self.column_count > 0
        {
            // Исправляем неполные строки, добавляя недостающие ячейки
            while self.current_column < self.column_count {
                let row = if self.is_inside_head {
                    0
                } else {
                    self.current_row
                };
                let col = self.current_column;
                self.cell_data.insert((row, col), String::new());
                self.current_column += 1;
            }
        }
    }

    pub fn start_cell(&mut self) {
        if self.is_inside_table {
            self.current_column += 1;
        }
    }

    pub fn end_cell(&mut self, data: Option<String>) {
        if let Some(text) = data {
            let row = if self.is_inside_head {
                0
            } else {
                self.current_row
            };
            let col = self.current_column - 1;
            self.cell_data.insert((row, col), text);
        }
    }

    pub fn get_current_alignment(&self) -> &AlignmentState {
        if self.current_column == 0 || self.current_column > self.alignments.len() {
            return &AlignmentState::None;
        }
        let col_index = self.current_column - 1;
        &self.alignments[col_index]
    }
}

// ========== Основной рендерер ==========
pub struct MarkdownRenderer {
    table_state: TableState,
    in_code_block: bool,
    code_content: String,
}

impl MarkdownRenderer {
    pub fn new() -> Self {
        Self {
            table_state: TableState::new(),
            in_code_block: false,
            code_content: String::new(),
        }
    }

    pub fn render_start_tag(&mut self, tag: Tag<'_>) -> String {
        let mut output = String::new();
        match tag {
            Tag::Heading { level, .. } => {
                output.push_str(format!("<{level}>").as_str());
            }
            Tag::Strikethrough => {
                output.push_str("<strike>");
            }
            Tag::List(Some(_)) => output.push_str("<ol>"),
            Tag::List(None) => output.push_str("<ul>"),
            Tag::Item => output.push_str("<li>"),
            Tag::Paragraph => output.push_str("<p>"),
            Tag::CodeBlock(kind) => {
                output.push_str(render_codeblock(kind).as_str());
            }
            Tag::Emphasis => output.push_str("<em>"),
            Tag::Strong => output.push_str("<strong>"),
            Tag::BlockQuote(kind) => output.push_str(render_blockquote(kind).as_str()),
            Tag::Link {
                link_type,
                dest_url,
                title,
                ..
            } => output.push_str(render_link(link_type, &dest_url, &title).as_str()),
            Tag::Image {
                link_type,
                dest_url,
                title,
                ..
            } => output.push_str(render_image(link_type, &dest_url, &title).as_str()),

            Tag::Table(alignments) => {
                self.table_state.start_table(alignments.clone());
                output.push_str("<table>");
            }

            Tag::TableHead => {
                self.table_state.start_head();
                output.push_str("<thead>");
            }

            Tag::TableCell => {
                self.table_state.start_cell();
                let alignment = self.table_state.get_current_alignment();
                let css_class = alignment.to_css_class();

                if css_class.is_empty() {
                    output.push_str("<td>");
                } else {
                    let _ = write!(output, "<td class=\"{css_class}\">");
                }
            }

            Tag::TableRow => {
                self.table_state.start_row();
                output.push_str("<tr>");
            }

            Tag::FootnoteDefinition(fref) => {
                let _ = write!(output, "<sup><a href=\"#{fref}\" id=\"{fref}\">");
            }

            Tag::HtmlBlock
            | Tag::DefinitionList
            | Tag::DefinitionListTitle
            | Tag::DefinitionListDefinition
            | Tag::Superscript
            | Tag::Subscript
            | Tag::MetadataBlock(_) => (),
        }

        output
    }

    pub fn render_end_tag(&mut self, tag: TagEnd) -> String {
        let mut output = String::new();
        match tag {
            TagEnd::Heading(level) => {
                output.push_str(format!("</{level}>").as_str());
            }
            TagEnd::Strikethrough => {
                output.push_str("</strike>");
            }
            TagEnd::List(true) => output.push_str("</ol>"),
            TagEnd::List(false) => output.push_str("</ul>"),
            TagEnd::Item => output.push_str("</li>"),
            TagEnd::Paragraph => output.push_str("</p>"),
            TagEnd::CodeBlock => {
                output.push_str("</code></pre></div>");
            }
            TagEnd::Emphasis => output.push_str("</em>"),
            TagEnd::Strong => output.push_str("</strong>"),
            TagEnd::BlockQuote(_) => output.push_str("</blockquote>"),
            TagEnd::Link => output.push_str("</a>"),
            TagEnd::Image => output.push_str("</img>"),

            TagEnd::Table => {
                self.table_state.end_table();
                output.push_str("</table>");
            }

            TagEnd::TableHead => {
                self.table_state.end_head();
                output.push_str("</thead>");
                if self.table_state.is_inside_table {
                    output.push_str("<tbody>");
                    self.table_state.is_inside_body = true;
                }
            }

            TagEnd::TableCell => {
                self.table_state.end_cell(None);
                output.push_str("</td>");
            }

            TagEnd::TableRow => {
                self.table_state.end_row();
                output.push_str("</tr>");
            }

            TagEnd::FootnoteDefinition => {
                output.push_str("</a></sup>");
            }

            TagEnd::HtmlBlock
            | TagEnd::DefinitionList
            | TagEnd::DefinitionListTitle
            | TagEnd::DefinitionListDefinition
            | TagEnd::Superscript
            | TagEnd::Subscript
            | TagEnd::MetadataBlock(_) => (),
        }

        output
    }
}

// ========== Основная функция ==========
pub fn markdown_to_html(markdown: &str, _file_path: &str) -> String {
    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_HEADING_ATTRIBUTES);
    options.insert(Options::ENABLE_GFM);

    let mut renderer = MarkdownRenderer::new();
    let mut html_output = String::new();

    let iterator = TextMergeStream::new(Parser::new_ext(markdown, options));

    for event in iterator {
        match event {
            Event::Start(tag) => {
                if let Tag::CodeBlock(_) = tag {
                    renderer.in_code_block = true;
                    renderer.code_content.clear();
                }
                html_output.push_str(renderer.render_start_tag(tag).as_str());
            }
            Event::End(tag) => {
                if let TagEnd::CodeBlock = tag {
                    renderer.in_code_block = false;
                    html_output.push_str(escape_html(&renderer.code_content).as_str());
                }
                html_output.push_str(renderer.render_end_tag(tag).as_str());
            }
            Event::Text(text) => {
                if renderer.in_code_block {
                    renderer.code_content.push_str(&text);
                } else {
                    html_output.push_str(&escape_html(&text));
                }
            }
            Event::TaskListMarker(true) => {
                html_output.push_str("<input type=\"checkbox\" checked/>");
            }
            Event::TaskListMarker(false) => {
                html_output.push_str("<input type=\"checkbox\"/>");
            }
            Event::Code(code) => {
                html_output.push_str(format!("<code>{}</code>", escape_html(&code)).as_str());
            }
            Event::SoftBreak => html_output.push(' '),
            Event::HardBreak => html_output.push_str("<br>"),

            Event::Html(html) => html_output.push_str(&html),

            Event::Rule => html_output.push_str("<hr>"),

            Event::FootnoteReference(fref) => {
                let _ = write!(
                    html_output,
                    "<sup><a href=\"#{fref}\" id=\"{fref}\">{fref}</a></sup>"
                );
            }

            Event::InlineHtml(_) | Event::InlineMath(_) | Event::DisplayMath(_) => (),
        }
    }

    html_output
}

// ========== Вспомогательные функции ==========
fn render_codeblock(block_kind: CodeBlockKind<'_>) -> String {
    let header = String::from(
        r#"
        <div class="code-block-wrapper">
            <div class="code-header">
                <span class="code-lang">%CODE%</span> <button class="copy-btn" onclick="copyCode(this)">Copy</button>
            </div>
    <pre>"#,
    );

    let mut output = String::new();

    match block_kind {
        CodeBlockKind::Indented => {
            output.push_str(&header.replace("%CODE%", "code"));
            output.push_str("<code>");
        }

        CodeBlockKind::Fenced(CowStr::Borrowed("mermaid")) => {
            output.push_str(&header.replace("%CODE%", "mermaid"));
            output.push_str("<code class=\"language-text mermaid\">");
        }

        CodeBlockKind::Fenced(CowStr::Borrowed("rawmermaid")) => {
            output.push_str(&header.replace("%CODE%", "mermaid"));
            output.push_str("<code class=\"language-text\">");
        }

        CodeBlockKind::Fenced(code_type) => {
            output.push_str(&header.replace("%CODE%", &code_type));
            output.push_str(format!("<code class=\"language-{code_type}\">").as_str());
        }
    }

    output
}

fn render_blockquote(quote_kind: Option<BlockQuoteKind>) -> String {
    let alert_class_and_title = match quote_kind {
        None => return "<blockquote>".to_string(),
        Some(BlockQuoteKind::Note) => ("note", "Note"),
        Some(BlockQuoteKind::Tip) => ("tip", "Tip"),
        Some(BlockQuoteKind::Important) => ("important", "Important"),
        Some(BlockQuoteKind::Warning) => ("warning", "Warning"),
        Some(BlockQuoteKind::Caution) => ("caution", "Caution"),
    };

    format!(
        "<blockquote><div class=\"markdown-alert markdown-alert-{}\"><p class=\"markdown-alert-title\">{}</p>",
        alert_class_and_title.0, alert_class_and_title.1
    )
}

fn render_link(link_type: LinkType, dest_url: &CowStr<'_>, title: &CowStr<'_>) -> String {
    match link_type {
        LinkType::Email => format!("<a href=\"mailto:{dest_url}\" title=\"{title}\">"),
        _ => format!("<a href=\"{dest_url}\" title=\"{title}\">"),
    }
}

fn render_image(_link_type: LinkType, dest_url: &CowStr<'_>, title: &CowStr<'_>) -> String {
    format!("<img src=\"{dest_url}\" alt=\"{title}\">")
}

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

// ========== Тесты ==========
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_table_with_alignments() {
        let markdown = r#"
| Header 1 | Header 2 | Header 3 |
|----------|:--------:|---------:|
| Left     | Center   | Right    |
| Data 1   | Data 2   | Data 3   |
"#;

        let html = markdown_to_html(markdown, "");

        // Проверяем, что HTML содержит правильные классы и атрибуты
        assert!(html.contains("<table>"));
        assert!(html.contains("<thead>"));
        assert!(html.contains("<tbody>"));
        assert!(html.contains("text-left"));
        assert!(html.contains("text-center"));
        assert!(html.contains("text-right"));
    }

    #[test]
    fn test_table_state_tracking() {
        let mut renderer = MarkdownRenderer::new();
        let alignments = vec![Alignment::Left, Alignment::Center, Alignment::Right];

        renderer.render_start_tag(Tag::Table(alignments.clone()));
        assert!(renderer.table_state.is_inside_table);
        assert_eq!(renderer.table_state.column_count, 3);

        renderer.render_start_tag(Tag::TableHead);
        assert!(renderer.table_state.is_inside_head);

        renderer.render_start_tag(Tag::TableRow);
        renderer.render_start_tag(Tag::TableCell);
        assert_eq!(
            renderer.table_state.get_current_alignment(),
            &AlignmentState::Left
        );

        renderer.render_end_tag(TagEnd::TableCell);
        renderer.render_start_tag(Tag::TableCell);
        assert_eq!(
            renderer.table_state.get_current_alignment(),
            &AlignmentState::Center
        );

        renderer.render_end_tag(TagEnd::TableCell);
        renderer.render_start_tag(Tag::TableCell);
        assert_eq!(
            renderer.table_state.get_current_alignment(),
            &AlignmentState::Right
        );
        renderer.render_end_tag(TagEnd::TableCell);
        renderer.render_end_tag(TagEnd::TableRow);
        renderer.render_end_tag(TagEnd::TableHead);

        // Проверяем, что тело было автоматически открыто
        renderer.render_start_tag(Tag::TableRow);
        renderer.render_start_tag(Tag::TableCell);
        assert_eq!(
            renderer.table_state.get_current_alignment(),
            &AlignmentState::Left
        );
        renderer.render_end_tag(TagEnd::TableCell);
        renderer.render_end_tag(TagEnd::TableRow);
        renderer.render_end_tag(TagEnd::Table);

        assert!(!renderer.table_state.is_inside_table);
    }
}