//! offload CLI - Flexible parallel test runner.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};

use anyhow::{Context, Result, anyhow};
use clap::{Parser, Subcommand};
use tracing::{Level, info, warn};
use tracing_subscriber::FmtSubscriber;

use offload::config::{
    self, CargoFrameworkConfig, Config, DefaultFrameworkConfig, DefaultProviderConfig,
    FrameworkConfig, GroupConfig, LocalProviderConfig, OffloadConfig, ProviderConfig,
    PytestFrameworkConfig, ReportConfig, SandboxConfig, VitestFrameworkConfig,
};
use offload::framework::{
    TestFramework, TestRecord, cargo::CargoFramework, default::DefaultFramework,
    pytest::PytestFramework, vitest::VitestFramework,
};
use offload::image_cache;
use offload::orchestrator::{Orchestrator, SandboxPool};
use offload::provider::{
    PrepareContext, SandboxProvider, default::DefaultProvider, local::LocalProvider,
    modal::ModalProvider,
};

/// A directory copy directive: local path -> sandbox path
#[derive(Debug, Clone)]
pub struct CopyDir {
    pub local: PathBuf,
    pub remote: PathBuf,
}

#[derive(Parser)]
#[command(name = "offload")]
#[command(about = "Flexible parallel test runner", long_about = None)]
#[command(version)]
struct Cli {
    /// Configuration file path
    #[arg(short, long, global = true, default_value = "offload.toml")]
    config: PathBuf,

    /// Verbose output
    #[arg(short, long, global = true)]
    verbose: bool,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Run tests
    Run {
        /// Override maximum parallel sandboxes
        #[arg(short, long)]
        parallel: Option<usize>,

        /// Only discover tests, don't run them
        #[arg(long)]
        collect_only: bool,

        /// Directories to copy to sandbox (format: /local/path:/sandbox/path)
        #[arg(long, value_name = "LOCAL:REMOTE")]
        copy_dir: Vec<String>,

        /// Environment variables to set in sandboxes (format: KEY=VALUE)
        #[arg(long = "env", value_name = "KEY=VALUE")]
        env_vars: Vec<String>,

        /// Skip cached image lookup during prepare (forces fresh build)
        #[arg(long)]
        no_cache: bool,

        /// Emit a Perfetto trace to {output_dir}/trace.json
        #[arg(long)]
        trace: bool,

        /// Show estimated sandbox cost after run.
        ///
        /// Note: This is calculated client-side using simple formulas and
        /// may not reflect actual billing, discounts, or pricing adjustments.
        #[arg(long)]
        show_estimated_cost: bool,

        /// Stop immediately when a test failure is detected
        #[arg(long)]
        fail_fast: bool,

        /// CI mode: replace progress bars with plain-text log lines
        #[arg(long)]
        ci: bool,

        /// Record test history after the run.
        ///
        /// Requires a [history] section in the config file.
        #[arg(long)]
        record_history: bool,
    },

    /// Discover tests without running them
    Collect {
        /// Output format (text, json)
        #[arg(short, long, default_value = "text")]
        format: String,
    },

    /// Validate configuration file
    Validate,

    /// Initialize a new configuration file
    Init {
        /// Provider type (local, default)
        #[arg(short, long, default_value = "local")]
        provider: String,

        /// Test framework (pytest, nextest, vitest, default)
        #[arg(short, long, default_value = "pytest")]
        framework: String,
    },

    /// Show checkpoint cache status for the current HEAD.
    CheckpointStatus {
        #[arg(long, default_value = "origin")]
        remote: String,
    },

    /// Build the sandbox image without running tests.
    ///
    /// Prepares the provider image (resolving cache, building if needed)
    /// and writes the image ID to git notes. Useful for CI warm-up or
    /// pre-building images before a test run.
    Build {
        /// Skip cached image lookup during prepare (forces fresh build)
        #[arg(long)]
        no_cache: bool,
    },

    /// View test run logs
    Logs {
        /// Show only failure logs
        #[arg(long)]
        failures: bool,

        /// Show only error logs
        #[arg(long)]
        errors: bool,

        /// Show only tests matching this exact ID (repeatable)
        #[arg(long)]
        test: Vec<String>,

        /// Show only tests whose ID matches this regex (substring match)
        #[arg(long)]
        test_regex: Option<String>,
    },

    /// History management commands
    History {
        #[command(subcommand)]
        subcommand: HistoryCommands,
    },

    /// Apply a git-format binary patch to the filesystem.
    ///
    /// Intended for use inside sandbox images during the thin-diff step.
    /// Uses diffy for patch application instead of git apply.
    ApplyDiff {
        /// Path to the patch file
        patch_file: PathBuf,

        /// Project root directory to apply patches relative to
        #[arg(long, default_value = ".")]
        project_root: PathBuf,
    },
}

/// Subcommands for history management.
#[derive(Subcommand)]
enum HistoryCommands {
    /// Git merge driver for history files.
    ///
    /// Usage: offload history merge <base> <ours> <theirs>
    ///
    /// This implements the git merge driver protocol. The merged result
    /// is written to the "ours" file.
    Merge {
        /// Base/ancestor version (%O)
        base: PathBuf,
        /// Our version (%A) - modified in place with merge result
        ours: PathBuf,
        /// Their version (%B)
        theirs: PathBuf,
    },

    /// Configure git merge driver for history files.
    ///
    /// Updates .gitattributes and configures the merge driver in .git/config.
    /// This enables automatic merging of offload-history.jsonl during git operations.
    SetupMergeDriver,
}

#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    // Set up logging
    let log_level = if cli.verbose {
        Level::INFO
    } else {
        Level::WARN
    };
    let subscriber = FmtSubscriber::builder()
        .with_max_level(log_level)
        .with_target(false)
        .finish();
    tracing::subscriber::set_global_default(subscriber)?;

    match cli.command {
        Commands::Run {
            parallel,
            collect_only,
            copy_dir,
            env_vars,
            no_cache,
            trace,
            show_estimated_cost,
            fail_fast,
            ci,
            record_history,
        } => {
            let ci = ci || std::env::var("CI").is_ok_and(|v| v == "true");
            run_tests(
                &cli.config,
                parallel,
                collect_only,
                copy_dir,
                env_vars,
                no_cache,
                cli.verbose,
                trace,
                show_estimated_cost,
                fail_fast,
                ci,
                record_history,
            )
            .await
        }
        Commands::Collect { format } => collect_tests(&cli.config, &format).await,
        Commands::Validate => validate_config(&cli.config),
        Commands::Init {
            provider,
            framework,
        } => init_config(&provider, &framework),
        Commands::CheckpointStatus { remote } => {
            let cwd = std::env::current_dir().context("failed to get current directory")?;
            let config_path_str = cli.config.to_string_lossy().to_string();
            image_cache::status_handler(&cwd, &config_path_str, &remote).await
        }
        Commands::Build { no_cache } => build_image(&cli.config, no_cache).await,
        Commands::Logs {
            failures,
            errors,
            test,
            test_regex,
        } => show_logs(&cli.config, failures, errors, &test, test_regex.as_deref()),
        Commands::History { subcommand } => match subcommand {
            HistoryCommands::Merge { base, ours, theirs } => {
                // Default reservoir size matches the history config default
                offload::history::merge::merge_history_files(&base, &ours, &theirs, 20)?;
                Ok(())
            }
            HistoryCommands::SetupMergeDriver => setup_merge_driver(),
        },
        Commands::ApplyDiff {
            patch_file,
            project_root,
        } => apply_diff(&patch_file, &project_root),
    }
}

/// Helper to get framework type name for validation.
fn framework_type_name(framework: &FrameworkConfig) -> &'static str {
    match framework {
        FrameworkConfig::Pytest(_) => "pytest",
        FrameworkConfig::Cargo(_) => "nextest",
        FrameworkConfig::Default(_) => "default",
        FrameworkConfig::Vitest(_) => "vitest",
    }
}

/// Discover tests for every group, tagging each with its group config.
async fn discover_all_tests(
    framework: &FrameworkConfig,
    groups: &HashMap<String, GroupConfig>,
) -> Result<Vec<TestRecord>> {
    let mut all_tests: Vec<TestRecord> = Vec::new();

    for (group_name, group_cfg) in groups {
        let tests = match framework {
            FrameworkConfig::Pytest(cfg) => {
                PytestFramework::new(cfg.clone())?
                    .discover(&[], &group_cfg.filters, group_name)
                    .await?
            }
            FrameworkConfig::Cargo(cfg) => {
                CargoFramework::new(cfg.clone())
                    .discover(&[], &group_cfg.filters, group_name)
                    .await?
            }
            FrameworkConfig::Default(cfg) => {
                DefaultFramework::new(cfg.clone())
                    .discover(&[], &group_cfg.filters, group_name)
                    .await?
            }
            FrameworkConfig::Vitest(cfg) => {
                VitestFramework::new(cfg.clone())?
                    .discover(&[], &group_cfg.filters, group_name)
                    .await?
            }
        };

        // Tag tests with group retry count
        let group_tests: Vec<TestRecord> = tests
            .into_iter()
            .map(|t| {
                t.with_retry_count(group_cfg.retry_count)
                    .with_schedule_individual(group_cfg.schedule_individual)
            })
            .collect();

        all_tests.extend(group_tests);
    }

    Ok(all_tests)
}

/// Discover tests concurrently with provider preparation, signalling completion.
async fn discover_with_signal(
    framework: &FrameworkConfig,
    groups: &HashMap<String, GroupConfig>,
    discovery_done: &AtomicBool,
) -> Result<Vec<TestRecord>> {
    eprintln!("[discover] Discovering tests...");
    let result = discover_all_tests(framework, groups).await;
    if let Ok(ref tests) = result {
        eprintln!(
            "[discover] found {} tests across {} groups",
            tests.len(),
            groups.len()
        );
    }
    discovery_done.store(true, Ordering::Release);
    result
}

/// Dispatch test execution to the appropriate framework, using the given provider.
#[allow(clippy::too_many_arguments)]
async fn dispatch_framework<P: offload::provider::SandboxProvider>(
    config: &Config,
    config_filename: &str,
    run_id: &str,
    all_tests: &[TestRecord],
    provider: P,
    copy_dirs: &[CopyDir],
    verbose: bool,
    tracer: &offload::trace::Tracer,
    show_estimated_cost: bool,
    fail_fast: bool,
    ci: bool,
    record_history: bool,
) -> Result<i32> {
    match &config.framework {
        FrameworkConfig::Pytest(f_cfg) => {
            run_all_tests(
                config,
                config_filename,
                run_id,
                all_tests,
                provider,
                PytestFramework::new(f_cfg.clone())?,
                copy_dirs,
                verbose,
                tracer,
                show_estimated_cost,
                fail_fast,
                ci,
                record_history,
            )
            .await
        }
        FrameworkConfig::Cargo(f_cfg) => {
            run_all_tests(
                config,
                config_filename,
                run_id,
                all_tests,
                provider,
                CargoFramework::new(f_cfg.clone()),
                copy_dirs,
                verbose,
                tracer,
                show_estimated_cost,
                fail_fast,
                ci,
                record_history,
            )
            .await
        }
        FrameworkConfig::Default(f_cfg) => {
            if fail_fast {
                warn!(
                    "--fail-fast: the default framework does not pass a stop flag to the test runner. Batches will still be cancelled on failure, but tests within a running batch will not stop early."
                );
            }
            run_all_tests(
                config,
                config_filename,
                run_id,
                all_tests,
                provider,
                DefaultFramework::new(f_cfg.clone()),
                copy_dirs,
                verbose,
                tracer,
                show_estimated_cost,
                fail_fast,
                ci,
                record_history,
            )
            .await
        }
        FrameworkConfig::Vitest(f_cfg) => {
            run_all_tests(
                config,
                config_filename,
                run_id,
                all_tests,
                provider,
                VitestFramework::new(f_cfg.clone())?,
                copy_dirs,
                verbose,
                tracer,
                show_estimated_cost,
                fail_fast,
                ci,
                record_history,
            )
            .await
        }
    }
}

/// Construct a [`PrepareContext`] and run the provider's prepare step.
///
/// Shared by `offload build` (standalone) and `offload run` (concurrent
/// with test discovery).
#[allow(clippy::too_many_arguments)]
async fn run_prepare<P: SandboxProvider>(
    provider: &mut P,
    repo: &Path,
    config: &Config,
    config_path: &Path,
    copy_dirs: &[(PathBuf, PathBuf)],
    no_cache: bool,
    tracer: &offload::trace::Tracer,
    discovery_done: &AtomicBool,
) -> Result<Option<String>> {
    let prepare_ctx = PrepareContext {
        copy_dirs,
        sandbox_init_cmd: config.offload.sandbox_init_cmd.as_deref(),
        post_patch_cmd: config.offload.post_patch_cmd.as_deref(),
        repo,
        config,
        config_path,
        no_cache,
        tracer,
        discovery_done,
    };

    provider
        .prepare(&prepare_ctx)
        .await
        .map_err(|e| anyhow::anyhow!(e))
}

/// Shared logic for Default and Modal provider arms: concurrent discovery +
/// prepare (which internally handles cache resolution, thin-diff, and
/// full-build fallback), then framework dispatch.
#[allow(clippy::too_many_arguments)]
async fn run_remote_provider<P: SandboxProvider>(
    repo: &Path,
    mut provider: P,
    config: &Config,
    copy_dir_tuples: &[(PathBuf, PathBuf)],
    config_filename: &str,
    run_id: &str,
    copy_dirs: &[CopyDir],
    no_cache: bool,
    verbose: bool,
    tracer: &offload::trace::Tracer,
    show_estimated_cost: bool,
    fail_fast: bool,
    config_path: &Path,
    ci: bool,
    record_history: bool,
) -> Result<Option<i32>> {
    let discovery_done = AtomicBool::new(false);

    let (all_tests, _) = tokio::try_join!(
        discover_with_signal(&config.framework, &config.groups, &discovery_done),
        run_prepare(
            &mut provider,
            repo,
            config,
            config_path,
            copy_dir_tuples,
            no_cache,
            tracer,
            &discovery_done,
        ),
    )?;

    if all_tests.is_empty() {
        return Ok(None);
    }

    dispatch_framework(
        config,
        config_filename,
        run_id,
        &all_tests,
        provider,
        copy_dirs,
        verbose,
        tracer,
        show_estimated_cost,
        fail_fast,
        ci,
        record_history,
    )
    .await
    .map(Some)
}

#[allow(clippy::too_many_arguments)]
async fn run_tests(
    config_path: &Path,
    parallel_override: Option<usize>,
    collect_only: bool,
    copy_dir_args: Vec<String>,
    env_vars: Vec<String>,
    no_cache: bool,
    verbose: bool,
    trace: bool,
    show_estimated_cost: bool,
    fail_fast: bool,
    ci: bool,
    record_history: bool,
) -> Result<()> {
    let tracer = if trace {
        offload::trace::Tracer::new()
    } else {
        offload::trace::Tracer::noop()
    };

    tracer.metadata_event(
        "process_name",
        offload::trace::PID_LOCAL,
        offload::trace::TID_MAIN,
        serde_json::json!({"name": "Offload (Local)"}),
    );
    tracer.metadata_event(
        "thread_name",
        offload::trace::PID_LOCAL,
        offload::trace::TID_MAIN,
        serde_json::json!({"name": "Main"}),
    );

    // Load configuration
    let mut config = config::load_config(config_path)
        .with_context(|| format!("Failed to load config from {}", config_path.display()))?;

    // Validate --record-history flag
    if record_history && config.history.is_none() {
        anyhow::bail!("--record-history requires a [history] section in the config file");
    }

    // Generate run ID for history recording
    let run_id = offload::generate_run_id();
    let config_filename = config_path
        .file_name()
        .and_then(|s| s.to_str())
        .unwrap_or("offload.toml")
        .to_string();

    // Apply overrides
    if let Some(parallel) = parallel_override {
        config.offload.max_parallel = parallel;
    }

    // Parse copy_dir arguments
    let copy_dirs: Vec<CopyDir> = copy_dir_args
        .iter()
        .map(|arg| {
            let parts: Vec<&str> = arg.splitn(2, ':').collect();
            if parts.len() != 2 {
                return Err(anyhow!(
                    "Invalid copy-dir format: '{}'. Expected format: /local/path:/sandbox/path",
                    arg
                ));
            }
            Ok(CopyDir {
                local: PathBuf::from(parts[0]),
                remote: PathBuf::from(parts[1]),
            })
        })
        .collect::<Result<Vec<_>>>()?;

    // Parse CLI env vars and merge into provider config (CLI overrides config)
    let cli_env: HashMap<String, String> = env_vars
        .iter()
        .filter_map(|s| {
            let mut parts = s.splitn(2, '=');
            match (parts.next(), parts.next()) {
                (Some(k), Some(v)) if !k.is_empty() => Some((k.to_string(), v.to_string())),
                _ => {
                    tracing::warn!("Ignoring invalid --env value: '{}' (expected KEY=VALUE)", s);
                    None
                }
            }
        })
        .collect();

    if !cli_env.is_empty() {
        info!("CLI --env vars: {:?}", cli_env.keys().collect::<Vec<_>>());
        match &mut config.provider {
            ProviderConfig::Local(cfg) => cfg.env.extend(cli_env),
            ProviderConfig::Default(cfg) => cfg.env.extend(cli_env),
            ProviderConfig::Modal(cfg) => cfg.env.extend(cli_env),
        }
    }

    info!("Loaded configuration from {}", config_path.display());

    // Handle collect-only: only discovery needed, no provider setup
    if collect_only {
        eprint!("Discovering tests... ");
        let all_tests = discover_all_tests(&config.framework, &config.groups).await?;
        eprintln!(
            "found {} tests across {} groups",
            all_tests.len(),
            config.groups.len()
        );
        for group_name in config.groups.keys() {
            let group_tests: Vec<_> = all_tests
                .iter()
                .filter(|t| t.group == *group_name)
                .collect();
            if !group_tests.is_empty() {
                println!("\nGroup '{}':", group_name);
                for test in group_tests {
                    println!("  {}", test.id);
                }
            }
        }
        return Ok(());
    }

    // Convert copy_dirs to tuples once (used by Default and Modal providers)
    let copy_dir_tuples: Vec<(PathBuf, PathBuf)> = copy_dirs
        .iter()
        .map(|cd| (cd.local.clone(), cd.remote.clone()))
        .collect();

    // Resolve cwd once for threading through git/checkpoint calls
    let cwd = std::env::current_dir().context("failed to get current directory")?;

    // Phase 1+2: Discover tests and prepare provider (concurrently where possible)
    let exit_code = match &config.provider {
        ProviderConfig::Local(p_cfg) => {
            // Local provider is synchronous -- no concurrency benefit
            eprint!("Discovering tests... ");
            let all_tests = discover_all_tests(&config.framework, &config.groups).await?;
            eprintln!(
                "found {} tests across {} groups",
                all_tests.len(),
                config.groups.len()
            );
            if all_tests.is_empty() {
                info!("No tests to run");
                return Ok(());
            }
            dispatch_framework(
                &config,
                &config_filename,
                &run_id,
                &all_tests,
                LocalProvider::new(p_cfg.clone()),
                &copy_dirs,
                verbose,
                &tracer,
                show_estimated_cost,
                fail_fast,
                ci,
                record_history,
            )
            .await?
        }
        ProviderConfig::Default(p_cfg) => {
            let provider = DefaultProvider::from_config(p_cfg.clone());
            match run_remote_provider(
                &cwd,
                provider,
                &config,
                &copy_dir_tuples,
                &config_filename,
                &run_id,
                &copy_dirs,
                no_cache,
                verbose,
                &tracer,
                show_estimated_cost,
                fail_fast,
                config_path,
                ci,
                record_history,
            )
            .await?
            {
                Some(code) => code,
                None => return Ok(()),
            }
        }
        ProviderConfig::Modal(p_cfg) => {
            let provider = ModalProvider::from_config(p_cfg.clone());
            match run_remote_provider(
                &cwd,
                provider,
                &config,
                &copy_dir_tuples,
                &config_filename,
                &run_id,
                &copy_dirs,
                no_cache,
                verbose,
                &tracer,
                show_estimated_cost,
                fail_fast,
                config_path,
                ci,
                record_history,
            )
            .await?
            {
                Some(code) => code,
                None => return Ok(()),
            }
        }
    };

    // Write trace file if tracing was enabled
    let trace_path = config.report.output_dir.join("trace.json");
    if let Err(e) = tracer.write_to_file(&trace_path) {
        eprintln!("Warning: failed to write trace file: {}", e);
    } else if trace {
        eprintln!("Trace written to {}", trace_path.display());
    }

    if exit_code != 0 {
        std::process::exit(exit_code);
    }

    Ok(())
}

/// Run all tests with a single orchestrator call.
/// Returns the exit code (0 = success, 1 = failures/not run, 2 = flaky only).
#[allow(clippy::too_many_arguments)]
async fn run_all_tests<P, D>(
    config: &config::Config,
    config_filename: &str,
    run_id: &str,
    tests: &[TestRecord],
    provider: P,
    framework: D,
    copy_dirs: &[CopyDir],
    verbose: bool,
    tracer: &offload::trace::Tracer,
    show_estimated_cost: bool,
    fail_fast: bool,
    ci: bool,
    record_history: bool,
) -> Result<i32>
where
    P: offload::provider::SandboxProvider,
    D: TestFramework,
{
    // Convert CopyDir to tuples
    let copy_dir_tuples: Vec<(PathBuf, PathBuf)> = copy_dirs
        .iter()
        .map(|cd| (cd.local.clone(), cd.remote.clone()))
        .collect();

    // Pre-populate sandbox pool
    let mut env = provider.base_env();
    env.push((
        "OFFLOAD_ROOT".to_string(),
        config
            .offload
            .sandbox_project_root
            .clone()
            .context("sandbox_project_root not set")?,
    ));

    let sandbox_config = SandboxConfig {
        id: format!("offload-{}", uuid::Uuid::new_v4()),
        working_dir: config
            .offload
            .working_dir
            .as_ref()
            .map(|p| p.to_string_lossy().to_string()),
        env,
        copy_dirs: copy_dir_tuples.clone(),
    };

    let mut sandbox_pool = SandboxPool::new();
    let _pool_span = tracer.span(
        "sandbox_pool_create",
        "local",
        offload::trace::PID_LOCAL,
        offload::trace::TID_MAIN,
    );
    sandbox_pool
        .populate(config.offload.max_parallel, &provider, &sandbox_config, ci)
        .await
        .context("Failed to create sandboxes")?;
    drop(_pool_span);

    let orchestrator = Orchestrator::new(
        config.clone(),
        config_filename.to_string(),
        run_id.to_string(),
        framework,
        verbose,
        tracer.clone(),
        show_estimated_cost,
        fail_fast,
        ci,
        record_history,
    );

    let result = orchestrator.run_with_tests(tests, sandbox_pool).await?;

    Ok(result.exit_code())
}

async fn build_image(config_path: &Path, no_cache: bool) -> Result<()> {
    let config = config::load_config(config_path)
        .with_context(|| format!("Failed to load config from {}", config_path.display()))?;

    let cwd = std::env::current_dir().context("failed to get current directory")?;
    let tracer = offload::trace::Tracer::noop();
    let discovery_done = AtomicBool::new(true);
    let copy_dir_tuples: Vec<(PathBuf, PathBuf)> = Vec::new();

    let result = match &config.provider {
        ProviderConfig::Local(_) => {
            eprintln!("Local provider does not build images. Nothing to do.");
            return Ok(());
        }
        ProviderConfig::Default(p_cfg) => {
            let mut provider = DefaultProvider::from_config(p_cfg.clone());
            run_prepare(
                &mut provider,
                &cwd,
                &config,
                config_path,
                &copy_dir_tuples,
                no_cache,
                &tracer,
                &discovery_done,
            )
            .await
        }
        ProviderConfig::Modal(p_cfg) => {
            let mut provider = ModalProvider::from_config(p_cfg.clone());
            run_prepare(
                &mut provider,
                &cwd,
                &config,
                config_path,
                &copy_dir_tuples,
                no_cache,
                &tracer,
                &discovery_done,
            )
            .await
        }
    };

    match result {
        Ok(Some(image_id)) => {
            println!("{}", image_id);
            Ok(())
        }
        Ok(None) => {
            eprintln!("No image ID returned (no prepare_command configured?)");
            Ok(())
        }
        Err(e) => Err(e.context("Build failed")),
    }
}

async fn collect_tests(config_path: &Path, format: &str) -> Result<()> {
    let config = config::load_config(config_path)?;

    let all_tests = discover_all_tests(&config.framework, &config.groups).await?;

    match format {
        "json" => {
            let json = serde_json::to_string_pretty(&all_tests)?;
            println!("{}", json);
        }
        _ => {
            println!(
                "Discovered {} tests across {} groups:",
                all_tests.len(),
                config.groups.len()
            );
            for group_name in config.groups.keys() {
                let group_tests: Vec<_> = all_tests
                    .iter()
                    .filter(|t| t.group == *group_name)
                    .collect();
                if !group_tests.is_empty() {
                    println!("\nGroup '{}':", group_name);
                    for test in group_tests {
                        println!("  {}", test.id);
                    }
                }
            }
        }
    }

    Ok(())
}

fn validate_config(config_path: &Path) -> Result<()> {
    match config::load_config(config_path) {
        Ok(config) => {
            println!("Configuration is valid!");
            println!();
            println!("Settings:");
            println!("  Max parallel: {}", config.offload.max_parallel);
            println!("  Test timeout: {}s", config.offload.test_timeout_secs);

            let provider_name = match &config.provider {
                ProviderConfig::Local(_) => "local",
                ProviderConfig::Default(_) => "default",
                ProviderConfig::Modal(_) => "modal",
            };
            println!("  Provider: {}", provider_name);

            let framework_name = framework_type_name(&config.framework);
            println!("  Framework: {}", framework_name);

            if let Some(ref init_cmd) = config.offload.sandbox_init_cmd {
                println!("  Sandbox init cmd: {}", init_cmd);
            }

            println!();
            println!("Groups:");
            for (group_name, group_config) in &config.groups {
                println!(
                    "  {}: retry_count = {}",
                    group_name, group_config.retry_count
                );
            }

            Ok(())
        }
        Err(e) => {
            eprintln!("Configuration error: {}", e);
            std::process::exit(1);
        }
    }
}

fn init_config(provider: &str, framework: &str) -> Result<()> {
    let provider_config = match provider {
        "local" => ProviderConfig::Local(LocalProviderConfig {
            working_dir: Some(PathBuf::from(".")),
            ..Default::default()
        }),
        "default" => ProviderConfig::Default(DefaultProviderConfig {
            create_command: "./scripts/create-sandbox.sh".into(),
            exec_command: "./scripts/exec-sandbox.sh {sandbox_id} {command}".into(),
            destroy_command: "./scripts/destroy-sandbox.sh {sandbox_id}".into(),
            prepare_command: None,
            download_command: None,
            working_dir: None,
            timeout_secs: 3600,
            copy_dirs: vec![],
            env: HashMap::new(),
            cpu_cores: 1.0,
        }),
        _ => {
            eprintln!("Unknown provider: {}. Use: local, default", provider);
            std::process::exit(1);
        }
    };

    let framework_config = match framework {
        "pytest" => FrameworkConfig::Pytest(PytestFrameworkConfig {
            paths: None,
            command: "python -m pytest".into(),
            test_id_format: "{name}".into(),
            ..Default::default()
        }),
        "nextest" => FrameworkConfig::Cargo(CargoFrameworkConfig {
            test_id_format: "{classname} {name}".into(),
            ..Default::default()
        }),
        "default" => FrameworkConfig::Default(DefaultFrameworkConfig {
            discover_command: "echo test1 test2".into(),
            run_command: "echo Running {tests}".into(),
            test_id_format: "{name}".into(),
            result_file: None,
            working_dir: None,
        }),
        "vitest" => FrameworkConfig::Vitest(VitestFrameworkConfig {
            command: "npx vitest".into(),
            test_id_format: "{classname} > {name}".into(),
            ..Default::default()
        }),
        _ => {
            eprintln!(
                "Unknown framework: {}. Use: pytest, nextest, vitest, default",
                framework
            );
            std::process::exit(1);
        }
    };

    let config = Config {
        offload: OffloadConfig {
            max_parallel: 10,
            test_timeout_secs: 900,
            working_dir: None,
            sandbox_repo_root: Some("/app".to_string()),
            sandbox_project_root: None,
            sandbox_init_cmd: None,
            post_patch_cmd: None,
        },
        provider: provider_config,
        framework: framework_config,
        groups: HashMap::from([(
            "default".to_string(),
            GroupConfig {
                retry_count: 0,
                filters: String::new(),
                ..Default::default()
            },
        )]),
        report: ReportConfig::default(),
        checkpoint: None,
        history: None,
    };

    let toml_content = toml::to_string_pretty(&config)?;
    let output = format!("# offload configuration file\n\n{}", toml_content);

    let path = PathBuf::from("offload.toml");
    if path.exists() {
        eprintln!("offload.toml already exists. Remove it first or edit manually.");
        std::process::exit(1);
    }

    std::fs::write(&path, output)?;
    println!("Created offload.toml");
    println!();
    println!("Edit the configuration as needed, then run:");
    println!("  offload run");

    Ok(())
}

/// Configure git merge driver for history files.
///
/// Updates .gitattributes and configures the merge driver in .git/config.
fn setup_merge_driver() -> Result<()> {
    // Check if .git directory exists
    if !Path::new(".git").exists() {
        anyhow::bail!("Not a git repository (no .git directory found)");
    }

    // Update .gitattributes
    let gitattributes_line = "offload-history.jsonl merge=offload-history";
    let gitattributes_path = Path::new(".gitattributes");

    let needs_update = if gitattributes_path.exists() {
        let contents =
            std::fs::read_to_string(gitattributes_path).context("Failed to read .gitattributes")?;
        !contents.contains(gitattributes_line)
    } else {
        true
    };

    if needs_update {
        use std::io::Write;
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(gitattributes_path)
            .context("Failed to open .gitattributes")?;
        writeln!(file, "{}", gitattributes_line)?;
        println!("Updated .gitattributes");
    } else {
        println!(".gitattributes already configured");
    }

    // Configure git merge driver using git config
    let name_status = std::process::Command::new("git")
        .args([
            "config",
            "merge.offload-history.name",
            "Offload test history merger",
        ])
        .status()
        .context("Failed to run git config for merge driver name")?;

    if !name_status.success() {
        anyhow::bail!("Failed to configure merge driver name");
    }

    let driver_status = std::process::Command::new("git")
        .args([
            "config",
            "merge.offload-history.driver",
            "offload history merge %O %A %B",
        ])
        .status()
        .context("Failed to run git config for merge driver")?;

    if !driver_status.success() {
        anyhow::bail!("Failed to configure merge driver command");
    }

    println!("Git merge driver configured");
    Ok(())
}

/// Apply a git-format binary patch to the filesystem.
///
/// Reads the patch file, parses it with `diffy::PatchSet`, and applies each
/// file patch (create, delete, modify, rename, copy) to disk under
/// `project_root`.
fn apply_diff(patch_file: &Path, project_root: &Path) -> Result<()> {
    use diffy::patch_set::{FileOperation, ParseOptions, PatchSet};

    let patch_bytes = std::fs::read(patch_file)
        .with_context(|| format!("failed to read patch file: {}", patch_file.display()))?;

    let patches = PatchSet::parse_bytes(&patch_bytes, ParseOptions::gitdiff());

    for result in patches {
        let file_patch = result.context("failed to parse patch entry")?;
        // Rename and Copy paths come from `rename from`/`rename to` (or `copy from`/`copy to`)
        // git headers which have no `a/`/`b/` prefix, so strip_prefix must be skipped for them.
        let operation = file_patch.operation();
        let needs_strip = !operation.is_rename() && !operation.is_copy();
        let stripped;
        let operation = if needs_strip {
            stripped = operation.strip_prefix(1);
            &stripped
        } else {
            operation
        };
        let new_mode = file_patch.new_mode().copied();

        match operation {
            FileOperation::Create(path) => {
                let path_str = path_from_utf8(path, "path in create patch")?;
                let target = project_root.join(path_str);
                let content = apply_patch_to_base(&[], file_patch.patch(), path_str)?;
                write_patched_file(&target, &content, new_mode)?;
            }
            FileOperation::Delete(path) => {
                let path_str = path_from_utf8(path, "path in delete patch")?;
                let target = project_root.join(path_str);
                if target.exists() {
                    std::fs::remove_file(&target)
                        .with_context(|| format!("failed to delete file: {}", target.display()))?;
                }
            }
            FileOperation::Modify { original, modified } => {
                let orig_str = path_from_utf8(original, "original path in patch")?;
                let mod_str = path_from_utf8(modified, "modified path in patch")?;
                let source = project_root.join(orig_str);
                // If paths differ (rare for Modify, but possible), write to the modified path
                let target = project_root.join(mod_str);
                let content =
                    read_and_patch(&source, file_patch.patch(), orig_str, "modification")?;
                write_patched_file(&target, &content, new_mode)?;
            }
            FileOperation::Rename { from, to } => {
                let from_str = path_from_utf8(from, "rename-from path in patch")?;
                let to_str = path_from_utf8(to, "rename-to path in patch")?;
                let source = project_root.join(from_str);
                let target = project_root.join(to_str);
                let content = read_and_patch(&source, file_patch.patch(), from_str, "rename")?;
                write_patched_file(&target, &content, new_mode)?;
                std::fs::remove_file(&source).with_context(|| {
                    format!("failed to remove source of rename: {}", source.display())
                })?;
            }
            FileOperation::Copy { from, to } => {
                let from_str = path_from_utf8(from, "copy-from path in patch")?;
                let to_str = path_from_utf8(to, "copy-to path in patch")?;
                let source = project_root.join(from_str);
                let target = project_root.join(to_str);
                let content = read_and_patch(&source, file_patch.patch(), from_str, "copy")?;
                write_patched_file(&target, &content, new_mode)?;
            }
        }
    }

    Ok(())
}

/// Decode a path from a patch's raw bytes, attaching a description on failure.
fn path_from_utf8<'a>(bytes: &'a [u8], description: &str) -> Result<&'a str> {
    std::str::from_utf8(bytes).with_context(|| format!("non-UTF-8 {description}"))
}

/// Read `source` and apply `patch_kind` to its contents. `label` describes the
/// operation (e.g. "rename") for error messages.
fn read_and_patch(
    source: &Path,
    patch_kind: &diffy::patch_set::PatchKind<'_, [u8]>,
    diag_path: &str,
    label: &str,
) -> Result<Vec<u8>> {
    let existing = std::fs::read(source)
        .with_context(|| format!("failed to read file for {label}: {}", source.display()))?;
    apply_patch_to_base(&existing, patch_kind, diag_path)
}

/// Write `content` to `target`, creating parent directories and applying the
/// requested file mode.
fn write_patched_file(
    target: &Path,
    content: &[u8],
    mode: Option<diffy::patch_set::FileMode>,
) -> Result<()> {
    if let Some(parent) = target.parent() {
        std::fs::create_dir_all(parent).with_context(|| {
            format!(
                "failed to create parent directories for {}",
                target.display()
            )
        })?;
    }
    std::fs::write(target, content)
        .with_context(|| format!("failed to write file: {}", target.display()))?;
    set_file_mode(target, mode)
}

/// Apply a patch (text or binary) to a base byte slice, returning the result.
fn apply_patch_to_base(
    base: &[u8],
    patch_kind: &diffy::patch_set::PatchKind<'_, [u8]>,
    file_path: &str,
) -> Result<Vec<u8>> {
    match patch_kind {
        diffy::patch_set::PatchKind::Text(text_patch) => diffy::apply_bytes(base, text_patch)
            .with_context(|| format!("failed to apply text patch to {file_path}")),
        diffy::patch_set::PatchKind::Binary(binary_patch) => binary_patch
            .apply(base)
            .with_context(|| format!("failed to apply binary patch to {file_path}")),
    }
}

/// Set file permissions based on the patch's file mode (unix only).
#[cfg(unix)]
fn set_file_mode(path: &Path, mode: Option<diffy::patch_set::FileMode>) -> Result<()> {
    use std::os::unix::fs::PermissionsExt;

    if let Some(diffy::patch_set::FileMode::Executable) = mode {
        let perms = std::fs::Permissions::from_mode(0o755);
        std::fs::set_permissions(path, perms).with_context(|| {
            format!("failed to set executable permission on {}", path.display())
        })?;
    }
    Ok(())
}

/// Set file permissions based on the patch's file mode (non-unix: no-op).
#[cfg(not(unix))]
fn set_file_mode(_path: &Path, _mode: Option<diffy::patch_set::FileMode>) -> Result<()> {
    Ok(())
}

fn show_logs(
    config_path: &Path,
    failures: bool,
    errors: bool,
    test_ids: &[String],
    test_regex: Option<&str>,
) -> Result<()> {
    let config = config::load_config(config_path)
        .with_context(|| format!("Failed to load config from {}", config_path.display()))?;

    let junit_path = config.report.output_dir.join(&config.report.junit_file);

    if !junit_path.is_file() {
        eprintln!(
            "No test results found at {}. Run `offload run` first.",
            junit_path.display()
        );
        std::process::exit(1);
    }

    let re = test_regex
        .map(regex::Regex::new)
        .transpose()
        .context("Invalid --test-regex pattern")?;

    let xml_content = std::fs::read_to_string(&junit_path)
        .with_context(|| format!("Failed to read {}", junit_path.display()))?;

    let testsuites = offload::report::parse_all_testsuites_xml(&xml_content);

    // Collect all testcases, deduplicating by test name (keep the one with failure/error info if any)
    use std::collections::BTreeMap;
    let mut tests: BTreeMap<String, &offload::report::TestcaseXml> = BTreeMap::new();
    for suite in &testsuites {
        for tc in &suite.testcases {
            let existing = tests.get(tc.name.as_str());
            // Prefer the entry that has failure/error info over a passing one
            let dominated = match existing {
                None => true,
                Some(prev) => {
                    (tc.failure.is_some() || tc.error.is_some())
                        && prev.failure.is_none()
                        && prev.error.is_none()
                }
            };
            if dominated {
                tests.insert(tc.name.clone(), tc);
            }
        }
    }

    // Filter by status flags, then by test selection
    let filtered: Vec<(&String, &&offload::report::TestcaseXml)> = tests
        .iter()
        .filter(|(name, tc)| {
            // Status filter
            let status_ok = if failures && errors {
                tc.failure.is_some() || tc.error.is_some()
            } else if failures {
                tc.failure.is_some()
            } else if errors {
                tc.error.is_some()
            } else {
                true
            };
            if !status_ok {
                return false;
            }

            // Exact ID filter
            if !test_ids.is_empty() && !test_ids.iter().any(|id| id == name.as_str()) {
                return false;
            }

            // Regex filter
            if let Some(ref re) = re
                && !re.is_match(name)
            {
                return false;
            }

            true
        })
        .collect();

    if filtered.is_empty() {
        eprintln!("No matching test results found in {}", junit_path.display());
        return Ok(());
    }

    for (name, tc) in &filtered {
        let status = if tc.failure.is_some() {
            "FAILED"
        } else if tc.error.is_some() {
            "ERROR"
        } else {
            "PASSED"
        };

        println!("=== {} [{}] ===", name, status);

        if let Some(ref failure) = tc.failure {
            if let Some(ref msg) = failure.message {
                println!("{}", msg);
            }
            if !failure.content.is_empty() {
                println!("{}", failure.content);
            }
        }
        if let Some(ref error) = tc.error {
            if let Some(ref msg) = error.message {
                println!("{}", msg);
            }
            if !error.content.is_empty() {
                println!("{}", error.content);
            }
        }
        println!();
    }

    Ok(())
}
