use anyhow::{Context, bail};
use ci_protocol::runner::{
    ArtifactCompleteRequest, ArtifactInitRequest, ClaimRequest, HeartbeatRequest, JobFinishRequest,
    JobFinishState, JobPayload, LogChunk, LogChunksRequest, RegisterRequest, StepReport,
    StepReportState,
};
use runner::client::RunnerClient;
use runner::executor::{JobOutcome, StepHooks, execute_job};
use runner::state::{credential_path, load_credential, store_credential};
use tracing::{info, warn};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    init_tracing();

    let mut cfg = config::runner::RunnerConfig::load().context("loading runner configuration")?;
    for (label, platform) in &cfg.actions_platforms {
        if (cfg.docker_enabled || platform == "-self-hosted") && !cfg.labels.contains(label) {
            cfg.labels.push(label.clone());
        }
    }
    if cfg.registration_token.trim().is_empty() && !credential_path(&cfg).exists() {
        bail!("runner needs APP_RUNNER__REGISTRATION_TOKEN or an existing credential file");
    }

    std::fs::create_dir_all(&cfg.workspace_root)
        .with_context(|| format!("creating workspace root {}", cfg.workspace_root))?;

    let client = RunnerClient::new(&cfg)?;
    let actions_available = runner::actions::available(&cfg);
    let protocol_versions = if actions_available {
        vec![2, 1]
    } else {
        vec![1]
    };
    let mut capabilities: Vec<String> = if actions_available {
        vec!["actions-v2".into()]
    } else {
        Vec::new()
    };
    if cfg.docker_enabled {
        capabilities.push("docker".into());
    }
    capabilities.extend(
        cfg.actions_groups
            .iter()
            .map(|group| format!("runner-group:{group}")),
    );
    let credential = match load_credential(&cfg)? {
        Some(credential) => credential,
        None => {
            info!(server = %cfg.server_url, "registering runner");
            let response = client
                .register(&RegisterRequest {
                    token: cfg.registration_token.clone(),
                    name: cfg.resolved_name(),
                    labels: cfg.labels.clone(),
                    capacity: cfg.capacity,
                    version: Some(env!("CARGO_PKG_VERSION").to_string()),
                    protocol_versions: protocol_versions.clone(),
                    capabilities: capabilities.clone(),
                })
                .await
                .context("registering runner")?;
            if !protocol_versions.contains(&response.protocol_version) {
                bail!(
                    "runner protocol mismatch: server {} / runner {}",
                    response.protocol_version,
                    ci_protocol::PROTOCOL_VERSION
                );
            }
            store_credential(&cfg, &response.credential)?;
            info!(runner_id = %response.runner_id, "runner registered");
            response.credential
        }
    };
    let client = client.with_credential(credential);

    let active = std::sync::Arc::new(std::sync::Mutex::new(
        std::collections::HashSet::<uuid::Uuid>::new(),
    ));
    let heartbeat_client = client.clone();
    let heartbeat_active = active.clone();
    let heartbeat_interval = cfg.heartbeat_interval_secs.max(1);
    let heartbeat_task = tokio::spawn(async move {
        let mut ticker = tokio::time::interval(std::time::Duration::from_secs(heartbeat_interval));
        ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
        loop {
            ticker.tick().await;
            let running: Vec<uuid::Uuid> = heartbeat_active
                .lock()
                .map(|set| set.iter().copied().collect())
                .unwrap_or_default();
            match heartbeat_client
                .heartbeat(&HeartbeatRequest {
                    running_job_ids: running,
                    protocol_versions: protocol_versions.clone(),
                    capabilities: capabilities.clone(),
                })
                .await
            {
                Ok(response) => {
                    if !response.cancel.is_empty() {
                        warn!(jobs = ?response.cancel, "server requested job cancellation");
                        runner::executor::request_cancel(&response.cancel);
                    }
                }
                Err(error) => warn!(%error, "heartbeat failed"),
            }
        }
    });

    info!(server = %cfg.server_url, "runner polling for jobs");
    let mut jobs = tokio::task::JoinSet::new();
    loop {
        while jobs.len() >= cfg.capacity.max(1) as usize {
            if let Some(result) = jobs.join_next().await {
                if let Err(error) = result {
                    warn!(%error,"runner job task failed");
                }
            }
        }
        let payload = match client.claim(&ClaimRequest { wait_secs: None }).await {
            Ok(response) => match response.job {
                Some(job) => job,
                None => continue,
            },
            Err(error) => {
                warn!(%error, "claim failed; backing off");
                tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                continue;
            }
        };
        info!(job_id = %payload.job_id, job = %payload.job_name, "job received");
        if let Ok(mut set) = active.lock() {
            set.insert(payload.job_id);
        }
        let job_client = client.clone();
        let job_config = cfg.clone();
        let job_active = active.clone();
        jobs.spawn(async move {
            execute_claimed_job(&job_client, &job_config, &payload).await;
            if let Ok(mut set) = job_active.lock() {
                set.remove(&payload.job_id);
            }
            runner::executor::clear_cancel(&payload.job_id);
        });
        if cfg.once {
            info!("runner configured for a single job; exiting");
            break;
        }
    }
    while let Some(result) = jobs.join_next().await {
        if let Err(error) = result {
            warn!(%error,"runner job task failed");
        }
    }

    heartbeat_task.abort();
    Ok(())
}

async fn execute_claimed_job(
    client: &RunnerClient,
    cfg: &config::runner::RunnerConfig,
    payload: &JobPayload,
) {
    let finish = if payload.actions.is_some() {
        runner::actions::execute(client, cfg, payload).await
    } else {
        match run_job(client, cfg, payload).await {
            JobOutcome::Success { exit_code } => JobFinishRequest {
                state: JobFinishState::Success,
                exit_code,
                ..Default::default()
            },
            JobOutcome::Failure { exit_code, error } => JobFinishRequest {
                state: JobFinishState::Failure,
                exit_code,
                error,
                ..Default::default()
            },
            JobOutcome::Cancelled => JobFinishRequest {
                state: JobFinishState::Cancelled,
                error: Some("cancelled".into()),
                ..Default::default()
            },
            JobOutcome::TimedOut => JobFinishRequest {
                state: JobFinishState::TimedOut,
                error: Some("job timed out".into()),
                ..Default::default()
            },
        }
    };
    for attempt in 0..3 {
        match client.finish(&payload.job_id, &finish).await {
            Ok(_) => return,
            Err(error) => {
                warn!(%error,attempt,"job finish report failed");
                tokio::time::sleep(std::time::Duration::from_secs(1)).await;
            }
        }
    }
}

async fn run_job(
    client: &RunnerClient,
    cfg: &config::runner::RunnerConfig,
    payload: &JobPayload,
) -> JobOutcome {
    if let Err(error) = client.start(&payload.job_id).await {
        warn!(%error, "job start report failed");
    }

    let seq = std::sync::Arc::new(std::sync::atomic::AtomicI32::new(0));
    let log_client = client.clone();
    let log_job = payload.job_id;
    // Mask secret values runner-side before they ever leave the machine; the
    // scheduler redacts again on ingest.
    let log_secrets: Vec<String> = payload.secrets.values().cloned().collect();
    let on_log = {
        let seq = std::sync::Arc::clone(&seq);
        let secrets = log_secrets.clone();
        move |text: String| {
            let client = log_client.clone();
            let job_id = log_job;
            let text = mask_secrets(text, &secrets);
            let num = seq.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            tokio::spawn(async move {
                let _ = client
                    .upload_logs(
                        job_id,
                        &LogChunksRequest {
                            chunks: vec![LogChunk { seq: num, text }],
                        },
                    )
                    .await;
            });
        }
    };

    let step_client = client.clone();
    let step_job = payload.job_id;
    let on_status =
        move |index: i32, signal: runner::executor::StepSignal, exit_code: Option<i32>| {
            let client = step_client.clone();
            let job_id = step_job;
            tokio::spawn(async move {
                let report = StepReport {
                    state: match signal {
                        runner::executor::StepSignal::Started => StepReportState::Started,
                        runner::executor::StepSignal::Finished => StepReportState::Finished,
                    },
                    exit_code,
                    name: None,
                    shell: None,
                    working_directory: None,
                    script: None,
                    timeout_secs: None,
                    started_at: (signal == runner::executor::StepSignal::Started)
                        .then(time::OffsetDateTime::now_utc),
                    finished_at: (signal == runner::executor::StepSignal::Finished)
                        .then(time::OffsetDateTime::now_utc),
                    ..Default::default()
                };
                let _ = client.report_step(job_id, index, &report).await;
            });
        };

    let artifact_client = client.clone();
    let artifact_job = payload.job_id;
    let on_artifact = move |name: String, path: String, size: i64, bytes: Vec<u8>| {
        let client = artifact_client.clone();
        let job_id = artifact_job;
        tokio::spawn(async move {
            let sha256 = runner::executor::sha256_hex(&bytes);
            let init = match client
                .init_artifact(
                    job_id,
                    &ArtifactInitRequest {
                        name: name.clone(),
                        size,
                        sha256: Some(sha256.clone()),
                    },
                )
                .await
            {
                Ok(init) => init,
                Err(error) => {
                    warn!(%error, %name, "artifact init failed");
                    return;
                }
            };
            if let Err(error) = runner::executor::upload_file(&init.upload_url, bytes).await {
                warn!(%error, %name, "artifact upload failed");
                return;
            }
            let _ = path;
            if let Err(error) = client
                .complete_artifact(
                    job_id,
                    init.artifact_id,
                    &ArtifactCompleteRequest {
                        size,
                        sha256: Some(sha256),
                    },
                )
                .await
            {
                warn!(%error, %name, "artifact completion failed");
            }
        });
    };

    let renew_client = client.clone();
    let renew_job = payload.job_id;
    let renew = move || {
        let client = renew_client.clone();
        let job_id = renew_job;
        tokio::spawn(async move {
            let _ = client.renew(&job_id).await;
        });
    };

    let hooks = StepHooks {
        on_status: std::sync::Arc::new(on_status),
        on_log: std::sync::Arc::new(on_log),
    };
    execute_job(
        cfg,
        payload,
        hooks,
        std::sync::Arc::new(on_artifact),
        std::sync::Arc::new(renew),
    )
    .await
}

fn mask_secrets(mut text: String, secrets: &[String]) -> String {
    for secret in secrets {
        if secret.len() >= 4 && text.contains(secret.as_str()) {
            text = text.replace(secret.as_str(), "***");
        }
    }
    text
}

fn init_tracing() {
    use tracing_subscriber::EnvFilter;
    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
    tracing_subscriber::fmt().with_env_filter(filter).init();
}