markon_core/

server.rs

use axum::{
    extract::{
        ws::{Message, WebSocket},
        Path as AxumPath, State, WebSocketUpgrade,
    },
    http::{header, StatusCode},
    response::{Html, IntoResponse, Response},
    routing::{delete, get, post},
    Json, Router,
};
use futures_util::{stream::StreamExt, SinkExt};
use qrcode::render::unicode::Dense1x2;
use qrcode::{EcLevel, QrCode};
use rusqlite::Connection;
use serde::{Deserialize, Serialize};
use std::fs;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use tera::Tera;
use tokio::net::TcpListener;
use tokio::sync::broadcast;

use crate::assets::{CssAssets, IconAssets, JsAssets, Templates};
use crate::i18n;
use crate::markdown::MarkdownRenderer;
use crate::search::{SearchQuery, SearchResult};
use crate::workspace::{
    generate_token, ServerLock, WorkspaceConfig, WorkspaceEntry, WorkspaceRegistry,
};

/// Initial workspace for the server (one per CLI path / GUI workspace entry).
pub struct WorkspaceInit {
    pub path: std::path::PathBuf,
    pub enable_search: bool,
    pub enable_viewed: bool,
    pub enable_edit: bool,
    pub shared_annotation: bool,
    /// Path within this workspace to open in the browser (e.g. "notes/file.md").
    pub initial_path: Option<String>,
}

/// Server configuration
pub struct ServerConfig {
    pub host: String,
    pub port: u16,
    pub theme: String,
    pub qr: Option<String>,
    pub open_browser: Option<String>,
    pub shared_annotation: bool,
    /// Random salt for workspace ID generation; None = auto-generate.
    pub salt: Option<String>,
    pub initial_workspaces: Vec<WorkspaceInit>,
    /// Pre-bound listener (GUI mode): server adopts this instead of binding fresh,
    /// eliminating the TOCTOU race between port discovery and actual bind.
    pub bound_listener: Option<std::net::TcpListener>,
    /// Externally-owned registry (GUI mode): share the same registry between
    /// the Tauri commands and the HTTP server so additions are immediately visible.
    pub registry: Option<Arc<WorkspaceRegistry>>,
    /// Management API token. None = auto-generate and write to lock file.
    pub management_token: Option<String>,
    /// UI language override: "zh", "en", or None (auto-detect via sys_locale).
    pub language: Option<String>,
    /// Custom keyboard shortcut overrides (JSON object, injected into browser pages).
    pub shortcuts_json: Option<String>,
    /// Custom CSS variable overrides (rendered as :root { --markon-*: value }).
    pub styles_css: Option<String>,
}

#[derive(Clone)]
pub struct AppState {
    pub theme: Arc<String>,
    pub tera: Arc<Tera>,
    pub shared_annotation: bool,
    pub db: Option<Arc<Mutex<Connection>>>,
    pub tx: Option<broadcast::Sender<String>>,
    pub workspace_registry: Arc<WorkspaceRegistry>,
    pub management_token: Arc<String>,
    /// Pre-built i18n JSON string for injection into templates.
    pub i18n_json: Arc<String>,
    /// Resolved UI language ("zh" or "en").
    pub i18n_lang: Arc<String>,
    /// Keyboard shortcut overrides JSON (empty string if none).
    pub shortcuts_json: Arc<String>,
    /// CSS variable overrides string.
    pub styles_css: Arc<String>,
}

fn detect_lang(override_lang: &Option<String>) -> String {
    match override_lang {
        Some(lang) => i18n::resolve_lang(lang).to_string(),
        None => i18n::resolve_lang("auto").to_string(),
    }
}

fn print_compact_qr(data: &str) -> Result<(), Box<dyn std::error::Error>> {
    // Use low error correction level for smaller QR codes
    let code = QrCode::with_error_correction_level(data.as_bytes(), EcLevel::L)?;

    // Render using Dense1x2 (Unicode half-blocks: ▀▄█) for compact display
    // Dense1x2 uses 2 vertical pixels per character (half-block characters)
    // This naturally compensates for terminal fonts where character height > width
    // Note: The aspect ratio depends on the terminal font - it won't be perfect square
    // on all terminals, but Dense1x2 provides the best balance between size and readability
    let string = code
        .render::<Dense1x2>()
        .quiet_zone(false) // No quiet zone to save space
        .build();

    // Add spacing: 4 spaces on the left, blank line below
    for line in string.lines() {
        println!("    {line}"); // 4 spaces on the left
    }
    println!(); // Blank line below

    Ok(())
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(tag = "type")]
enum WebSocketMessage {
    #[serde(rename = "all_annotations")]
    AllAnnotations { annotations: Vec<serde_json::Value> },
    #[serde(rename = "new_annotation")]
    NewAnnotation { annotation: serde_json::Value },
    #[serde(rename = "delete_annotation")]
    DeleteAnnotation { id: String },
    #[serde(rename = "clear_annotations")]
    ClearAnnotations,
    #[serde(rename = "viewed_state")]
    ViewedState { state: serde_json::Value },
    #[serde(rename = "update_viewed_state")]
    UpdateViewedState { state: serde_json::Value },
}

pub async fn start(config: ServerConfig) -> Result<(), String> {
    let ServerConfig {
        host,
        port,
        theme,
        qr,
        open_browser,
        shared_annotation,
        salt,
        initial_workspaces,
        bound_listener,
        registry,
        management_token,
        language,
        shortcuts_json,
        styles_css,
    } = config;

    // Initialize Tera template engine from embedded resources.
    let mut tera = Tera::default();
    for file_name in Templates::iter() {
        if let Some(file) = Templates::get(&file_name) {
            match std::str::from_utf8(&file.data) {
                Ok(content) => {
                    if let Err(e) = tera.add_raw_template(&file_name, content) {
                        return Err(format!("Failed to add template '{file_name}': {e}"));
                    }
                }
                Err(e) => {
                    return Err(format!("Failed to read template '{file_name}': {e}"));
                }
            }
        }
    }

    let (db, tx) = if shared_annotation {
        let db_path = std::env::var("MARKON_SQLITE_PATH").unwrap_or_else(|_| {
            let home = dirs::home_dir().expect("Cannot find home directory");
            home.join(".markon/annotation.sqlite")
                .to_string_lossy()
                .to_string()
        });
        let parent_dir = std::path::Path::new(&db_path).parent().unwrap();
        if !parent_dir.exists() {
            fs::create_dir_all(parent_dir).expect("Failed to create database directory");
        }
        let conn = Connection::open(&db_path).expect("Failed to open database");
        conn.execute(
            "CREATE TABLE IF NOT EXISTS annotations (
                id TEXT PRIMARY KEY,
                file_path TEXT NOT NULL,
                data TEXT NOT NULL
            )",
            [],
        )
        .expect("Failed to create annotations table");
        conn.execute(
            "CREATE TABLE IF NOT EXISTS viewed_state (
                file_path TEXT PRIMARY KEY,
                state TEXT NOT NULL,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )",
            [],
        )
        .expect("Failed to create viewed_state table");
        let db = Arc::new(Mutex::new(conn));
        let tx = broadcast::channel(100).0;
        (Some(db), Some(tx))
    } else {
        (None, None)
    };

    // Build workspace registry and register initial workspaces.
    let effective_salt = salt.unwrap_or_else(|| format!("markon:{port}"));
    let registry = registry.unwrap_or_else(|| Arc::new(WorkspaceRegistry::new(effective_salt)));

    // Track first workspace's URL path for browser/QR.
    let mut first_workspace_url_path: Option<String> = None;

    for ws_init in initial_workspaces {
        // Expand leading ~ (ASCII) or ～ (full-width) before canonicalize.
        let raw_str = ws_init.path.to_string_lossy().into_owned();
        let normalized = if raw_str.starts_with('～') {
            raw_str.replacen('～', "~", 1)
        } else {
            raw_str
        };
        let raw = std::path::PathBuf::from(&normalized);
        let expanded = if normalized.starts_with("~/") || normalized == "~" {
            if let Some(home) = dirs::home_dir() {
                if normalized == "~" {
                    home
                } else {
                    home.join(&normalized[2..])
                }
            } else {
                raw
            }
        } else {
            raw
        };
        let path = expanded.canonicalize().unwrap_or(expanded);
        let id = registry.add(WorkspaceConfig {
            path,
            enable_search: ws_init.enable_search,
            enable_viewed: ws_init.enable_viewed,
            enable_edit: ws_init.enable_edit,
            shared_annotation: ws_init.shared_annotation,
        });
        if first_workspace_url_path.is_none() {
            let url_path = match ws_init.initial_path {
                Some(ref p) => format!("/{id}/{}", p.trim_start_matches('/')),
                None => format!("/{id}/"),
            };
            first_workspace_url_path = Some(url_path);
        }
    }

    let token = Arc::new(management_token.unwrap_or_else(generate_token));

    let state = AppState {
        theme: Arc::new(theme),
        tera: Arc::new(tera),
        shared_annotation,
        db,
        tx,
        workspace_registry: registry,
        management_token: token.clone(),
        i18n_json: Arc::new(i18n::load_i18n()),
        i18n_lang: Arc::new(detect_lang(&language)),
        shortcuts_json: Arc::new(shortcuts_json.unwrap_or_default()),
        styles_css: Arc::new(styles_css.unwrap_or_default()),
    };

    // Management API: requires loopback source IP + valid token header.
    let mgmt = Router::new()
        .route("/api/workspace", post(add_workspace_handler))
        .route(
            "/api/workspace/{id}",
            delete(remove_workspace_handler).put(update_workspace_handler),
        )
        .route("/api/workspaces", get(list_workspaces_handler))
        .route("/api/save", post(save_file_handler))
        .layer(axum::middleware::from_fn_with_state(
            state.clone(),
            require_local_and_token,
        ));

    let mut app = Router::new()
        // Static assets (literal prefix beats /{workspace_id}/ param)
        .route("/favicon.ico", get(serve_favicon))
        .route("/_/favicon.ico", get(serve_favicon))
        .route("/_/favicon.svg", get(serve_favicon_svg))
        .route("/_/css/{filename}", get(serve_css))
        .route("/_/js/{*path}", get(serve_js))
        .route("/_/ws/{workspace_id}", get(config_ws_handler))
        // Search API (read-only, no auth needed)
        .route("/search", get(search_handler))
        // Workspace content routes
        .route("/{workspace_id}/", get(handle_workspace_root))
        .route("/{workspace_id}/{*path}", get(handle_workspace_path))
        // Everything else → 404
        .fallback(|| async { StatusCode::NOT_FOUND })
        .merge(mgmt);

    if shared_annotation {
        app = app.route("/_/ws", get(ws_handler));
    }

    let app = app.with_state(state);

    let listener = if let Some(std_listener) = bound_listener {
        std_listener
            .set_nonblocking(true)
            .map_err(|e| format!("Failed to set non-blocking: {e}"))?;
        TcpListener::from_std(std_listener)
            .map_err(|e| format!("Failed to convert listener: {e}"))?
    } else {
        let addr = format!("{}:{}", host, port)
            .parse::<SocketAddr>()
            .map_err(|e| format!("Invalid host address '{}': {}", host, e))?;
        TcpListener::bind(&addr)
            .await
            .map_err(|e| format!("Failed to bind to {addr}: {e}"))?
    };
    let addr = listener
        .local_addr()
        .map_err(|e| format!("Failed to get local address: {e}"))?;
    println!("listening on http://{addr}");
    if let Some(ref p) = first_workspace_url_path {
        println!("workspace: http://{addr}{p}");
    }

    // Write lock file so CLI can discover this server.
    let _lock_guard = {
        if let Err(e) = (ServerLock {
            port: addr.port(),
            token: token.as_ref().clone(),
        })
        .write()
        {
            eprintln!("[server] Failed to write lock file: {e}");
        }
        struct LockGuard;
        impl Drop for LockGuard {
            fn drop(&mut self) {
                ServerLock::remove();
            }
        }
        LockGuard
    };

    // Helper: build a full URL from a base option string.
    let make_url = |base_option: &str, ws_path: &Option<String>| -> String {
        let base = if base_option == "local" {
            format!("http://{addr}")
        } else {
            base_option.to_string()
        };
        match ws_path {
            Some(p) => format!("{}{}", base.trim_end_matches('/'), p),
            None => format!("{}/", base.trim_end_matches('/')),
        }
    };

    let custom_base = qr
        .as_ref()
        .filter(|u| u.as_str() != "missing")
        .or_else(|| open_browser.as_ref().filter(|u| u.as_str() != "local"));
    if let Some(base) = custom_base {
        println!(
            "accessible at {}",
            make_url(base, &first_workspace_url_path)
        );
    }

    if let Some(ref qr_option) = qr {
        println!();
        let qr_url = if qr_option == "missing" {
            format!("http://{addr}")
        } else {
            make_url(qr_option, &first_workspace_url_path)
        };
        if let Err(e) = print_compact_qr(&qr_url) {
            eprintln!("Failed to generate QR code: {e}");
        }
    }

    if let Some(ref base_opt) = open_browser {
        let url = make_url(base_opt, &first_workspace_url_path);
        if let Err(e) = open::that(&url) {
            eprintln!("Failed to open browser: {e}");
        }
    }

    axum::serve(
        listener,
        app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
    )
    .await
    .map_err(|e| format!("Server error: {e}"))?;
    Ok(())
}

/// Lightweight always-on WebSocket per workspace — pushes a "reload" text frame
/// whenever workspace flags change. No auth required (read-only notification).
async fn config_ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    AxumPath(workspace_id): AxumPath<String>,
) -> impl IntoResponse {
    let Some(ws_entry) = state.workspace_registry.get(&workspace_id) else {
        return StatusCode::NOT_FOUND.into_response();
    };
    let mut rx = ws_entry.config_tx.subscribe();
    ws.on_upgrade(move |mut socket| async move {
        while let Ok(()) = rx.recv().await {
            if socket
                .send(axum::extract::ws::Message::Text("reload".into()))
                .await
                .is_err()
            {
                break;
            }
        }
    })
}

/// Middleware: management API only accepts loopback source + valid token header.
async fn require_local_and_token(
    State(state): State<AppState>,
    axum::extract::ConnectInfo(addr): axum::extract::ConnectInfo<std::net::SocketAddr>,
    req: axum::extract::Request,
    next: axum::middleware::Next,
) -> Response {
    if !addr.ip().is_loopback() {
        return StatusCode::FORBIDDEN.into_response();
    }
    let ok = req
        .headers()
        .get("X-Markon-Token")
        .and_then(|v| v.to_str().ok())
        .map(|t| t == state.management_token.as_str())
        .unwrap_or(false);
    if !ok {
        return StatusCode::UNAUTHORIZED.into_response();
    }
    next.run(req).await
}

async fn ws_handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
    ws.on_upgrade(|socket| handle_socket(socket, state))
}

async fn handle_socket(socket: WebSocket, state: AppState) {
    let (mut sender, mut receiver) = socket.split();

    // Subscribe to the broadcast channel
    let mut rx = state.tx.as_ref().unwrap().subscribe();

    // Wait for the first message from client to get the file path
    let file_path = match receiver.next().await {
        Some(Ok(Message::Text(text))) => {
            // Expect first message to be file path
            // Convert Utf8Bytes to String
            text.to_string()
        }
        _ => {
            eprintln!("[WebSocket] Failed to receive file path from client");
            return;
        }
    };

    // Send all existing annotations to the new client
    let annotations = {
        let db = state.db.as_ref().unwrap().lock().unwrap();
        let mut stmt = db
            .prepare("SELECT data FROM annotations WHERE file_path = ?1")
            .unwrap();
        let rows = stmt
            .query_map([&file_path.as_str()], |row| row.get::<_, String>(0))
            .unwrap();
        let mut annotations = Vec::new();
        for row in rows {
            let data: serde_json::Value = serde_json::from_str(&row.unwrap()).unwrap();
            annotations.push(data);
        }
        eprintln!(
            "[WebSocket] Sending {} annotations for file_path: {}",
            annotations.len(),
            file_path
        );
        annotations
    };

    let initial_msg = WebSocketMessage::AllAnnotations { annotations };
    if sender
        .send(Message::Text(
            serde_json::to_string(&initial_msg).unwrap().into(),
        ))
        .await
        .is_err()
    {
        return;
    }

    // Send existing viewed state to the new client
    let viewed_state = {
        let db = state.db.as_ref().unwrap().lock().unwrap();
        let state_json = db
            .query_row(
                "SELECT state FROM viewed_state WHERE file_path = ?1",
                [&file_path.as_str()],
                |row| row.get::<_, String>(0),
            )
            .unwrap_or_else(|_| "{}".to_string());
        serde_json::from_str(&state_json).unwrap_or(serde_json::json!({}))
    };

    let viewed_msg = WebSocketMessage::ViewedState {
        state: viewed_state,
    };
    if sender
        .send(Message::Text(
            serde_json::to_string(&viewed_msg).unwrap().into(),
        ))
        .await
        .is_err()
    {
        return;
    }

    // Task to forward broadcast messages to the client
    let mut send_task = tokio::spawn(async move {
        while let Ok(msg) = rx.recv().await {
            if sender.send(Message::Text(msg.into())).await.is_err() {
                break;
            }
        }
    });

    // Task to handle incoming messages from the client
    let mut recv_task = {
        let state = state.clone();
        tokio::spawn(async move {
            while let Some(Ok(Message::Text(text))) = receiver.next().await {
                let msg: WebSocketMessage = match serde_json::from_str(&text) {
                    Ok(msg) => msg,
                    Err(_) => continue,
                };

                let db = state.db.as_ref().unwrap().lock().unwrap();
                match msg {
                    WebSocketMessage::NewAnnotation { annotation } => {
                        let id = annotation["id"].as_str().unwrap().to_string();
                        let data = serde_json::to_string(&annotation).unwrap();
                        db.execute(
                            "INSERT OR REPLACE INTO annotations (id, file_path, data) VALUES (?1, ?2, ?3)",
                            [&id.as_str(), &file_path.as_str(), &data.as_str()],
                        )
                        .unwrap();
                        let broadcast_msg = WebSocketMessage::NewAnnotation { annotation };
                        state
                            .tx
                            .as_ref()
                            .unwrap()
                            .send(serde_json::to_string(&broadcast_msg).unwrap())
                            .unwrap();
                    }
                    WebSocketMessage::DeleteAnnotation { id } => {
                        db.execute(
                            "DELETE FROM annotations WHERE id = ?1 AND file_path = ?2",
                            [&id.as_str(), &file_path.as_str()],
                        )
                        .unwrap();
                        let broadcast_msg = WebSocketMessage::DeleteAnnotation { id };
                        state
                            .tx
                            .as_ref()
                            .unwrap()
                            .send(serde_json::to_string(&broadcast_msg).unwrap())
                            .unwrap();
                    }
                    WebSocketMessage::ClearAnnotations => {
                        eprintln!("[WebSocket] Clearing annotations for file_path: {file_path}");

                        // Clear annotations
                        match db.execute(
                            "DELETE FROM annotations WHERE file_path = ?1",
                            [&file_path.as_str()],
                        ) {
                            Ok(affected_rows) => {
                                eprintln!(
                                    "[WebSocket] Deleted {affected_rows} annotation rows for file_path: {file_path}"
                                );
                            }
                            Err(e) => {
                                eprintln!("[WebSocket] Failed to clear annotations: {e}");
                            }
                        }

                        // Also clear viewed state
                        match db.execute(
                            "DELETE FROM viewed_state WHERE file_path = ?1",
                            [&file_path.as_str()],
                        ) {
                            Ok(affected_rows) => {
                                eprintln!(
                                    "[WebSocket] Deleted {affected_rows} viewed_state rows for file_path: {file_path}"
                                );
                            }
                            Err(e) => {
                                eprintln!("[WebSocket] Failed to clear viewed_state: {e}");
                            }
                        }

                        // Broadcast clear_annotations message
                        let broadcast_msg = WebSocketMessage::ClearAnnotations;
                        state
                            .tx
                            .as_ref()
                            .unwrap()
                            .send(serde_json::to_string(&broadcast_msg).unwrap())
                            .unwrap();

                        // Broadcast empty viewed_state
                        let empty_viewed_state = WebSocketMessage::ViewedState {
                            state: serde_json::Value::Object(serde_json::Map::new()),
                        };
                        state
                            .tx
                            .as_ref()
                            .unwrap()
                            .send(serde_json::to_string(&empty_viewed_state).unwrap())
                            .unwrap();
                    }
                    WebSocketMessage::UpdateViewedState {
                        state: viewed_state,
                    } => {
                        let state_json = serde_json::to_string(&viewed_state).unwrap();
                        db.execute(
                            "INSERT OR REPLACE INTO viewed_state (file_path, state, updated_at) VALUES (?1, ?2, CURRENT_TIMESTAMP)",
                            [&file_path.as_str(), &state_json.as_str()],
                        )
                        .unwrap();

                        // Broadcast to other clients
                        let broadcast_msg = WebSocketMessage::ViewedState {
                            state: viewed_state,
                        };
                        state
                            .tx
                            .as_ref()
                            .unwrap()
                            .send(serde_json::to_string(&broadcast_msg).unwrap())
                            .unwrap();
                    }
                    _ => {}
                }
            }
        })
    };

    // Wait for either task to complete
    tokio::select! {
        _ = (&mut send_task) => recv_task.abort(),
        _ = (&mut recv_task) => send_task.abort(),
    };
}

// ── Workspace content handlers ────────────────────────────────────────────────

async fn handle_workspace_root(
    State(state): State<AppState>,
    AxumPath(workspace_id): AxumPath<String>,
) -> impl IntoResponse {
    let Some(ws) = state.workspace_registry.get(&workspace_id) else {
        return StatusCode::NOT_FOUND.into_response();
    };
    render_directory_listing(&workspace_id, &ws, None, &state)
}

async fn handle_workspace_path(
    State(state): State<AppState>,
    AxumPath((workspace_id, path)): AxumPath<(String, String)>,
) -> impl IntoResponse {
    let Some(ws) = state.workspace_registry.get(&workspace_id) else {
        return StatusCode::NOT_FOUND.into_response();
    };

    let decoded = urlencoding::decode(&path).unwrap_or_else(|_| path.clone().into());
    let full_path = ws.root.join(decoded.trim_start_matches('/'));

    let canonical = match full_path.canonicalize() {
        Ok(p) => p,
        Err(_) => {
            return (StatusCode::NOT_FOUND, format!("Path not found: {decoded}")).into_response()
        }
    };

    if !canonical.starts_with(&ws.root) {
        return (StatusCode::FORBIDDEN, "Access denied").into_response();
    }

    if canonical.is_file() {
        if canonical
            .extension()
            .is_some_and(|e| e.to_string_lossy().to_lowercase() == "md")
        {
            render_markdown_file(&canonical.to_string_lossy(), &workspace_id, &ws, &state)
        } else {
            serve_file(&canonical)
        }
    } else if canonical.is_dir() {
        render_directory_listing(&workspace_id, &ws, Some(&decoded), &state)
    } else {
        (StatusCode::NOT_FOUND, "Path not found").into_response()
    }
}

// ── Workspace management API ──────────────────────────────────────────────────

#[derive(Deserialize)]
struct AddWorkspaceRequest {
    path: String,
    #[serde(default)]
    enable_search: bool,
    #[serde(default)]
    enable_viewed: bool,
    #[serde(default)]
    enable_edit: bool,
    #[serde(default)]
    shared_annotation: bool,
}

#[derive(Serialize)]
struct AddWorkspaceResponse {
    id: String,
}

async fn add_workspace_handler(
    State(state): State<AppState>,
    Json(req): Json<AddWorkspaceRequest>,
) -> impl IntoResponse {
    use std::path::PathBuf;
    let path = PathBuf::from(&req.path);
    let path = match path.canonicalize() {
        Ok(p) => p,
        Err(e) => return (StatusCode::BAD_REQUEST, format!("Invalid path: {e}")).into_response(),
    };
    let id = state.workspace_registry.add(WorkspaceConfig {
        path,
        enable_search: req.enable_search,
        enable_viewed: req.enable_viewed,
        enable_edit: req.enable_edit,
        shared_annotation: req.shared_annotation,
    });
    Json(AddWorkspaceResponse { id }).into_response()
}

async fn remove_workspace_handler(
    State(state): State<AppState>,
    AxumPath(id): AxumPath<String>,
) -> impl IntoResponse {
    if state.workspace_registry.remove(&id) {
        StatusCode::OK
    } else {
        StatusCode::NOT_FOUND
    }
}

#[derive(Deserialize)]
struct UpdateWorkspaceRequest {
    #[serde(default)]
    enable_search: bool,
    #[serde(default)]
    enable_viewed: bool,
    #[serde(default)]
    enable_edit: bool,
    #[serde(default)]
    shared_annotation: bool,
}

async fn update_workspace_handler(
    State(state): State<AppState>,
    AxumPath(id): AxumPath<String>,
    Json(req): Json<UpdateWorkspaceRequest>,
) -> impl IntoResponse {
    if state.workspace_registry.update_flags(
        &id,
        req.enable_search,
        req.enable_viewed,
        req.enable_edit,
        req.shared_annotation,
    ) {
        StatusCode::OK
    } else {
        StatusCode::NOT_FOUND
    }
}

async fn list_workspaces_handler(State(state): State<AppState>) -> impl IntoResponse {
    Json(state.workspace_registry.info_list())
}

// ── Search handler ────────────────────────────────────────────────────────────

#[derive(Deserialize)]
struct WorkspaceSearchQuery {
    ws: String,
    #[serde(flatten)]
    q: SearchQuery,
}

async fn search_handler(
    State(state): State<AppState>,
    axum::extract::Query(query): axum::extract::Query<WorkspaceSearchQuery>,
) -> impl IntoResponse {
    if query.q.q.is_empty() {
        return Json(Vec::<SearchResult>::new());
    }
    let Some(ws) = state.workspace_registry.get(&query.ws) else {
        return Json(Vec::new());
    };
    if !ws.enable_search.load(std::sync::atomic::Ordering::Relaxed) {
        return Json(Vec::new());
    }
    let idx = ws.search_index.lock().unwrap().clone();
    let Some(idx) = idx else {
        return Json(Vec::new()); // still indexing
    };
    let results = idx.search(&query.q.q, 20).unwrap_or_else(|e| {
        eprintln!("[search] error: {e}");
        Vec::new()
    });
    Json(results)
}

fn render_markdown_file(
    file_path: &str,
    workspace_id: &str,
    ws: &WorkspaceEntry,
    state: &AppState,
) -> Response {
    match fs::read_to_string(file_path) {
        Ok(markdown_input) => {
            let renderer = MarkdownRenderer::new(&state.theme);
            let (html_content, has_mermaid, toc) = renderer.render(&markdown_input);

            let title = std::path::Path::new(file_path)
                .file_name()
                .map(|n| n.to_string_lossy().to_string())
                .unwrap_or_else(|| file_path.to_string());

            let mut context = tera::Context::new();
            context.insert("title", &format!("markon - {title}"));
            context.insert("file_path", file_path);
            context.insert("workspace_id", workspace_id);
            context.insert("theme", state.theme.as_str());
            context.insert("content", &html_content);
            // Back link: parent dir of this file within the workspace.
            let back_link = std::path::Path::new(file_path)
                .parent()
                .and_then(|p| p.strip_prefix(&ws.root).ok())
                .map(|rel| {
                    let rel_str = rel.to_string_lossy();
                    if rel_str.is_empty() {
                        format!("/{workspace_id}/")
                    } else {
                        format!("/{workspace_id}/{}/", rel_str)
                    }
                })
                .unwrap_or_else(|| format!("/{workspace_id}/"));
            context.insert("back_link", &back_link);
            context.insert("show_back_link", &true);
            context.insert("has_mermaid", &has_mermaid);
            context.insert("toc", &toc);
            context.insert(
                "shared_annotation",
                &ws.shared_annotation
                    .load(std::sync::atomic::Ordering::Relaxed),
            );
            context.insert("enable_viewed", &ws.enable_viewed);
            context.insert("enable_search", &ws.enable_search);
            context.insert("enable_edit", &ws.enable_edit);

            if ws.enable_edit.load(std::sync::atomic::Ordering::Relaxed) {
                // JSON-encode and HTML-escape so </script> in content can't break the page.
                let json = serde_json::to_string(&markdown_input)
                    .unwrap_or_default()
                    .replace('<', "\\u003c")
                    .replace('>', "\\u003e")
                    .replace('&', "\\u0026");
                context.insert("markdown_content_json", &json);
            }

            context.insert("i18n_json", state.i18n_json.as_str());
            context.insert("i18n_lang", state.i18n_lang.as_str());
            context.insert("shortcuts_json", state.shortcuts_json.as_str());
            context.insert("styles_css", state.styles_css.as_str());

            match state.tera.render("layout.html", &context) {
                Ok(html) => Html(html).into_response(),
                Err(e) => (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Template error: {e}"),
                )
                    .into_response(),
            }
        }
        Err(e) => {
            let mut context = tera::Context::new();
            context.insert("title", "Error");
            context.insert("theme", state.theme.as_str());
            context.insert(
                "content",
                &format!(
                    r#"<p style="color: red;">Error reading file '{file_path}': {e}</p>
                       <a href="/">← Back to file list</a>"#
                ),
            );
            context.insert("show_back_link", &false);
            context.insert("has_mermaid", &false);
            context.insert("i18n_json", state.i18n_json.as_str());
            context.insert("i18n_lang", state.i18n_lang.as_str());
            context.insert("shortcuts_json", state.shortcuts_json.as_str());
            context.insert("styles_css", state.styles_css.as_str());

            match state.tera.render("layout.html", &context) {
                Ok(html) => Html(html).into_response(),
                Err(e) => (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Template error: {e}"),
                )
                    .into_response(),
            }
        }
    }
}

fn render_directory_listing(
    workspace_id: &str,
    ws: &WorkspaceEntry,
    dir_param: Option<&str>,
    state: &AppState,
) -> Response {
    use std::path::PathBuf;

    let current_dir = if let Some(dir_str) = dir_param {
        let p = PathBuf::from(dir_str);
        if p.is_absolute() {
            p
        } else {
            ws.root.join(&p)
        }
    } else {
        ws.root.clone()
    };

    let current_dir = match current_dir.canonicalize() {
        Ok(p) => p,
        Err(e) => {
            return (StatusCode::BAD_REQUEST, format!("Invalid directory: {e}")).into_response()
        }
    };

    #[derive(serde::Serialize)]
    struct Entry {
        name: String,
        is_dir: bool,
        link: String,
    }

    let mut entries: Vec<Entry> = match fs::read_dir(&current_dir) {
        Ok(dir_entries) => dir_entries
            .filter_map(|e| e.ok())
            .filter_map(|entry| {
                let path = entry.path();
                let name = entry.file_name().to_string_lossy().to_string();
                if name.starts_with('.') {
                    return None;
                }
                // Use file_type() — avoids stat() syscall that can block on AutoFS mount points.
                let file_type = match entry.file_type() {
                    Ok(ft) => ft,
                    Err(_) => return None,
                };
                let is_dir = file_type.is_dir();
                let rel = path
                    .strip_prefix(&ws.root)
                    .unwrap_or(&path)
                    .to_string_lossy()
                    .to_string();
                if is_dir {
                    Some(Entry {
                        name,
                        is_dir: true,
                        link: format!("/{workspace_id}/{rel}/"),
                    })
                } else {
                    let is_md = path
                        .extension()
                        .is_some_and(|e| e.to_string_lossy().to_lowercase() == "md");
                    if is_md {
                        Some(Entry {
                            name,
                            is_dir: false,
                            link: format!("/{workspace_id}/{rel}"),
                        })
                    } else {
                        None
                    }
                }
            })
            .collect(),
        Err(e) => {
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Error reading directory: {e}"),
            )
                .into_response()
        }
    };

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

    let show_parent = current_dir != ws.root;
    let parent_link: Option<String> = if show_parent {
        current_dir.parent().map(|parent| {
            let rel = parent
                .strip_prefix(&ws.root)
                .map(|p| p.to_string_lossy().to_string())
                .unwrap_or_default();
            if rel.is_empty() {
                format!("/{workspace_id}/")
            } else {
                format!("/{workspace_id}/{rel}/")
            }
        })
    } else {
        None
    };

    let mut context = tera::Context::new();
    context.insert("theme", state.theme.as_str());
    context.insert("workspace_id", workspace_id);
    context.insert("current_dir", &current_dir.display().to_string());
    context.insert("entries", &entries);
    context.insert("show_parent", &show_parent);
    context.insert("parent_link", &parent_link);
    context.insert("enable_search", &ws.enable_search);
    context.insert("i18n_json", state.i18n_json.as_str());
    context.insert("i18n_lang", state.i18n_lang.as_str());
    context.insert("shortcuts_json", state.shortcuts_json.as_str());
    context.insert("styles_css", state.styles_css.as_str());

    match state.tera.render("directory.html", &context) {
        Ok(html) => Html(html).into_response(),
        Err(e) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Template error: {e}"),
        )
            .into_response(),
    }
}

async fn serve_favicon() -> impl IntoResponse {
    // Redirect /_/favicon.ico to /_/favicon.svg
    (
        StatusCode::MOVED_PERMANENTLY,
        [(header::LOCATION, "/_/favicon.svg")],
    )
        .into_response()
}

async fn serve_favicon_svg() -> impl IntoResponse {
    serve_static_file("favicon.svg", IconAssets::get, "image/svg+xml")
}

async fn serve_css(AxumPath(filename): AxumPath<String>) -> impl IntoResponse {
    serve_static_file(&filename, CssAssets::get, "text/css")
}

async fn serve_js(AxumPath(path): AxumPath<String>) -> impl IntoResponse {
    serve_static_file(&path, JsAssets::get, "application/javascript")
}

fn serve_static_file<F>(filename: &str, getter: F, content_type: &str) -> Response
where
    F: FnOnce(&str) -> Option<rust_embed::EmbeddedFile>,
{
    match getter(filename) {
        Some(file) => (
            StatusCode::OK,
            [(header::CONTENT_TYPE, content_type)],
            file.data.into_owned(),
        )
            .into_response(),
        None => (StatusCode::NOT_FOUND, "File not found").into_response(),
    }
}

fn serve_file(path: &std::path::Path) -> Response {
    match fs::read(path) {
        Ok(content) => {
            // Detect MIME type based on file extension
            let mime_type = path
                .extension()
                .and_then(|ext| ext.to_str())
                .and_then(get_mime_type)
                .unwrap_or("application/octet-stream");

            (StatusCode::OK, [(header::CONTENT_TYPE, mime_type)], content).into_response()
        }
        Err(e) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Error reading file: {e}"),
        )
            .into_response(),
    }
}

fn get_mime_type(ext: &str) -> Option<&'static str> {
    match ext.to_lowercase().as_str() {
        // Images
        "png" => Some("image/png"),
        "jpg" | "jpeg" => Some("image/jpeg"),
        "gif" => Some("image/gif"),
        "webp" => Some("image/webp"),
        "svg" => Some("image/svg+xml"),
        "ico" => Some("image/x-icon"),
        "bmp" => Some("image/bmp"),

        // Audio
        "mp3" => Some("audio/mpeg"),
        "wav" => Some("audio/wav"),
        "ogg" => Some("audio/ogg"),
        "m4a" => Some("audio/mp4"),
        "flac" => Some("audio/flac"),

        // Video
        "mp4" => Some("video/mp4"),
        "webm" => Some("video/webm"),
        "ogv" => Some("video/ogg"),
        "avi" => Some("video/x-msvideo"),
        "mov" => Some("video/quicktime"),
        "mkv" => Some("video/x-matroska"),

        // Documents
        "pdf" => Some("application/pdf"),
        "txt" => Some("text/plain"),
        "json" => Some("application/json"),
        "xml" => Some("application/xml"),

        // Archives
        "zip" => Some("application/zip"),
        "tar" => Some("application/x-tar"),
        "gz" => Some("application/gzip"),
        "7z" => Some("application/x-7z-compressed"),

        _ => None,
    }
}

// ── File editing API ──────────────────────────────────────────────────────────

#[derive(Deserialize)]
struct SaveFileRequest {
    workspace_id: String,
    file_path: String,
    content: String,
}

#[derive(Serialize)]
struct SaveFileResponse {
    success: bool,
    message: String,
}

async fn save_file_handler(
    State(state): State<AppState>,
    Json(payload): Json<SaveFileRequest>,
) -> impl IntoResponse {
    let ws = match state.workspace_registry.get(&payload.workspace_id) {
        Some(w) => w,
        None => {
            return Json(SaveFileResponse {
                success: false,
                message: "Workspace not found".into(),
            })
            .into_response()
        }
    };

    if !ws.enable_edit.load(std::sync::atomic::Ordering::Relaxed) {
        return Json(SaveFileResponse {
            success: false,
            message: "Edit feature is not enabled".into(),
        })
        .into_response();
    }

    let decoded = match urlencoding::decode(&payload.file_path) {
        Ok(p) => p,
        Err(_) => {
            return Json(SaveFileResponse {
                success: false,
                message: "Invalid file path encoding".into(),
            })
            .into_response()
        }
    };

    let full_path = ws.root.join(decoded.trim_start_matches('/'));
    let canonical = match full_path.canonicalize() {
        Ok(p) => p,
        Err(_) => {
            return Json(SaveFileResponse {
                success: false,
                message: format!("File not found: {decoded}"),
            })
            .into_response()
        }
    };

    if !canonical.starts_with(&ws.root) {
        return Json(SaveFileResponse {
            success: false,
            message: "Access denied".into(),
        })
        .into_response();
    }
    if !canonical.is_file() {
        return Json(SaveFileResponse {
            success: false,
            message: "Path is not a file".into(),
        })
        .into_response();
    }
    if canonical
        .extension()
        .is_none_or(|e| e.to_string_lossy().to_lowercase() != "md")
    {
        return Json(SaveFileResponse {
            success: false,
            message: "Only .md files can be edited".into(),
        })
        .into_response();
    }
    if fs::metadata(&canonical)
        .map(|m| m.permissions().readonly())
        .unwrap_or(true)
    {
        return Json(SaveFileResponse {
            success: false,
            message: "File is read-only".into(),
        })
        .into_response();
    }

    match fs::write(&canonical, &payload.content) {
        Ok(_) => Json(SaveFileResponse {
            success: true,
            message: "File saved successfully".into(),
        })
        .into_response(),
        Err(e) => Json(SaveFileResponse {
            success: false,
            message: format!("Failed to save: {e}"),
        })
        .into_response(),
    }
}