mdpreview/src/main.rs

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

use syntect::easy::HighlightLines;
use syntect::highlighting::ThemeSet;
use syntect::html::{styled_line_to_highlighted_html, IncludeBackground};
use syntect::parsing::SyntaxSet;

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

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

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();

    println!("Загрузка баз синтаксисов и тем...");
    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),
    };

    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);

    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    println!("Сервер запущен на http://{}", addr);
    println!("Пример: http://127.0.0.1:3000/files/readme.md");

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

async fn root() -> Result<Html<String>, StatusCode> {
    render_directory_index(&PathBuf::from("./notes"), "").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("./notes");
    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("./notes").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 {
                if 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;
    }

    println!("Watcher запущен для директории: {:?}", watch_path);

    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 = 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;")
}