AnyBot/src/diagram.rs

// Mermaid diagrams in chat (```mermaid fences), like the original's own Mermaid-compatible
// engine: pure Rust (mermaid-rs-renderer → SVG), drawn with resvg using AnyBot's fonts so
// Chinese labels show. Flowcharts, sequence, class, state, ER, gantt, pie, mindmap, timeline…

use std::sync::OnceLock;

/// The original refuses diagrams over 50 000 characters.
pub const MAX_SOURCE: usize = 50_000;

fn fontdb() -> std::sync::Arc<resvg::usvg::fontdb::Database> {
    static DB: OnceLock<std::sync::Arc<resvg::usvg::fontdb::Database>> = OnceLock::new();
    DB.get_or_init(|| {
        let mut db = resvg::usvg::fontdb::Database::new();
        db.load_font_data(crate::HUNINN.to_vec());
        db.load_system_fonts();
        // Every family the diagram asks for falls back to 粉圓體 (it has Latin and CJK).
        let found: Option<String> = db.faces().find(|f| f.families.iter().any(|(n, _)| n.to_lowercase().contains("huninn") || n.contains("粉圓"))).and_then(|f| f.families.first().map(|(n, _)| n.clone()));
        if let Some(name) = found {
            db.set_sans_serif_family(name.clone());
            db.set_serif_family(name);
        }
        std::sync::Arc::new(db)
    })
    .clone()
}

/// Mermaid source → SVG text.
pub fn svg(source: &str, dark: bool) -> Result<String, String> {
    if source.len() > MAX_SOURCE {
        return Err("圖表太大".into());
    }
    let mut opts = mermaid_rs_renderer::RenderOptions::default();
    if dark {
        opts.theme = mermaid_rs_renderer::Theme::dark();
    }
    opts.theme.font_family = "sans-serif".into();
    // A panic inside the layout must not take the window down.
    let src = source.to_string();
    std::panic::catch_unwind(move || mermaid_rs_renderer::render_strict(&src, opts).map_err(|e| e.to_string()))
        .map_err(|_| "圖表排版失敗".to_string())?
}

/// Mermaid source → RGBA pixels at `scale` (2.0 for sharp text on Retina).
pub fn render(source: &str, dark: bool, scale: f32) -> Result<image::RgbaImage, String> {
    let svg = svg(source, dark)?;
    let mut opt = resvg::usvg::Options::default();
    opt.fontdb = fontdb();
    opt.font_family = "sans-serif".into();
    let tree = resvg::usvg::Tree::from_str(&svg, &opt).map_err(|e| e.to_string())?;
    let size = tree.size();
    let (w, h) = ((size.width() * scale).ceil().max(1.0) as u32, (size.height() * scale).ceil().max(1.0) as u32);
    if w as u64 * h as u64 > 40_000_000 {
        return Err("圖表太大".into());
    }
    let mut pm = resvg::tiny_skia::Pixmap::new(w, h).ok_or("圖表太大")?;
    resvg::render(&tree, resvg::tiny_skia::Transform::from_scale(scale, scale), &mut pm.as_mut());
    // tiny-skia is premultiplied; egui wants straight alpha.
    let mut px = pm.take();
    for c in px.chunks_exact_mut(4) {
        let a = c[3] as u32;
        if a > 0 && a < 255 {
            for v in &mut c[..3] {
                *v = ((*v as u32 * 255 + a / 2) / a).min(255) as u8;
            }
        }
    }
    image::RgbaImage::from_raw(w, h, px).ok_or_else(|| "圖表轉換失敗".into())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn flowchart_and_errors() {
        let img = render("flowchart TD\n  A[開始] --> B{要不要出門？}\n  B -->|要| C[帶傘]\n  B -->|不要| D[在家看書]\n  C --> E[結束]\n  D --> E", false, 1.0).unwrap();
        assert!(img.width() > 100 && img.height() > 100);
        // Some ink in the picture (not all transparent).
        assert!(img.pixels().filter(|p| p.0[3] > 0).count() > 1000);
        if let Ok(out) = std::env::var("ANYBOT_DIAGRAM_OUT") {
            img.save(out).unwrap();
        }
        assert!(render("this is not mermaid at all {{{", false, 1.0).is_err());
        assert!(render(&"x".repeat(MAX_SOURCE + 1), false, 1.0).is_err());
    }
}