mod buildctl_runner;
mod orchestrator_client;
mod retry;
mod sizing;

use anyhow::{Context, Result};
use clap::Parser;
use orchestrator_client::OrchestratorClient;
use std::path::PathBuf;
use tracing::{error, info, warn};

#[derive(Parser)]
#[command(name = "buildkit-launcher", about = "Ephemeral buildkit runner — ADR-039")]
struct Args {
    /// Workspace name (used for P95 lookup and runner labelling)
    #[arg(long, env = "BUILDKIT_WORKSPACE")]
    workspace: String,

    /// Build context directory (must be reachable from this host)
    #[arg(long)]
    context: PathBuf,

    /// Fully-qualified image reference to push after build
    #[arg(long)]
    image: String,

    /// Dockerfile path relative to context
    #[arg(long, default_value = "Dockerfile")]
    dockerfile: String,

    /// Cache-from reference (optional, e.g. type=registry,ref=zot.example/cache)
    #[arg(long)]
    cache_from: Option<String>,

    /// Cache-to reference (optional)
    #[arg(long)]
    cache_to: Option<String>,

    /// SSH key to use when connecting to the runner
    #[arg(long, env = "BUILDKIT_SSH_KEY")]
    ssh_key: PathBuf,

    /// Language hint for size defaults (rust, go, java, etc.)
    #[arg(long)]
    language: Option<String>,

    /// Golden image snapshot ID or name to boot the runner from
    #[arg(long, env = "BUILDKIT_RUNNER_IMAGE", default_value = "buildkit-runner-latest")]
    runner_image: String,

    /// Orchestrator base URL
    #[arg(long, env = "ORCHESTRATOR_URL", default_value = "http://localhost:9011")]
    orchestrator_url: String,
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env()
                .add_directive("buildkit_launcher=info".parse()?),
        )
        .init();

    let args = Args::parse();
    let client = OrchestratorClient::new(args.orchestrator_url.clone());

    // Resolve runner size
    let p95 = client.get_p95(&args.workspace).await.unwrap_or(None);
    let (p95_cpu, p95_mem) = p95
        .as_ref()
        .map(|s| (Some(s.cpu_p95), Some(s.memory_mb_p95)))
        .unwrap_or((None, None));

    let size = sizing::resolve(
        &args.context,
        p95_cpu,
        p95_mem,
        args.language.as_deref(),
    )
    .context("failed to resolve runner size")?;

    info!(
        cpu = size.cpu,
        memory_gb = size.memory_gb,
        workspace = %args.workspace,
        "resolved runner size"
    );

    let result = run_build(&args, &client, size.clone()).await;

    match result {
        Ok(()) => Ok(()),
        Err(e) if buildctl_runner::is_oom(&e) && retry::oom_retry_allowed(0) => {
            warn!(
                error = %e,
                limit = retry::MAX_OOM_RETRIES,
                "OOM detected — retrying at next size tier"
            );

            let retry_size = retry::next_size_tier(&size)
                .context("OOM on largest tier — no retry possible")?;

            info!(
                cpu = retry_size.cpu,
                memory_gb = retry_size.memory_gb,
                "retrying build at next tier"
            );

            run_build(&args, &client, retry_size).await
        }
        Err(e) => Err(e),
    }
}

async fn run_build(args: &Args, client: &OrchestratorClient, size: sizing::RunnerSize) -> Result<()> {
    let runner = client
        .spawn_runner(&size, &args.workspace, &args.runner_image)
        .await
        .context("failed to spawn runner")?;

    info!(
        lease_id = %runner.lease_id,
        ssh_host = %runner.ssh_host,
        ssh_port = runner.ssh_port,
        expires_at = %runner.expires_at,
        "runner spawned"
    );

    let destroy = |lease_id: String| {
        let client = client.clone();
        tokio::spawn(async move {
            if let Err(e) = client.destroy_runner(&lease_id).await {
                error!(lease_id = %lease_id, error = %e, "failed to destroy runner");
            }
        })
    };

    // Rsync build context to runner
    if let Err(e) = buildctl_runner::rsync_context(
        &runner.ssh_host,
        runner.ssh_port,
        &args.context,
        &args.ssh_key,
    )
    .await
    {
        destroy(runner.lease_id.clone());
        return Err(e.context("rsync build context failed"));
    }

    // Run buildctl
    let build_result = buildctl_runner::run_buildctl(
        &runner.ssh_host,
        runner.ssh_port,
        &args.ssh_key,
        &args.image,
        &args.dockerfile,
        args.cache_from.as_deref(),
        args.cache_to.as_deref(),
    )
    .await;

    // Always destroy the runner, even on build failure
    destroy(runner.lease_id.clone());

    let output = build_result.context("buildctl failed")?;

    info!(
        duration_secs = output.duration_secs,
        mem_peak_mb = output.mem_peak_mb,
        "build complete — recording metrics"
    );

    if let Err(e) = client
        .record_metrics(
            &args.workspace,
            output.cpu_secs,
            output.mem_peak_mb,
            output.duration_secs,
        )
        .await
    {
        warn!(error = %e, "failed to record build metrics (non-fatal)");
    }

    Ok(())
}
