//! The runner loop: claim a job, run it, ship its logs and artifacts, and report the result.

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use forge_ci::mask::{MaskStream, Masker};
use forge_ci::proto::{FinishRequest, JobAssignment, Outcome};
use tokio::sync::{Semaphore, mpsc};
use tokio::task::{JoinHandle, JoinSet};
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;

use crate::artifacts;
use crate::client::{ClientError, ForgeClient};
use crate::config::RunnerConfig;
use crate::executor::{ExecRequest, ExecResult, Executor};
use crate::logs::{self, PumpConfig};
use crate::update::{self, Installed, UpdateState};
use crate::workspace;

/// Log lines queued between the executor and the log shipper.
const LOG_CHANNEL: usize = 256;
/// How many times the runner tries to report a finished job.
const FINISH_ATTEMPTS: u32 = 5;
/// Largest single artifact the runner will upload. The forge enforces its own total cap.
const MAX_ARTIFACT_BYTES: u64 = 512 * 1024 * 1024;

/// Timing knobs. Production uses the defaults. Tests shorten them.
#[derive(Clone, Copy, Debug)]
pub struct RunOptions {
    pub heartbeat_every: Duration,
    pub claim_backoff: Duration,
    pub pump: PumpConfig,
    /// When set, replaces `timeout_minutes` from the job.
    pub timeout_override: Option<Duration>,
}

impl Default for RunOptions {
    fn default() -> Self {
        Self {
            heartbeat_every: Duration::from_secs(10),
            claim_backoff: Duration::from_secs(5),
            pump: PumpConfig::default(),
            timeout_override: None,
        }
    }
}

/// How [`Runner::run`] ended.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RunEnd {
    /// Shutdown was requested.
    Stopped,
    /// The binary was replaced. The caller exits with [`update::RESTART_EXIT_CODE`].
    Updated,
}

/// A runner bound to one forge and one executor.
#[derive(Clone)]
pub struct Runner {
    client: ForgeClient,
    config: Arc<RunnerConfig>,
    executor: Arc<dyn Executor>,
    options: RunOptions,
}

impl Runner {
    pub fn new(
        config: RunnerConfig,
        executor: Arc<dyn Executor>,
        options: RunOptions,
    ) -> Result<Self, ClientError> {
        let client = ForgeClient::new(&config.url)?;
        Ok(Self {
            client,
            config: Arc::new(config),
            executor,
            options,
        })
    }

    pub fn client(&self) -> &ForgeClient {
        &self.client
    }

    /// Claims and runs jobs, up to `concurrency` at a time, until `shutdown` is cancelled. Jobs
    /// still running at shutdown are cancelled and reported as cancelled.
    ///
    /// When the forge advertises a different binary and `auto_update` is on, no new jobs are
    /// claimed. Once the running jobs finish, the binary is replaced and the loop returns
    /// [`RunEnd::Updated`].
    pub async fn run(&self, shutdown: CancellationToken) -> RunEnd {
        let semaphore = Arc::new(Semaphore::new(self.config.concurrency));
        let mut jobs: JoinSet<()> = JoinSet::new();
        // Hashed once at startup. Without it, or with auto_update off, updates stay disabled.
        let exe = if self.config.auto_update {
            match update::locate_self() {
                Ok(found) => Some(found),
                Err(e) => {
                    tracing::warn!(error = %e, "auto update disabled, cannot read own binary");
                    None
                }
            }
        } else {
            None
        };
        let mut updates = UpdateState::new(
            self.config.auto_update,
            exe.as_ref().map(|(_, sha)| sha.clone()),
        );
        tracing::info!(
            name = %self.config.name,
            concurrency = self.config.concurrency,
            auto_update = updates.enabled(),
            "runner started"
        );
        let end = loop {
            if let Some((exe_path, _)) = &exe
                && let Some(sha) = updates.pending().map(str::to_owned)
            {
                if !jobs.is_empty() {
                    // Claiming is paused until the running jobs finish.
                    tokio::select! {
                        _ = shutdown.cancelled() => break RunEnd::Stopped,
                        Some(joined) = jobs.join_next() => {
                            if let Err(e) = joined {
                                tracing::error!(error = %e, "job task panicked");
                            }
                        }
                    }
                    continue;
                }
                match update::install(&self.client, &self.config.runner_secret, exe_path, &sha)
                    .await
                {
                    Ok(Installed::Replaced) => {
                        tracing::info!(sha = %sha, "binary updated, restarting");
                        break RunEnd::Updated;
                    }
                    Ok(Installed::Listed) => {
                        tracing::warn!(sha = %sha, "forge binary is on the bad list, not updating");
                        updates.give_up(&sha);
                    }
                    Err(e) => {
                        tracing::warn!(sha = %sha, error = %e, "update failed, keeping this binary");
                        updates.give_up(&sha);
                    }
                }
                continue;
            }

            let permit = tokio::select! {
                _ = shutdown.cancelled() => break RunEnd::Stopped,
                permit = semaphore.clone().acquire_owned() => match permit {
                    Ok(permit) => permit,
                    Err(_) => break RunEnd::Stopped,
                },
            };
            let claimed = tokio::select! {
                _ = shutdown.cancelled() => {
                    drop(permit);
                    break RunEnd::Stopped;
                }
                r = self.client.claim(&self.config.runner_secret, &self.config.labels) => r,
            };
            match claimed {
                Ok(reply) => {
                    if updates.observe(reply.binary_sha256.as_deref())
                        && let Some(sha) = updates.pending()
                    {
                        tracing::info!(sha = %sha, "forge advertises a new binary, updating when idle");
                    }
                    match reply.job {
                        Some(job) => {
                            tracing::info!(
                                job_id = %job.job_id,
                                job_name = %job.job_name,
                                repo = %job.repo,
                                "claimed job"
                            );
                            let runner = self.clone();
                            let cancel = shutdown.child_token();
                            jobs.spawn(async move {
                                let _permit = permit;
                                runner.process(job, cancel).await;
                            });
                        }
                        None => drop(permit),
                    }
                }
                Err(e) => {
                    drop(permit);
                    tracing::warn!(error = %e, "claim failed, backing off");
                    tokio::select! {
                        _ = shutdown.cancelled() => break RunEnd::Stopped,
                        _ = tokio::time::sleep(self.options.claim_backoff) => {}
                    }
                }
            }
            while let Some(joined) = jobs.try_join_next() {
                if let Err(e) = joined {
                    tracing::error!(error = %e, "job task panicked");
                }
            }
        };
        while let Some(joined) = jobs.join_next().await {
            if let Err(e) = joined {
                tracing::error!(error = %e, "job task panicked");
            }
        }
        tracing::info!("runner stopped");
        end
    }

    /// Runs one claimed job end to end and reports its result. Returns the final outcome.
    pub async fn process(&self, job: JobAssignment, shutdown: CancellationToken) -> Outcome {
        // Mask the job's secret values and its own token, both raw and base64 (done by Masker).
        let mut secret_values: Vec<String> = job.secrets.values().cloned().collect();
        secret_values.push(job.job_token.clone());
        let masker = Masker::new(secret_values);

        let (tx, rx) = mpsc::channel::<Vec<u8>>(LOG_CHANNEL);
        let pump = {
            let client = self.client.clone();
            let job_id = job.job_id.clone();
            let token = job.job_token.clone();
            let cfg = self.options.pump;
            tokio::spawn(async move {
                logs::pump(&client, &job_id, &token, MaskStream::new(masker), rx, cfg).await
            })
        };

        send_line(
            &tx,
            format!(
                "forge-runner {}: job {} on runner {}, image {}, commit {}",
                env!("CARGO_PKG_VERSION"),
                job.job_name,
                self.config.name,
                job.image,
                job.commit
            ),
        )
        .await;

        let timeout = self
            .options
            .timeout_override
            .unwrap_or_else(|| Duration::from_secs(60 * u64::from(job.timeout_minutes)));
        let job_cancel = shutdown.child_token();
        let timed_out = Arc::new(AtomicBool::new(false));
        let heartbeat = self.spawn_heartbeat(&job, job_cancel.clone());
        let timer = {
            let cancel = job_cancel.clone();
            let flag = timed_out.clone();
            tokio::spawn(async move {
                tokio::time::sleep(timeout).await;
                flag.store(true, Ordering::SeqCst);
                cancel.cancel();
            })
        };

        let mut workspace_dir = None;
        let prepared = workspace::prepare(&self.config.work_dir, &job, &self.config.url).await;
        let mut result = match prepared {
            Ok(dir) => {
                let result = self
                    .executor
                    .run(ExecRequest {
                        job: &job,
                        workspace: &dir,
                        logs: tx.clone(),
                        cancel: job_cancel.clone(),
                    })
                    .await;
                workspace_dir = Some(dir);
                result
            }
            Err(e) => ExecResult {
                outcome: Outcome::Error,
                exit_code: None,
                message: Some(format!("workspace: {e}")),
            },
        };

        // Artifacts only come from a successful run.
        let upload_error = match workspace_dir.as_deref() {
            Some(dir) if result.outcome == Outcome::Success && !job.artifacts.is_empty() => {
                self.upload_artifacts(&job, dir).await.err()
            }
            _ => None,
        };
        if let Some(message) = upload_error {
            result = ExecResult {
                outcome: Outcome::Error,
                exit_code: result.exit_code,
                message: Some(message),
            };
        }

        heartbeat.abort();
        timer.abort();
        if timed_out.load(Ordering::SeqCst) && result.outcome == Outcome::Cancelled {
            result = ExecResult {
                outcome: Outcome::Failure,
                exit_code: None,
                message: Some(format!("timed out after {} minutes", job.timeout_minutes)),
            };
        }

        if let Some(message) = &result.message {
            send_line(&tx, format!("runner: {message}")).await;
        }
        drop(tx);
        match pump.await {
            Ok(Ok(bytes)) => tracing::debug!(job_id = %job.job_id, bytes, "logs shipped"),
            Ok(Err(e)) => tracing::warn!(job_id = %job.job_id, error = %e, "log shipping failed"),
            Err(e) => tracing::error!(job_id = %job.job_id, error = %e, "log task panicked"),
        }

        self.finish(&job, &result).await;
        if let Some(dir) = workspace_dir
            && let Err(e) = workspace::cleanup(&dir).await
        {
            tracing::warn!(job_id = %job.job_id, error = %e, "workspace cleanup failed");
        }
        tracing::info!(job_id = %job.job_id, outcome = ?result.outcome, "job finished");
        result.outcome
    }

    fn spawn_heartbeat(&self, job: &JobAssignment, cancel: CancellationToken) -> JoinHandle<()> {
        let client = self.client.clone();
        let job_id = job.job_id.clone();
        let token = job.job_token.clone();
        let every = self.options.heartbeat_every;
        tokio::spawn(async move {
            let mut ticker = tokio::time::interval_at(Instant::now() + every, every);
            loop {
                ticker.tick().await;
                match client.heartbeat(&job_id, &token).await {
                    Ok(beat) if beat.cancel => {
                        tracing::info!(job_id = %job_id, "forge requested cancel");
                        cancel.cancel();
                        return;
                    }
                    Ok(_) => {}
                    Err(e) => tracing::warn!(job_id = %job_id, error = %e, "heartbeat failed"),
                }
            }
        })
    }

    /// Collects the job's artifacts and uploads them. The error is a message for the job record.
    async fn upload_artifacts(
        &self,
        job: &JobAssignment,
        dir: &std::path::Path,
    ) -> Result<(), String> {
        let found = artifacts::collect(dir, &job.artifacts, MAX_ARTIFACT_BYTES)
            .map_err(|e| format!("artifacts: {e}"))?;
        for artifact in &found {
            self.client
                .upload_artifact(&job.job_id, &job.job_token, &artifact.name, &artifact.path)
                .await
                .map_err(|e| format!("artifact {}: {e}", artifact.name))?;
        }
        tracing::info!(job_id = %job.job_id, count = found.len(), "artifacts uploaded");
        Ok(())
    }

    async fn finish(&self, job: &JobAssignment, result: &ExecResult) {
        let req = FinishRequest {
            outcome: result.outcome,
            exit_code: result.exit_code,
            message: result.message.clone(),
        };
        for attempt in 1..=FINISH_ATTEMPTS {
            match self.client.finish(&job.job_id, &job.job_token, &req).await {
                Ok(()) => return,
                Err(e) => {
                    tracing::warn!(job_id = %job.job_id, attempt, error = %e, "finish failed");
                    tokio::time::sleep(Duration::from_secs(2 * u64::from(attempt))).await;
                }
            }
        }
        tracing::error!(
            job_id = %job.job_id,
            "could not report the job result; the forge will reap it after heartbeats stop"
        );
    }
}

async fn send_line(tx: &mpsc::Sender<Vec<u8>>, line: String) {
    let mut bytes = line.into_bytes();
    bytes.push(b'\n');
    let _ = tx.send(bytes).await;
}