mdpreview/src/main.rs

use axum::{
    Router,
    extract::{Path, State},
    http::{HeaderMap, StatusCode, header},
    response::{Html, IntoResponse, Sse},
    routing::get,
};
use futures::StreamExt;
use notify::{Config, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
use pulldown_cmark::{CodeBlockKind, Event, Options, Tag, html};
use std::convert::Infallible;
use std::net::{SocketAddr, ToSocketAddrs};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::fs;
use tokio::sync::broadcast;
use tokio_stream::wrappers::BroadcastStream;

use syntect::easy::HighlightLines;
use syntect::highlighting::ThemeSet;
use syntect::html::{IncludeBackground, styled_line_to_highlighted_html};
use syntect::parsing::SyntaxSet;
use tower_http::trace::TraceLayer;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

use clap::Parser;
use std::io;

// Загрузка шаблонов при компиляции
const TEMPLATE_FILE: &str = include_str!("../templates/file.html");
const TEMPLATE_DIR: &str = include_str!("../templates/dir.html");

#[derive(clap::Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Host to listen
    #[arg(long, default_value_t = String::from("localhost"))]
    host: String,

    /// Port to listen
    #[arg(short, long, default_value_t = 8000)]
    port: u16,

    /// Markdown documents directory root
    #[arg(short, long)]
    root: PathBuf,
}

#[derive(Clone)]
struct AppState {
    syntax_set: Arc<SyntaxSet>,
    theme_set: Arc<ThemeSet>,
    tx: Arc<broadcast::Sender<String>>,
    root: PathBuf,
}

#[tokio::main]
async fn main() {
    let args = Args::parse();

    if !args.root.is_dir() {
        eprintln!("Root {root} is not a directory", root = args.root.display());
        std::process::exit(1);
    }

    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new(
            std::env::var("RUST_LOG").unwrap_or_else(|_| "info,tower_http=info".into()),
        ))
        .with(tracing_subscriber::fmt::layer())
        .init();

    let ss = SyntaxSet::load_defaults_newlines();
    let ts = ThemeSet::load_defaults();

    let (tx, _rx) = broadcast::channel::<String>(100);

    let state = AppState {
        syntax_set: Arc::new(ss),
        theme_set: Arc::new(ts),
        tx: Arc::new(tx),
        root: args.root,
    };

    let watcher_state = state.clone();
    tokio::spawn(async move {
        run_file_watcher(watcher_state).await;
    });

    let app = Router::new()
        .route("/", get(root))
        .route("/*path", get(serve_file))
        .route("/events/*path", get(sse_handler))
        .with_state(state)
        .layer(TraceLayer::new_for_http());

    let addr = resolve_addr(&args.host, args.port).unwrap();

    println!("Сервер запущен на http://{}", addr);

    let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

fn resolve_addr(host: &str, port: u16) -> io::Result<SocketAddr> {
    let addr_str = format!("{}:{}", host, port);

    addr_str
        .to_socket_addrs()?
        .next()
        .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "Не удалось разрешить адрес"))
}

async fn root(State(state): State<AppState>) -> Result<Html<String>, StatusCode> {
    render_directory_index(&state.root, "").await
}

async fn sse_handler(
    State(state): State<AppState>,
    Path(full_path): Path<String>,
) -> impl IntoResponse {
    let mut headers = HeaderMap::new();
    headers.insert(
        header::CONTENT_TYPE,
        header::HeaderValue::from_static("text/event-stream"),
    );
    headers.insert(
        header::CACHE_CONTROL,
        header::HeaderValue::from_static("no-cache"),
    );
    headers.insert(
        header::CONNECTION,
        header::HeaderValue::from_static("keep-alive"),
    );

    let rx = state.tx.subscribe();
    let requested_path = full_path.clone();

    let stream = BroadcastStream::new(rx).filter_map(move |res| {
        let req_path = requested_path.clone();
        async move {
            match res {
                Ok(changed_path) => {
                    if changed_path.contains(&req_path) {
                        Some(Ok::<axum::response::sse::Event, Infallible>(
                            axum::response::sse::Event::default()
                                .event("reload")
                                .data(""),
                        ))
                    } else {
                        None
                    }
                }
                Err(_) => None,
            }
        }
    });

    let sse = Sse::new(stream).keep_alive(
        axum::response::sse::KeepAlive::new()
            .interval(Duration::from_secs(15))
            .text("ping"),
    );

    (headers, sse)
}

async fn serve_file(
    State(state): State<AppState>,
    Path(full_path): Path<String>,
) -> Result<Html<String>, StatusCode> {
    if full_path.is_empty() {
        return Err(StatusCode::NOT_FOUND);
    }

    let mut requested_path = PathBuf::from(&state.root);
    requested_path.push(&full_path);

    // Безопасность путей
    let safe_path = match fs::canonicalize(&requested_path).await {
        Ok(p) => p,
        Err(_) => return Err(StatusCode::NOT_FOUND),
    };

    let base_dir = match fs::canonicalize(&state.root).await {
        Ok(p) => p,
        Err(_) => return Err(StatusCode::INTERNAL_SERVER_ERROR),
    };

    if !safe_path.starts_with(&base_dir) {
        eprintln!("Попытка выхода за пределы директории: {:?}", safe_path);
        return Err(StatusCode::FORBIDDEN);
    }

    let metadata = match fs::metadata(&safe_path).await {
        Ok(m) => m,
        Err(e) => {
            eprintln!("Ошибка получения метаданных: {}", e);
            return Err(StatusCode::INTERNAL_SERVER_ERROR);
        }
    };

    if metadata.is_dir() {
        return render_directory_index(&safe_path, &full_path).await;
    }

    let content = match fs::read_to_string(&safe_path).await {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Ошибка чтения: {}", e);
            return Err(StatusCode::INTERNAL_SERVER_ERROR);
        }
    };

    // Логика кнопки "Назад"
    let back_link = if let Some(pos) = full_path.rfind('/') {
        let parent = &full_path[..pos];
        if parent.is_empty() {
            "/".to_string()
        } else {
            format!("/{}", parent)
        }
    } else {
        "/".to_string()
    };

    let back_button_html = format!(
        r#"<div style="margin-bottom: 20px;">
            <a href="{}" style="display: inline-flex; align-items: center; color: #90a4ae; text-decoration: none; font-size: 0.95em; transition: color 0.2s;"
               onmouseover="this.style.color='#ffffff'" onmouseout="this.style.color='#90a4ae'">
                <span style="margin-right: 8px; font-size: 1.2em;">←</span> Назад
            </a>
        </div>"#,
        back_link
    );

    let html_content = markdown_to_html(&content, &state.syntax_set, &state.theme_set, &full_path);

    // Заполнение шаблона
    let final_html = TEMPLATE_FILE
        .replace("{{CONTENT}}", &html_content)
        .replace("{{SSE_URL}}", &format!("/events/{}", full_path))
        .replace("{{BACK_BUTTON}}", &back_button_html);

    Ok(Html(final_html))
}

async fn render_directory_index(
    dir_path: &PathBuf,
    request_path: &str,
) -> Result<Html<String>, StatusCode> {
    let mut entries = match fs::read_dir(dir_path).await {
        Ok(list) => list,
        Err(e) => {
            eprintln!("Ошибка чтения директории: {}", e);
            return Err(StatusCode::FORBIDDEN);
        }
    };

    let mut files: Vec<(String, String, bool)> = Vec::new();

    while let Some(entry) = entries
        .next_entry()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
    {
        let file_name = entry.file_name().to_string_lossy().to_string();

        if file_name.starts_with('.') {
            continue;
        }

        let is_dir = entry.metadata().await.map(|m| m.is_dir()).unwrap_or(false);

        let mut link_path = request_path.to_string();
        if !link_path.ends_with('/') {
            link_path.push('/');
        }
        link_path.push_str(&file_name);

        files.push((file_name, link_path, is_dir));
    }

    files.sort_by(|a, b| match (a.2, b.2) {
        (true, false) => std::cmp::Ordering::Less,
        (false, true) => std::cmp::Ordering::Greater,
        _ => a.0.cmp(&b.0),
    });

    let mut list_html = String::from("<ul>");

    if !request_path.is_empty() && request_path != "files" {
        let parent_path = request_path.rsplit_once('/').map(|(p, _)| p).unwrap_or("");
        let parent_link = if parent_path.is_empty() {
            "/".to_string()
        } else {
            format!("/{}", parent_path)
        };

        list_html.push_str(&format!(
            r#"<li>
                <a href="{}" class="back-link">📁 ..</a>
            </li>"#,
            parent_link
        ));
    }

    for (name, link, is_dir) in files {
        let icon = if is_dir { "📁" } else { "📄" };
        list_html.push_str(&format!(
            r#"<li>
                <a href="/{}" class="file-link">
                    <span class="icon">{}</span>
                    <span>{}</span>
                </a>
            </li>"#,
            link.trim_start_matches('/'),
            icon,
            name
        ));
    }
    list_html.push_str("</ul>");

    let title_path = if request_path.is_empty() {
        ""
    } else {
        request_path.trim_start_matches('/')
    };

    let final_html = TEMPLATE_DIR
        .replace("{{TITLE_PATH}}", title_path)
        .replace("{{FILE_LIST}}", &list_html);

    Ok(Html(final_html))
}

async fn run_file_watcher(state: AppState) {
    let (tx_fs, mut rx_fs) = tokio::sync::mpsc::channel::<PathBuf>(100);

    let tx_fs_clone = tx_fs.clone();
    let mut watcher = RecommendedWatcher::new(
        move |res: Result<notify::Event, notify::Error>| {
            if let Ok(event) = res
                && matches!(event.kind, EventKind::Modify(_))
            {
                for path in event.paths {
                    let _ = tx_fs_clone.blocking_send(path);
                }
            }
        },
        Config::default(),
    )
    .expect("Failed to create watcher");

    let watch_path = PathBuf::from("./notes");
    if let Err(e) = watcher.watch(&watch_path, RecursiveMode::Recursive) {
        eprintln!("Ошибка настройки watcher: {}", e);
        return;
    }

    while let Some(path) = rx_fs.recv().await {
        if let Some(path_str) = path.to_str() {
            let _ = state.tx.send(path_str.to_string());
        }
    }
}

fn markdown_to_html(markdown: &str, ss: &SyntaxSet, ts: &ThemeSet, _file_path: &str) -> String {
    let theme = &ts.themes["base16-ocean.dark"];

    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);
    options.insert(Options::ENABLE_SMART_PUNCTUATION);

    let parser = pulldown_cmark::Parser::new_ext(markdown, options);

    let mut processed_events: Vec<Event> = Vec::new();
    let mut in_code_block = false;
    let mut current_lang: Option<String> = None;
    let mut current_code = String::new();

    for event in parser {
        match event {
            Event::Start(Tag::CodeBlock(kind)) => {
                in_code_block = true;
                current_code.clear();
                current_lang = if let CodeBlockKind::Fenced(l) = kind {
                    Some(l.to_string())
                } else {
                    None
                };
            }
            Event::End(Tag::CodeBlock(_)) => {
                in_code_block = false;
                let is_mermaid = current_lang.as_deref() == Some("mermaid");

                if is_mermaid {
                    let escaped_code = escape_html(&current_code);
                    let mermaid_html = format!(
                        r#"<div class="code-block-wrapper mermaid-wrapper">
                            <div class="code-header">
                                <span class="code-lang">Mermaid Diagram</span>
                                <button class="copy-btn" onclick="copyCode(this)">Copy</button>
                            </div>
                            <div class="mermaid" style="background: transparent; padding: 20px; text-align: center;">{}</div>
                        </div>"#,
                        escaped_code
                    );
                    processed_events.push(Event::Html(mermaid_html.into()));
                } else {
                    let lang_display = current_lang.as_deref().unwrap_or("text");
                    let lang_escaped = escape_html(lang_display);

                    let highlighted_html = if let Some(lang) = &current_lang {
                        if let Some(syntax) = ss.find_syntax_by_token(lang) {
                            let mut h = HighlightLines::new(syntax, theme);
                            let mut result_html = String::new();
                            for line in current_code.lines() {
                                let line_with_newline = format!("{}\n", line);
                                match h.highlight_line(&line_with_newline, ss) {
                                    Ok(regions) => {
                                        let html_line = styled_line_to_highlighted_html(
                                            &regions[..],
                                            IncludeBackground::No,
                                        )
                                        .unwrap_or_else(|_| escape_html(&line_with_newline));
                                        result_html.push_str(&html_line);
                                    }
                                    Err(_) => {
                                        result_html.push_str(&escape_html(&line_with_newline))
                                    }
                                }
                            }
                            result_html
                        } else {
                            escape_html(&current_code)
                        }
                    } else {
                        escape_html(&current_code)
                    };

                    let code_container = format!(
                        r#"<div class="code-block-wrapper">
                            <div class="code-header">
                                <span class="code-lang">{}</span>
                                <button class="copy-btn" onclick="copyCode(this)">Copy</button>
                            </div>
                            <pre style="margin: 0; border-radius: 0 0 6px 6px;"><code>{}</code></pre>
                        </div>"#,
                        lang_escaped, highlighted_html
                    );
                    processed_events.push(Event::Html(code_container.into()));
                }
            }
            Event::Text(text) if in_code_block => {
                current_code.push_str(&text);
            }
            _ => {
                if !in_code_block {
                    processed_events.push(event);
                }
            }
        }
    }

    let mut body_html = String::new();
    html::push_html(&mut body_html, processed_events.into_iter());
    body_html
}

fn escape_html(text: &str) -> String {
    text.replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
        .replace('\'', "&#39;")
}