stash/src/main.rs

use std::path::PathBuf;
use clap::Parser;
use colored::*;
use stash::{
    Config,
    config::{default_data_dir, display_path, resolve_db_path, validate_db_location},
    start,
};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

#[derive(Parser, Debug)]
#[command(
    name = "stash",
    author = "stash developers",
    version,
    about = "🚀 本地静态页面托管与多项目管理工具",
    long_about = "支持本地静态页面按目录隔离托管，支持前端直接调用同源接口，提供清爽的项目导航列表。"
)]
struct Cli {
    /// 静态页面托管根目录（支持相对路径和绝对路径）
    #[arg(short, long, default_value = "./public", value_name = "DIR")]
    root: PathBuf,

    /// 数据目录（默认是与二进制同级的 data/，用于存放 stash.db）
    #[arg(long, value_name = "DIR")]
    data: Option<PathBuf>,

    /// 监听端口
    #[arg(short, long, default_value_t = 8080, value_name = "PORT")]
    port: u16,

    /// 绑定监听地址（默认 127.0.0.1，如需局域网访问可传 0.0.0.0）
    #[arg(long, default_value = "127.0.0.1", value_name = "HOST")]
    host: String,

    /// 关闭根目录项目导航列表功能（仅托管静态文件）
    #[arg(long, default_value_t = false)]
    no_listing: bool,

    /// 关闭 CORS 跨域标头注入
    #[arg(long, default_value_t = false)]
    no_cors: bool,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // 1. 初始化 tracing 日志
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "info".into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    // 2. 解析命令行输入
    let cli = Cli::parse();

    // 检查目录是否存在，不存在则友好提示并自动创建默认 public
    if !cli.root.exists() {
        if cli.root == PathBuf::from("./public") {
            println!(
                "{} 默认托管目录 {} 不存在，正在自动创建...",
                "[INFO]".cyan(),
                "./public".yellow()
            );
            std::fs::create_dir_all(&cli.root)?;
        } else {
            eprintln!(
                "{} 指定的托管目录不存在: {}",
                "[ERROR]".red().bold(),
                cli.root.display().to_string().yellow()
            );
            std::process::exit(1);
        }
    }

    let abs_root = cli.root.canonicalize().unwrap_or_else(|_| cli.root.clone());

    // 3. 解析数据目录：默认跟随二进制（双击启动时 CWD 不可控），可用 --data 覆盖
    let data_dir = match cli.data {
        Some(dir) => dir,
        None => default_data_dir(),
    };

    // 先校验再建目录：配置错了就不该在磁盘上留下任何痕迹
    let db_path = resolve_db_path(&data_dir);
    if let Err(error) = validate_db_location(&abs_root, &db_path) {
        eprintln!("{} {}", "[ERROR]".red().bold(), error);
        std::process::exit(1);
    }

    std::fs::create_dir_all(&data_dir)?;
    let data_dir = data_dir.canonicalize().unwrap_or(data_dir);
    let db_path = resolve_db_path(&data_dir);

    // 4. 构造配置并启动
    let config = Config {
        root: abs_root.clone(),
        data_dir: data_dir.clone(),
        db_path: db_path.clone(),
        port: cli.port,
        host: cli.host,
        directory_listing: !cli.no_listing,
        enable_cors: !cli.no_cors,
    };

    let handle = match start(config).await {
        Ok(h) => h,
        Err(err) => {
            eprintln!("{} 启动失败: {}", "[FATAL]".red().bold(), err);
            std::process::exit(1);
        }
    };

    // 5. 打印启动成功 Banner
    println!();
    println!("{}", "==================================================".bright_blue());
    println!("  ✨ {} {}", "Stash Server".green().bold(), format!("v{}", env!("CARGO_PKG_VERSION")).dimmed());
    println!("{}", "==================================================".bright_blue());
    println!("  ➜ 本地访问:    {}", format!("http://{}", handle.local_addr).cyan().bold().underline());
    println!("  ➜ 托管路径:    {}", display_path(&abs_root).white());
    println!("  ➜ 内容库:      {}", display_path(&db_path).white());
    println!("  ➜ 前端资源:    {}", "已内嵌 (stash-web/dist)".white());
    println!("  ➜ 接口列表:    {}", format!("http://{}/api/projects", handle.local_addr).dimmed());
    println!("{}", "--------------------------------------------------".bright_blue());
    println!("  按 {} 停止服务", "Ctrl+C".yellow().bold());
    println!();

    // 6. 跨平台捕获退出信号
    wait_for_shutdown_signal().await;

    println!("\n{} 正在停止服务并释放端口...", "[SYS]".yellow());
    handle.stop();

    // 留出时间让 tokio 优雅关闭 socket 句柄
    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
    println!("{} 服务已安全退出。\n", "[DONE]".green().bold());

    Ok(())
}

/// 监听停机信号（支持 Windows Ctrl+C 以及 Linux/macOS SIGINT/SIGTERM）
async fn wait_for_shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("无法注册 Ctrl+C 信号处理器");
    };

    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("无法注册 SIGTERM 信号处理器")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }
}