fcmc 0.9.4

use bollard::Docker;
use clap::{Parser, ValueEnum};
use colored::*;
use fcmc::{ChallengeMeta, GameBoxMeta};
use std::{
    fs,
    io::{self, Write},
    path::PathBuf,
};

#[derive(Debug, Clone, ValueEnum)]
#[value(rename_all = "snake_case")]
enum GenFormat {
    #[value(alias = "c")]
    Challenge,
    #[value(alias = "g")]
    Gamebox,
    #[value(alias = "t")]
    Target,
}
#[derive(Parser, Debug)]
#[command(name = "fcmc", about = "FloatCTF 题目配置检查和管理工具")]
#[command(version = env!("CARGO_PKG_VERSION"))]
struct Args {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Parser, Debug, Clone)]
#[command(rename_all = "snake_case")]
enum Commands {
    /// 检查题目配置文件是否合法
    Check {
        /// 配置文件目录 (里面需要包含 meta.toml)
        #[arg(short, long)]
        path: Option<String>,
    },
    /// 构建题目镜像
    Build {
        /// 配置文件目录 (里面需要包含 meta.toml)
        #[arg(short, long)]
        path: Option<String>,
        /// 生成模板类型: challenge (c) | gamebox (g) | target (t)
        #[arg(short, long, default_value = "challenge")]
        format: GenFormat,
    },
    /// 生成新的题目模板
    Gen {
        /// 新题目的名称
        #[arg(short, long)]
        name: String,

        /// 输出目录
        #[arg(short, long, default_value = ".")]
        output: String,

        /// 生成模板类型: challenge (c) | gamebox (g) | target (t)
        #[arg(short, long, default_value = "challenge")]
        format: GenFormat,

        /// gamebox 模板 (仅 format=gamebox 时生效)
        #[arg(short, long, default_value = "false")]
        template: bool,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    match args.command {
        Commands::Check { path } => {
            // 如果指定了 --file，用它作为目录；否则默认当前目录
            let mut pass = true;

            let dir = path.unwrap_or_else(|| ".".to_string());
            let path = PathBuf::from(&dir).join("meta.toml");

            println!("\n================ 配置检查报告 ================\n");
            println!("配置文件: {:?}", path);

            let content = fs::read_to_string(&path)?;
            let cfg = toml::from_str::<ChallengeMeta>(&content);
            //  检查配置文件是否合法
            println!("\n[解析结果]");
            match cfg {
                Ok(cfg) => {
                    println!("  {}   配置文件解析成功", "OK".green());

                    println!("\n[附件检查]");
                    // 检查附件
                    if let Some(attachment) = &cfg.attachment {
                        let attachment_path = PathBuf::from(&dir).join(attachment);
                        if attachment_path.exists() {
                            let size = fs::metadata(&attachment_path)?.len();
                            println!(
                                "  {}   附件存在: {:?} ({} bytes)",
                                "OK".green(),
                                attachment_path,
                                size
                            );
                        } else {
                            println!("  {}   附件不存在: {:?}", "ERR".red(), attachment_path);
                            pass = false;
                        }
                    } else {
                        println!("  {}   未配置附件", "WARN".yellow());
                    }

                    // test docker
                    println!("\n[Docker 检查]");
                    if cfg.docker.is_some() {
                        match test_docker(&cfg).await {
                            Ok(_) => (),
                            Err(e) => {
                                println!("  {}   Docker 测试失败: {}", "ERR".red(), e);
                                println!("{} fcmc build -p {}", "  尝试构建镜像:".cyan(), dir);
                                pass = false;
                            }
                        }
                    } else {
                        println!("  {}   未配置 Docker", "WARN".yellow());
                    }
                }
                Err(e) => {
                    println!("  {}   配置文件解析失败: {}", "ERR".red(), e);
                    pass = false;
                }
            }

            if pass {
                println!("\n----------------------------------------------");
                println!("最终结果: {}", "通过".green());
                println!("==============================================\n");
            } else {
                println!("\n----------------------------------------------");
                println!("最终结果: {}", "失败".red());
                println!("==============================================\n");
            }
        }
        Commands::Gen {
            name,
            output,
            format,
            template,
        } => match format {
            GenFormat::Challenge => {
                ChallengeMeta::generate_template(&name, &output).await?;
            }
            GenFormat::Gamebox => {
                if template {
                    GameBoxMeta::generate_basic_template(&name, &output).await?;
                } else {
                    GameBoxMeta::generate_template(&name, &output).await?;
                }
            }
            GenFormat::Target => {
                todo!("target 模板生成");
            }
        },
        Commands::Build { path, format } => {
            let dir = path.unwrap_or_else(|| ".".to_string());
            let path = PathBuf::from(&dir).join("meta.toml");

            let content = fs::read_to_string(&path)?;
            let src_dir = PathBuf::from(&dir).join("src");

            let docker = Docker::connect_with_defaults()?;

            match format {
                GenFormat::Challenge => {
                    let cfg = toml::from_str::<ChallengeMeta>(&content)?;
                    cfg.build_image(&docker, &src_dir).await?;
                }
                GenFormat::Gamebox => {
                    let cfg = toml::from_str::<GameBoxMeta>(&content)?;
                    cfg.build_image(&docker, &src_dir).await?;
                }
                GenFormat::Target => {
                    todo!("target build");
                }
            }
        }
    }

    Ok(())
}

async fn test_docker(cm: &ChallengeMeta) -> anyhow::Result<()> {
    let docker = Docker::connect_with_defaults()?;

    let flag = "flag{test-test-test-test}";

    match cm.create_and_start(&docker, "local_test", flag).await {
        Ok(port) => {
            let name = format!("{}_local_test", cm.name);
            println!("  {} 已启动容器: {}", "OK".green(), name);
            println!("  {} 访问地址: http://localhost:{}", "INFO".cyan(), port);

            // kill
            print!("  INFO 按回车键继续以清理容器...");
            io::stdout().flush()?;
            let mut input = String::new();
            io::stdin().read_line(&mut input)?;

            match fcmc::stop_and_remove(&docker, "local_test").await {
                Ok(_) => println!("  {} 已清理容器: {}", "OK".green(), name),
                Err(e) => println!("  {} 容器清理失败: {}", "ERR".red(), e),
            }
        }
        Err(e) => {
            return Err(e);
        }
    }

    Ok(())
}