tetratto/crates/shared/src/markdown.rs

use ammonia::Builder;
use markdown::{to_html_with_options, Options, CompileOptions, ParseOptions, Constructs};
use std::collections::HashSet;

/// Render markdown input into HTML
pub fn render_markdown(input: &str) -> String {
    let input = &parse_alignment(input);
    let options = Options {
        compile: CompileOptions {
            allow_any_img_src: false,
            allow_dangerous_html: true,
            allow_dangerous_protocol: true,
            gfm_task_list_item_checkable: false,
            gfm_tagfilter: false,
            ..Default::default()
        },
        parse: ParseOptions {
            constructs: Constructs {
                math_flow: true,
                math_text: true,
                ..Constructs::gfm()
            },
            gfm_strikethrough_single_tilde: false,
            math_text_single_dollar: false,
            mdx_expression_parse: None,
            mdx_esm_parse: None,
            ..Default::default()
        },
    };

    let html = match to_html_with_options(input, &options) {
        Ok(h) => h,
        Err(e) => e.to_string(),
    };

    let mut allowed_attributes = HashSet::new();
    allowed_attributes.insert("id");
    allowed_attributes.insert("class");
    allowed_attributes.insert("ref");
    allowed_attributes.insert("aria-label");
    allowed_attributes.insert("lang");
    allowed_attributes.insert("title");
    allowed_attributes.insert("align");
    allowed_attributes.insert("src");

    Builder::default()
        .generic_attributes(allowed_attributes)
        .add_tags(&[
            "video", "source", "img", "b", "span", "p", "i", "strong", "em", "a", "align",
        ])
        .rm_tags(&["script", "style", "link", "canvas"])
        .add_tag_attributes("a", &["href", "target"])
        .add_url_schemes(&["atto"])
        .clean(&html)
        .to_string()
        .replace(
            "src=\"http",
            "loading=\"lazy\" src=\"/api/v1/util/proxy?url=http",
        )
        .replace("<video loading=", "<video controls loading=")
}

fn parse_alignment_line(line: &str, output: &mut String, buffer: &mut String, is_in_pre: bool) {
    if is_in_pre {
        output.push_str(&format!("{line}\n"));
        return;
    }

    let mut is_alignment_waiting: bool = false;
    let mut alignment_center: bool = false;
    let mut has_dash: bool = false;
    let mut escape: bool = false;

    for char in line.chars() {
        if alignment_center && char != '-' {
            // last char was <, but we didn't receive a hyphen directly after
            alignment_center = false;
            buffer.push('<');
        }

        if has_dash && char != '>' {
            // the last char was -, meaning we need to flip has_dash and push the char since we haven't used it
            has_dash = false;
            buffer.push('-');
        }

        match char {
            '\\' => {
                escape = true;
                continue;
            }
            '-' => {
                if escape {
                    buffer.push(char);
                    escape = false;
                    continue;
                }

                if alignment_center && is_alignment_waiting {
                    // this means the previous element was <, so we're wrapping up alignment now
                    alignment_center = false;
                    is_alignment_waiting = false;
                    output.push_str(&format!("<align class=\"center\">{buffer}</align>"));
                    buffer.clear();
                    continue;
                }

                has_dash = true;

                if !is_alignment_waiting {
                    // we need to go ahead and push/clear the buffer so we don't capture the stuff that came before this
                    // this only needs to be done on the first of these for a single alignment block
                    output.push_str(&format!("{buffer}\n"));
                    buffer.clear();
                }
            }
            '<' => {
                if escape {
                    buffer.push(char);
                    escape = false;
                    continue;
                }

                alignment_center = true;
                continue;
            }
            '>' => {
                if escape {
                    buffer.push(char);
                    escape = false;
                    continue;
                }

                if has_dash {
                    has_dash = false;

                    // if we're already waiting for aligmment, this means this is the SECOND aligner arrow
                    if is_alignment_waiting {
                        is_alignment_waiting = false;
                        output.push_str(&format!("<align class=\"right\">{buffer}</align>"));
                        buffer.clear();
                        continue;
                    }

                    // we're now waiting for the next aligner
                    is_alignment_waiting = true;
                    continue;
                } else {
                    buffer.push('>');
                }
            }
            _ => buffer.push(char),
        }

        escape = false;
    }

    output.push_str(&format!("{buffer}\n"));
    buffer.clear();
}

pub fn parse_alignment(input: &str) -> String {
    let lines = input.split("\n");

    let mut is_in_pre: bool = false;
    let mut output = String::new();
    let mut buffer = String::new();

    for line in lines {
        match line {
            "```" => {
                is_in_pre = !is_in_pre;
                output.push_str(&format!("{line}\n"));
            }
            _ => parse_alignment_line(line, &mut output, &mut buffer, is_in_pre),
        }
    }

    output.push_str(&buffer);
    output
}