use super::ast::{Block, MermaidResult};
use futures::FutureExt;
use futures::future::BoxFuture;
const MERMAID_INK_BASE: &str = "https://mermaid.ink/img/";
const PREVALIDATE_TIMEOUT_SECS: u64 = 10;
const ERROR_NOTE_MAX_CHARS: usize = 400;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct MediaEntry {
pub(crate) id: String,
pub(crate) url: String,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct MermaidFence {
pub(crate) start: usize,
pub(crate) end: usize,
pub(crate) source: String,
}
pub(crate) fn base64url(input: &str) -> String {
use base64::Engine as _;
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(input.as_bytes())
}
pub(crate) fn has_mermaid_fence(text: &str) -> bool {
let mut in_fence = false;
for line in text.lines() {
let t = line.trim();
if let Some(rest) = t.strip_prefix("```") {
if in_fence {
in_fence = false; } else {
if rest.trim().eq_ignore_ascii_case("mermaid") {
return true;
}
in_fence = true;
}
}
}
false
}
pub(crate) fn find_mermaid_fences(text: &str) -> Vec<MermaidFence> {
let mut fences = Vec::new();
let mut in_fence = false;
let mut is_mermaid = false;
let mut block_start = 0usize;
let mut source_start = 0usize;
let mut pos = 0usize;
for line in text.split_inclusive('\n') {
let line_end = pos + line.len();
let trimmed = line.trim();
if let Some(rest) = trimmed.strip_prefix("```") {
if in_fence {
if is_mermaid {
fences.push(MermaidFence {
start: block_start,
end: line_end,
source: text[source_start..pos].to_string(),
});
}
in_fence = false;
is_mermaid = false;
} else {
block_start = pos;
source_start = line_end;
is_mermaid = rest.trim().eq_ignore_ascii_case("mermaid");
in_fence = true;
}
}
pos = line_end;
}
fences
}
pub(crate) fn should_render_mermaid(text: &str) -> bool {
let tg = &crate::config::Config::current().channels.telegram;
tg.rich_messages && tg.mermaid_render && has_mermaid_fence(text)
}
pub(crate) async fn prevalidate(source: &str) -> MermaidResult {
let url = format!("{}{}", MERMAID_INK_BASE, base64url(source));
let client = match reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(PREVALIDATE_TIMEOUT_SECS))
.build()
{
Ok(c) => c,
Err(_) => return MermaidResult::Failed("diagram renderer unavailable".into()),
};
let resp = match client.get(&url).send().await {
Ok(r) => r,
Err(e) => {
let note = if e.is_timeout() {
"diagram renderer timed out".to_string()
} else {
"diagram renderer unreachable".to_string()
};
return MermaidResult::Failed(note);
}
};
let status = resp.status().as_u16();
let content_type = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_string();
if is_image_response(status, &content_type) {
return MermaidResult::Image(url);
}
let body = resp.text().await.unwrap_or_default();
MermaidResult::Failed(error_note(status, &body))
}
pub(crate) fn is_image_response(status: u16, content_type: &str) -> bool {
(200..300).contains(&status) && content_type.to_lowercase().starts_with("image/")
}
pub(crate) fn error_note(status: u16, body: &str) -> String {
let trimmed = body.trim();
if trimmed.is_empty() {
return format!("diagram renderer returned HTTP {status}");
}
trimmed.chars().take(ERROR_NOTE_MAX_CHARS).collect()
}
pub(crate) fn replacement_for(
outcome: &MermaidResult,
index: usize,
source: &str,
) -> (String, Option<MediaEntry>) {
match outcome {
MermaidResult::Image(url) => {
let id = format!("diag{index}");
(
format!(""),
Some(MediaEntry {
id,
url: url.clone(),
}),
)
}
MermaidResult::Failed(err) => (markdown_failure_block(err, source), None),
}
}
pub(crate) fn resolve_markdown_media(text: &str) -> BoxFuture<'static, (String, Vec<MediaEntry>)> {
let text = text.to_string();
async move {
let fences = find_mermaid_fences(&text);
if fences.is_empty() {
return (text, Vec::new());
}
let mut result = text.clone();
let mut media = Vec::new();
for (i, fence) in fences.iter().enumerate().rev() {
let outcome = prevalidate(&fence.source).await;
let (replacement, entry) = replacement_for(&outcome, i, &fence.source);
if let Some(e) = entry {
media.push(e);
}
result.replace_range(fence.start..fence.end, &replacement);
}
media.reverse();
(result, media)
}
.boxed()
}
pub(crate) fn resolve_blocks(blocks: Vec<Block>) -> BoxFuture<'static, Vec<Block>> {
async move {
let mut out = Vec::with_capacity(blocks.len());
for block in blocks {
out.push(resolve_block(block).await);
}
out
}
.boxed()
}
fn resolve_block(block: Block) -> BoxFuture<'static, Block> {
async move {
match block {
Block::Code {
lang: Some(lang),
text,
} if is_mermaid_lang(&lang) => {
let result = prevalidate(&text).await;
Block::Mermaid {
source: text,
result,
}
}
Block::Quote(inner) => Block::Quote(resolve_blocks(inner).await),
Block::List(mut list) => {
for item in &mut list.items {
item.children = resolve_blocks(std::mem::take(&mut item.children)).await;
}
Block::List(list)
}
Block::Details {
summary,
blocks,
open,
} => Block::Details {
summary,
blocks: resolve_blocks(blocks).await,
open,
},
other => other,
}
}
.boxed()
}
fn is_mermaid_lang(lang: &str) -> bool {
lang.trim().eq_ignore_ascii_case("mermaid")
}
pub(crate) fn markdown_failure_block(err: &str, source: &str) -> String {
format!(
"> ⚠️ **Mermaid diagram could not be rendered**\n\n```\n{err}\n\nSource:\n{source}\n```"
)
}
pub(crate) fn image_html(url: &str) -> String {
format!("<figure><img src=\"{}\"/></figure>", escape(url))
}
pub(crate) fn failure_html(err: &str, source: &str) -> String {
format!(
"<b>⚠️ Mermaid diagram could not be rendered</b>\n<blockquote>{}</blockquote>\n<pre><code>{}</code></pre>",
escape(err),
escape(source)
)
}
fn escape(t: &str) -> String {
t.replace('&', "&")
.replace('<', "<")
.replace('>', ">")
}