Update main.rs with Mermaid support

This commit is contained in:
2026-03-18 21:32:45 +08:00
parent f14e1b2c60
commit 8ed58fe9de

507
src/main.rs Normal file
View File

@@ -0,0 +1,507 @@
use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use pulldown_cmark::{html, Options, Parser as MarkdownParser};
use std::io::{Read, Write};
use std::net::TcpListener;
use std::path::{Path, PathBuf};
#[derive(Parser)]
#[command(name = "md_reader")]
#[command(about = "A simple Markdown reader", version)]
struct Cli {
#[command(subcommand)]
command: Commands,
}
#[derive(Subcommand)]
enum Commands {
Render {
#[arg(help = "Markdown file to render")]
file: String,
#[arg(short, long, help = "Output file path")]
output: Option<String>,
},
Batch {
#[arg(help = "Directory containing markdown files")]
directory: String,
#[arg(short, long, help = "Output directory")]
output: Option<String>,
},
Server {
#[arg(short, long, default_value = "8080")]
port: u16,
#[arg(short, long, help = "Directory to serve")]
path: Option<String>,
},
Preview {
#[arg(help = "Markdown file to preview")]
file: String,
#[arg(short, long, default_value = "Preview")]
title: Option<String>,
#[arg(short = 'w', long, default_value = "900")]
width: u32,
#[arg(short = 'e', long, default_value = "700")]
height: u32,
},
}
fn render_markdown(content: &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);
let parser = MarkdownParser::new_ext(content, options);
let mut html_output = String::new();
html::push_html(&mut html_output, parser);
let html_output = process_mermaid_blocks(&html_output);
html_output
}
fn process_mermaid_blocks(html: &str) -> String {
let mut result = String::new();
let mut code_content = Vec::new();
let mut skip_until_close = false;
for line in html.lines() {
if line.contains("<pre><code")
&& (line.contains("mermaid") || line.contains("language-mermaid"))
{
skip_until_close = true;
let parts: Vec<&str> = line.split("<pre><code").collect();
if !parts[0].is_empty() {
result.push_str(parts[0]);
}
if let Some(rest) = parts.get(1) {
let after_lang: String = rest.split('>').skip(1).collect::<Vec<_>>().join(">");
if !after_lang.is_empty() && !after_lang.contains("<pre>") {
code_content.push(after_lang.trim_end_matches("</code>").to_string());
}
}
continue;
}
if skip_until_close {
if line.contains("</code>") {
let clean = line
.replace("</code>", "")
.replace("</pre>", "")
.trim_end()
.to_string();
if !clean.is_empty() {
code_content.push(clean);
}
skip_until_close = false;
result.push_str(&format!(
"<div class=\"mermaid\">{}</div>",
code_content.join("\n").trim()
));
code_content.clear();
} else {
code_content.push(line.to_string());
}
continue;
}
result.push_str(line);
result.push('\n');
}
result
}
fn wrap_in_html(title: &str, content: &str) -> String {
format!(
r#"<!DOCTYPE html>
<html lang="zh-TW">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{}</title>
<script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>
<style>
body {{
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
max-width: 800px;
margin: 0 auto;
padding: 20px;
line-height: 1.6;
color: #333;
}}
pre {{
background: #f4f4f4;
padding: 16px;
border-radius: 8px;
overflow-x: auto;
}}
code {{
background: #f4f4f4;
padding: 2px 6px;
border-radius: 4px;
font-family: "SF Mono", Monaco, Consolas, monospace;
}}
pre code {{
background: none;
padding: 0;
}}
blockquote {{
border-left: 4px solid #ddd;
margin: 0;
padding-left: 16px;
color: #666;
}}
table {{
border-collapse: collapse;
width: 100%;
}}
th, td {{
border: 1px solid #ddd;
padding: 8px;
text-align: left;
}}
th {{
background: #f4f4f4;
}}
img {{
max-width: 100%;
}}
.mermaid {{
background: white;
text-align: center;
padding: 20px;
cursor: pointer;
position: relative;
border: 1px solid #ddd;
border-radius: 8px;
margin: 16px 0;
}}
.mermaid:hover {{
border-color: #007bff;
box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}}
.mermaid svg {{
display: block;
margin: 0 auto;
}}
.mermaid-download {{
position: absolute;
top: 8px;
right: 8px;
background: #007bff;
color: white;
border: none;
padding: 4px 12px;
border-radius: 4px;
font-size: 12px;
cursor: pointer;
opacity: 0;
transition: opacity 0.2s;
}}
.mermaid:hover .mermaid-download {{
opacity: 1;
}}
.mermaid-download:hover {{
background: #0056b3;
}}
</style>
</head>
<body>
{}
<script>
async function renderMermaidDiagrams() {{
const mermaidDivs = document.querySelectorAll('.mermaid');
for (let i = 0; i < mermaidDivs.length; i++) {{
const div = mermaidDivs[i];
const code = div.textContent.trim();
const id = 'mermaid-' + i;
try {{
const {{ svg }} = await mermaid.render(id, code);
div.innerHTML = svg + '<button class="mermaid-download" onclick="downloadSvg(this)" data-svg="' + encodeURIComponent(svg) + '">Download SVG</button>';
}} catch (err) {{
console.error('Mermaid render error:', err);
div.innerHTML = '<pre style="text-align:left;color:red;">Error: ' + err.message + '</pre>';
}}
}}
}}
function downloadSvg(btn) {{
const svg = decodeURIComponent(btn.getAttribute('data-svg'));
const blob = new Blob([svg], {{ type: 'image/svg+xml' }});
const url = URL.createObjectURL(blob);
const a = document.createElement('a');
a.href = url;
a.download = 'diagram-' + Date.now() + '.svg';
document.body.appendChild(a);
a.click();
document.body.removeChild(a);
URL.revokeObjectURL(url);
}}
mermaid.initialize({{ startOnLoad: false }});
renderMermaidDiagrams();
</script>
</body>
</html>"#,
title, content
)
}
fn read_file(path: &Path) -> Result<String> {
std::fs::read_to_string(path)
.with_context(|| format!("Failed to read file: {}", path.display()))
}
fn process_file(file_path: &str, output: Option<String>) -> Result<()> {
let path = Path::new(file_path);
let content = read_file(path)?;
let html = render_markdown(&content);
let title = path
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("Markdown");
let full_html = wrap_in_html(title, &html);
let out_path = if let Some(o) = output {
o
} else {
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
let out_dir = home.join("docs/html");
std::fs::create_dir_all(&out_dir)?;
out_dir
.join(format!("{}.html", title))
.to_string_lossy()
.to_string()
};
std::fs::write(&out_path, &full_html)
.with_context(|| format!("Failed to write to {}", out_path))?;
println!("Saved to {}", out_path);
Ok(())
}
fn batch_process(input_dir: &str, output_dir: Option<String>) -> Result<()> {
let out_dir = if let Some(o) = output_dir {
PathBuf::from(o)
} else {
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
home.join("docs/html")
};
std::fs::create_dir_all(&out_dir)?;
let dir_path = Path::new(input_dir);
let mut count = 0;
for entry in std::fs::read_dir(dir_path)? {
let entry = entry?;
let path = entry.path();
if path.extension().and_then(|s| s.to_str()) == Some("md") {
let content = read_file(&path)?;
let html = render_markdown(&content);
let title = path
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("Markdown");
let full_html = wrap_in_html(title, &html);
let out_path = out_dir.join(format!("{}.html", title));
std::fs::write(&out_path, &full_html)?;
println!("Saved: {}", out_path.display());
count += 1;
}
}
println!("\nConverted {} files to {}", count, out_dir.display());
Ok(())
}
fn start_server(port: u16, root_path: Option<String>) -> Result<()> {
let addr = format!("0.0.0.0:{}", port);
let listener = TcpListener::bind(&addr)?;
println!("Server running at http://localhost:{}/", port);
println!("Press Ctrl+C to stop");
let root = root_path.unwrap_or_else(|| ".".to_string());
for stream in listener.incoming() {
let mut stream = stream?;
let mut buffer = [0; 8192];
let bytes_read = stream.read(&mut buffer)?;
let request = String::from_utf8_lossy(&buffer[..bytes_read]);
let lines: Vec<&str> = request.lines().collect();
if let Some(request_line) = lines.first() {
let parts: Vec<&str> = request_line.split_whitespace().collect();
if parts.len() >= 2 {
let path_info = if parts[1] == "/" {
"/index.html"
} else {
parts[1]
};
let file_path = Path::new(&root).join(path_info.trim_start_matches('/'));
let response = handle_request(&file_path);
let _ = stream.write_all(response.as_bytes());
let _ = stream.flush();
}
}
}
Ok(())
}
fn handle_request(path: &Path) -> String {
if path.is_file() {
match read_file(path) {
Ok(content) => {
let ext = path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("text/plain");
if ext == "md" {
let html = render_markdown(&content);
return make_response("200 OK", "text/html", &wrap_in_html("Markdown", &html));
}
let ct = match ext {
"html" => "text/html",
"css" => "text/css",
"js" => "application/javascript",
"json" => "application/json",
"png" => "image/png",
"jpg" | "jpeg" => "image/jpeg",
"svg" => "image/svg+xml",
"gif" => "image/gif",
_ => "text/plain",
};
make_response("200 OK", ct, &content)
}
Err(_) => make_response("404 Not Found", "text/plain", "File not found"),
}
} else if path.is_dir() {
let index_path = path.join("index.html");
if index_path.exists() {
if let Ok(content) = read_file(&index_path) {
return make_response("200 OK", "text/html", &content);
}
}
make_response("200 OK", "text/html", &list_directory(path))
} else {
make_response("404 Not Found", "text/plain", "File not found")
}
}
fn make_response(status: &str, content_type: &str, body: &str) -> String {
format!(
"HTTP/1.1 {}\r\n\
Content-Type: {}\r\n\
Content-Length: {}\r\n\
Connection: close\r\n\
\r\n\
{}",
status,
content_type,
body.len(),
body
)
}
fn list_directory(path: &Path) -> String {
let mut files = Vec::new();
files.push(format!("<h1>Index of {}</h1>", path.display()));
files.push("<ul>".to_string());
if let Ok(entries) = std::fs::read_dir(path) {
let mut entries: Vec<_> = entries.filter_map(|e| e.ok()).collect();
entries.sort_by_key(|a| a.file_name());
for entry in entries {
let name = entry.file_name().to_string_lossy().to_string();
let display_name = if entry.path().is_dir() {
format!("{}/", name)
} else {
name.clone()
};
files.push(format!(
"<li><a href=\"{}\">{}</a></li>",
name, display_name
));
}
}
files.push("</ul>".to_string());
files.join("\n")
}
fn preview_markdown(file_path: &str, title: &str, _width: u32, _height: u32) -> Result<()> {
let path = Path::new(file_path);
let content = read_file(path)?;
let html = render_markdown(&content);
let full_html = wrap_in_html(title, &html);
let temp_dir = std::env::temp_dir();
let safe_title: String = title
.replace(' ', "_")
.chars()
.filter(|c| c.is_alphanumeric() || *c == '_' || *c == '-' || *c == '.')
.collect();
let temp_file = temp_dir.join(format!("{}.html", safe_title));
println!("Temp file: {}", temp_file.display());
std::fs::write(&temp_file, &full_html)
.with_context(|| format!("Failed to write temp file: {}", temp_file.display()))?;
open::that(&temp_file).with_context(|| format!("Failed to open file in browser"))?;
println!("Opening preview in browser... (close browser window to exit)");
Ok(())
}
fn main() {
let cli = Cli::parse();
match cli.command {
Commands::Render { file, output } => {
if let Err(e) = process_file(&file, output) {
eprintln!("Error: {}", e);
std::process::exit(1);
}
}
Commands::Batch { directory, output } => {
if let Err(e) = batch_process(&directory, output) {
eprintln!("Error: {}", e);
std::process::exit(1);
}
}
Commands::Server { port, path } => {
if let Err(e) = start_server(port, path) {
eprintln!("Server error: {}", e);
std::process::exit(1);
}
}
Commands::Preview {
file,
title,
width,
height,
} => {
let window_title = title.unwrap_or_else(|| {
Path::new(&file)
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("Preview")
.to_string()
});
if let Err(e) = preview_markdown(&file, &window_title, width, height) {
eprintln!("Preview error: {}", e);
std::process::exit(1);
}
}
}
}