wiki-maker/src/main.rs

use axum::{
    Router,
    extract::{Path, State},
    http::StatusCode,
    response::{Html, IntoResponse, Response},
    routing::get,
};
use clap::{Parser, Subcommand};
use lazy_static::lazy_static;
use pulldown_cmark::{Options, Parser as MarkdownParser, html};
use std::sync::Arc;
use std::{io::Cursor, path::PathBuf};
use syntect::{highlighting::ThemeSet, parsing::SyntaxSet};
use tera::{Context, Tera};

mod codeblocks;
use codeblocks::*;

lazy_static! {
    pub static ref TEMPLATES: Tera = {
        let mut tera = Tera::default();
        tera.add_raw_templates(vec![
            ("page.html", include_str!("../templates/page.html")),
            ("style.css", include_str!("../templates/style.css")),
        ])
        .unwrap();
        tera
    };
    pub static ref SYNTAX_SET: SyntaxSet = SyntaxSet::load_defaults_newlines();
    pub static ref THEME_SET: ThemeSet = {
        let mut set = ThemeSet::load_defaults();

        let theme_bytes = include_bytes!("../themes/Catppuccin-Macchiato.tmTheme");

        let mut cursor = Cursor::new(theme_bytes);
        match syntect::highlighting::ThemeSet::load_from_reader(&mut cursor) {
            Ok(theme) => {
                set.themes.insert("Catppuccin Macchiato".to_string(), theme);
            }
            Err(e) => {
                tracing::error!("Failed to load embedded theme: {}", e);
            }
        }
        set
    };
}

#[derive(Parser)]
#[command(author, version, about = "A simple markdown book server")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Serve the markdown files in a directory
    Serve {
        /// Path to the directory containing SUMMARY.md
        path: PathBuf,
        /// Port to listen on
        #[arg(short, long, default_value = "3456")]
        port: u16,
    },
}

struct AppState {
    docs_dir: PathBuf,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    lazy_static::initialize(&TEMPLATES);

    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    let cli = Cli::parse();

    match cli.command {
        Commands::Serve { path, port } => {
            let abs_path = std::fs::canonicalize(&path)?;

            let shared_state = Arc::new(AppState { docs_dir: abs_path });

            let app = Router::new()
                .route("/", get(render_summary))
                .route("/{page}", get(render_page))
                .route("/style.css", get(serve_css))
                .with_state(shared_state);

            let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", port)).await?;
            tracing::info!("Listening on http://localhost:{}", port);
            axum::serve(listener, app).await?;
        }
    }

    Ok(())
}

async fn render_summary(State(state): State<Arc<AppState>>) -> impl IntoResponse {
    render_md_file("SUMMARY.md", state).await
}

async fn render_page(
    State(state): State<Arc<AppState>>,
    Path(page): Path<String>,
) -> impl IntoResponse {
    let filename = if page.ends_with(".md") {
        page
    } else {
        format!("{}.md", page)
    };
    render_md_file(&filename, state).await
}

async fn serve_css() -> impl IntoResponse {
    match TEMPLATES.render("style.css", &tera::Context::new()) {
        Ok(css) => Response::builder()
            .header("content-type", "text/css")
            .body(css.into())
            .unwrap(),
        Err(_) => (StatusCode::NOT_FOUND, "CSS not found").into_response(),
    }
}

async fn render_md_file(filename: &str, state: Arc<AppState>) -> Html<String> {
    let file_path = state.docs_dir.join(filename);

    let content = match tokio::fs::read_to_string(&file_path).await {
        Ok(c) => c,
        Err(_) => return Html("<h1>404</h1><p>Page not found</p>".to_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 renderer = CodeblockRenderer::new(parser);

    let mut html_output = String::new();
    html::push_html(&mut html_output, renderer);

    let (prev_page, next_page) = get_nav_links(&state.docs_dir, filename);

    let mut context = Context::new();
    context.insert("title", filename);
    context.insert("content", &html_output);
    context.insert("prev_page", &prev_page);
    context.insert("next_page", &next_page);

    match TEMPLATES.render("page.html", &context) {
        Ok(rendered) => Html(rendered),
        Err(e) => Html(format!("<h1>Template Error</h1><pre>{}</pre>", e)),
    }
}

fn get_nav_links(dir: &PathBuf, current_file: &str) -> (Option<String>, Option<String>) {
    let mut files: Vec<String> = std::fs::read_dir(dir)
        .unwrap()
        .filter_map(|entry| {
            let path = entry.ok()?.path();
            if path.extension()? == "md" && path.file_name()? != "SUMMARY.md" {
                Some(path.file_name()?.to_str()?.to_string())
            } else {
                None
            }
        })
        .collect();

    files.sort();

    if current_file == "SUMMARY.md" {
        return (None, files.first().cloned());
    }

    let pos = files.iter().position(|f| f == current_file);

    match pos {
        Some(i) => {
            let prev = if i == 0 {
                // If first page, point back to summary page
                Some(".".to_string())
            } else {
                // Otherwise point to the previous file in the list
                files.get(i - 1).cloned()
            };

            let next = files.get(i + 1).cloned();
            (prev, next)
        }

        None => (None, None),
    }
}