use clap::CommandFactory;
use clap::Parser;
use clap_complete::Shell;
use clap_complete::generate;
use codex_arg0::arg0_dispatch_or_else;
use codex_chatgpt::apply_command::ApplyCommand;
use codex_chatgpt::apply_command::run_apply_command;
use codex_cli::LandlockCommand;
use codex_cli::SeatbeltCommand;
use codex_cli::login::run_login_status;
use codex_cli::login::run_login_with_api_key;
use codex_cli::login::run_login_with_chatgpt;
use codex_cli::proto;
use codex_common::CliConfigOverrides;
use codex_exec::Cli as ExecCli;
use codex_tui::Cli as TuiCli;
use std::path::PathBuf;

use crate::proto::ProtoCli;
mod concurrent;

/// Codex CLI
///
/// If no subcommand is specified, options will be forwarded to the interactive CLI.
#[derive(Debug, Parser)]
#[clap(
    author,
    version,
    // If a sub‑command is given, ignore requirements of the default args.
    subcommand_negates_reqs = true
)]
struct MultitoolCli {
    #[clap(flatten)]
    pub config_overrides: CliConfigOverrides,

    /// Experimental:Launch a concurrent task in a separate Git worktree using the given prompt.
    /// Creates worktree under $CODEX_HOME/worktrees/<repo>/codex/<slug> and runs `codex exec` in full-auto mode.
    #[arg(long = "concurrent", value_name = "PROMPT")]
    pub concurrent: Option<String>,

    /// When using --concurrent, also attempt to auto-merge the resulting changes
    /// back into the current working tree as unstaged modifications via
    /// a 3-way git apply. Disable with --automerge=false.
    #[arg(long = "automerge", default_value_t = true, action = clap::ArgAction::Set)]
    pub automerge: bool,

    /// Run the same --concurrent prompt N times in separate worktrees and keep them all.
    /// Intended to generate multiple candidate solutions without auto-merging.
    #[arg(long = "best-of-n", value_name = "N", default_value_t = 1)]
    pub best_of_n: usize,

    #[clap(flatten)]
    interactive: TuiCli,

    #[clap(subcommand)]
    subcommand: Option<Subcommand>,
}

#[derive(Debug, clap::Subcommand)]
enum Subcommand {
    /// Run Codex non-interactively.
    #[clap(visible_alias = "e")]
    Exec(ExecCli),

    /// Manage login.
    Login(LoginCommand),

    /// Experimental: run Codex as an MCP server.
    Mcp,

    /// Run the Protocol stream via stdin/stdout
    #[clap(visible_alias = "p")]
    Proto(ProtoCli),

    /// Generate shell completion scripts.
    Completion(CompletionCommand),

    /// Internal debugging commands.
    Debug(DebugArgs),

    /// Apply the latest diff produced by Codex agent as a `git apply` to your local working tree.
    #[clap(visible_alias = "a")]
    Apply(ApplyCommand),
}

#[derive(Debug, Parser)]
struct CompletionCommand {
    /// Shell to generate completions for
    #[clap(value_enum, default_value_t = Shell::Bash)]
    shell: Shell,
}

#[derive(Debug, Parser)]
struct DebugArgs {
    #[command(subcommand)]
    cmd: DebugCommand,
}

#[derive(Debug, clap::Subcommand)]
enum DebugCommand {
    /// Run a command under Seatbelt (macOS only).
    Seatbelt(SeatbeltCommand),

    /// Run a command under Landlock+seccomp (Linux only).
    Landlock(LandlockCommand),
}

#[derive(Debug, Parser)]
struct LoginCommand {
    #[clap(skip)]
    config_overrides: CliConfigOverrides,

    #[arg(long = "api-key", value_name = "API_KEY")]
    api_key: Option<String>,

    #[command(subcommand)]
    action: Option<LoginSubcommand>,
}

#[derive(Debug, clap::Subcommand)]
enum LoginSubcommand {
    /// Show login status.
    Status,
}

fn main() -> anyhow::Result<()> {
    arg0_dispatch_or_else(|codex_linux_sandbox_exe| async move {
        cli_main(codex_linux_sandbox_exe).await?;
        Ok(())
    })
}

async fn cli_main(codex_linux_sandbox_exe: Option<PathBuf>) -> anyhow::Result<()> {
    let cli = MultitoolCli::parse();

    // Handle --concurrent at the root level.
    if let Some(prompt) = cli.concurrent.clone() {
        if cli.subcommand.is_some() {
            eprintln!("--concurrent cannot be used together with a subcommand");
            std::process::exit(2);
        }
        let runs = if cli.best_of_n == 0 { 1 } else { cli.best_of_n };
        if runs > 1 {
            println!(
                "Running best-of-n with {runs} runs; auto-merge will be disabled and worktrees kept."
            );

            // Launch all runs concurrently and collect results as they finish.
            let mut join_set = tokio::task::JoinSet::new();
            for _ in 0..runs {
                let prompt = prompt.clone();
                let overrides = cli.config_overrides.clone();
                let sandbox = codex_linux_sandbox_exe.clone();
                join_set.spawn(async move {
                    concurrent::run_concurrent_flow_quiet_no_automerge(prompt, overrides, sandbox)
                        .await
                });
            }

            let mut results: Vec<concurrent::ConcurrentRunResult> = Vec::with_capacity(runs);
            while let Some(join_result) = join_set.join_next().await {
                match join_result {
                    Ok(Ok(res)) => {
                        println!(
                            "task finished for branch: {}\n, directory: {}",
                            res.branch,
                            res.worktree_dir.display()
                        );
                        results.push(res);
                    }
                    Ok(Err(err)) => {
                        eprintln!("concurrent task failed: {err}");
                    }
                    Err(join_err) => {
                        eprintln!("failed to join concurrent task: {join_err}");
                    }
                }
            }

            println!("\nBest-of-n summary:");
            for r in &results {
                let status = match r.exec_exit_code {
                    Some(0) => "OK",
                    Some(_code) => "FAIL",
                    None => "OK",
                };
                let log = r
                    .log_file
                    .as_ref()
                    .map(|p| p.to_string_lossy().to_string())
                    .unwrap_or_else(|| "<no log>".to_string());
                println!(
                    "[{status}] branch={} worktree={} log={}",
                    r.branch,
                    r.worktree_dir.display(),
                    log
                );
            }
        } else {
            concurrent::run_concurrent_flow(
                prompt,
                cli.config_overrides,
                codex_linux_sandbox_exe,
                cli.automerge,
                false,
            )
            .await?;
        }
        return Ok(());
    }

    if cli.best_of_n > 1 {
        eprintln!("--best-of-n requires --concurrent <PROMPT>");
        std::process::exit(2);
    }

    match cli.subcommand {
        None => {
            let mut tui_cli = cli.interactive;
            prepend_config_flags(&mut tui_cli.config_overrides, cli.config_overrides);
            let usage = codex_tui::run_main(tui_cli, codex_linux_sandbox_exe).await?;
            println!("{}", codex_core::protocol::FinalOutput::from(usage));
        }
        Some(Subcommand::Exec(mut exec_cli)) => {
            prepend_config_flags(&mut exec_cli.config_overrides, cli.config_overrides);
            codex_exec::run_main(exec_cli, codex_linux_sandbox_exe).await?;
        }
        Some(Subcommand::Mcp) => {
            codex_mcp_server::run_main(codex_linux_sandbox_exe).await?;
        }
        Some(Subcommand::Login(mut login_cli)) => {
            prepend_config_flags(&mut login_cli.config_overrides, cli.config_overrides);
            match login_cli.action {
                Some(LoginSubcommand::Status) => {
                    run_login_status(login_cli.config_overrides).await;
                }
                None => {
                    if let Some(api_key) = login_cli.api_key {
                        run_login_with_api_key(login_cli.config_overrides, api_key).await;
                    } else {
                        run_login_with_chatgpt(login_cli.config_overrides).await;
                    }
                }
            }
        }
        Some(Subcommand::Proto(mut proto_cli)) => {
            prepend_config_flags(&mut proto_cli.config_overrides, cli.config_overrides);
            proto::run_main(proto_cli).await?;
        }
        Some(Subcommand::Completion(completion_cli)) => {
            print_completion(completion_cli);
        }
        Some(Subcommand::Debug(debug_args)) => match debug_args.cmd {
            DebugCommand::Seatbelt(mut seatbelt_cli) => {
                prepend_config_flags(&mut seatbelt_cli.config_overrides, cli.config_overrides);
                codex_cli::debug_sandbox::run_command_under_seatbelt(
                    seatbelt_cli,
                    codex_linux_sandbox_exe,
                )
                .await?;
            }
            DebugCommand::Landlock(mut landlock_cli) => {
                prepend_config_flags(&mut landlock_cli.config_overrides, cli.config_overrides);
                codex_cli::debug_sandbox::run_command_under_landlock(
                    landlock_cli,
                    codex_linux_sandbox_exe,
                )
                .await?;
            }
        },
        Some(Subcommand::Apply(mut apply_cli)) => {
            prepend_config_flags(&mut apply_cli.config_overrides, cli.config_overrides);
            run_apply_command(apply_cli, None).await?;
        }
    }

    Ok(())
}

/// Prepend root-level overrides so they have lower precedence than
/// CLI-specific ones specified after the subcommand (if any).
fn prepend_config_flags(
    subcommand_config_overrides: &mut CliConfigOverrides,
    cli_config_overrides: CliConfigOverrides,
) {
    subcommand_config_overrides
        .raw_overrides
        .splice(0..0, cli_config_overrides.raw_overrides);
}

fn print_completion(cmd: CompletionCommand) {
    let mut app = MultitoolCli::command();
    let name = "codex";
    generate(cmd.shell, &mut app, name, &mut std::io::stdout());
}