//! `flow runner` (docs/m0.md §5): claim a job from OxygenHub, fetch the

//! repository into a cache with the job's token, write the state into a

//! workspace, run it in its tier, stream the log, report the receipt.

//! One job at a time.

use crate::home::{Home, RunnerConfig, RunnerKey};

use crate::{microvm, workspace};

use anyhow::{bail, Context, Result};

use flow_core::api::{Claim, ClaimRequest, Decline, Finish, LogBatch, Record, Register, Registered, RunnerInfo, MAX_RECORD_TEXT};

use flow_core::{scope_tree_hash, Receipt};

use oxy_core::hash::ChangeHash;

use oxy_core::repo::Repo;

use std::path::{Path, PathBuf};

use std::time::{Duration, Instant};

/// A request the hub refused (4xx). On a job's endpoints, 409 means the

/// job no longer runs: it was cancelled, or given up.

#[derive(Debug)]

pub struct Refused {

    pub code: u16,

    pub text: String,

}

impl std::fmt::Display for Refused {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        write!(f, "HTTP {}: {}", self.code, self.text)

    }

}

impl std::error::Error for Refused {}

/// Whether `e` is the hub saying the job no longer runs.

fn job_gone(e: &anyhow::Error) -> bool {

    e.downcast_ref::<Refused>().is_some_and(|r| r.code == 409)

}

/// What a job ends with when its run was cancelled: nothing is reported,

/// as the hub has ended it already.

const CANCELLED: &str = "the run was cancelled";

fn cancelled() -> Finish {

    Finish {

        exit_code: None,

        duration_ms: 0,

        receipt: None,

        signed_receipt: None,

        error: Some(CANCELLED.into()),

        peak_disk_mib: None,

        peak_memory_mib: None,

    }

}

pub struct Hub {

    base: String,

    token: String,

    agent: ureq::Agent,

}

impl Hub {

    pub fn new(cfg: &RunnerConfig) -> Result<Hub> {

        let tls = oxy_core::protocol::tls_client_config(cfg.ca.as_deref()).map_err(|e| anyhow::anyhow!("{e}"))?;

        let agent = ureq::AgentBuilder::new().tls_config(tls).timeout(Duration::from_secs(60)).build();

        Ok(Hub { base: cfg.hub.trim_end_matches('/').to_string(), token: cfg.token.clone(), agent })

    }

    /// POST to the hub. A server error (5xx) means nothing was committed,

    /// so it is tried again after 1, 2 and 4 seconds before giving up: a

    /// busy hub costs a job a few seconds, not its log lines or its claim.

    fn post(&self, path: &str, body: &impl serde::Serialize) -> Result<Option<ureq::Response>> {

        self.post_with(path, body, &RETRY_DELAYS)

    }

    fn post_with(&self, path: &str, body: &impl serde::Serialize, delays: &[Duration]) -> Result<Option<ureq::Response>> {

        let url = format!("{}/api/v1{path}", self.base);

        let r = retrying(&url, delays, false, || {

            self.agent.post(&url).set("Authorization", &format!("Bearer {}", self.token)).send_json(body).map_err(Box::new)

        })?;

        Ok(if r.status() == 204 { None } else { Some(r) })

    }

    /// POST what a job leaves behind (oxygenflow#49): through a hub that is

    /// away for a few minutes, as while a deploy restarts it, so a job's

    /// declines and uploads outlast it.

    fn post_patiently(&self, path: &str, body: &impl serde::Serialize) -> Result<Option<ureq::Response>> {

        let url = format!("{}/api/v1{path}", self.base);

        let r = retrying(&url, &PATIENT_DELAYS, true, || {

            self.agent.post(&url).set("Authorization", &format!("Bearer {}", self.token)).send_json(body).map_err(Box::new)

        })?;

        Ok(if r.status() == 204 { None } else { Some(r) })

    }

    pub fn claim(&self, req: &ClaimRequest) -> Result<Option<Claim>> {

        match self.post("/flow/claim", req)? {

            Some(r) => Ok(Some(r.into_json()?)),

            None => Ok(None),

        }

    }

    /// Hand job `job` back, saying why this runner will not run it.

    pub fn decline(&self, job: i64, reason: &str) -> Result<()> {

        self.post_patiently(&format!("/flow/jobs/{job}/decline"), &Decline { reason: reason.into() })?;

        Ok(())

    }

    /// Hand back a job this runner holds as it leaves (oxygenflow#76): the

    /// hub queues it again ahead of new work for another runner.

    pub fn release(&self, job: i64, reason: &str) -> Result<()> {

        self.post_patiently(&format!("/flow/jobs/{job}/release"), &Decline { reason: reason.into() })?;

        Ok(())

    }

    /// Say this runner is there while it claims nothing: its slots are full,

    /// or it is draining (oxygenhub#243). Once, no retries: the next comes

    /// soon. A hub without the endpoint (404) is not an error.

    /// A ten-minute certificate for `identity`, asked with a job's token

    /// (keyless, oxygenhub#241), and the identity of the hub's key that

    /// issued it, from `/.well-known/oxyhub-key`.

    pub fn certify(

        &self,

        job_token: &str,

        identity: &oxy_core::identity::Identity,

        proof: &oxy_core::identity::Signature,

    ) -> Result<(oxy_core::cert::SignedCertificate, oxy_core::identity::Identity)> {

        use oxy_core::dsse::{base64, base64_decode};

        let body = serde_json::json!({

            "identity": base64(&identity.encode()),

            "proof": base64(&oxy_core::cbor::encode(&proof.to_value())),

            "scope": "seal",

        });

        let url = format!("{}/api/v1/keyless/certify", self.base);

        let r = self

            .agent

            .post(&url)

            .set("Authorization", &format!("Bearer {job_token}"))

            .send_json(body)

            .map_err(|e| anyhow::anyhow!("asking the hub to certify the key: {e}"))?;

        let v: serde_json::Value = r.into_json()?;

        let cert = base64_decode(v["certificate"].as_str().unwrap_or(""))

            .context("the hub's answer has no certificate")

            .and_then(|b| oxy_core::cert::SignedCertificate::decode(&b).map_err(|e| anyhow::anyhow!("{e}")))?;

        let mut bytes = Vec::new();

        std::io::Read::read_to_end(

            &mut self.agent.get(&format!("{}/.well-known/oxyhub-key", self.base)).call()?.into_reader(),

            &mut bytes,

        )?;

        let issuer = oxy_core::identity::Identity::decode(&bytes)?;

        if issuer.hash() != cert.cert.issuer || cert.verify(&issuer).is_err() {

            bail!("the certificate is not signed by the hub's published key");

        }

        if cert.cert.key_id() != identity.hash() {

            bail!("the hub certified another key");

        }

        Ok((cert, issuer))

    }

    pub fn heartbeat(&self) -> Result<()> {

        match self.post_with("/flow/runners/heartbeat", &serde_json::json!({}), &[]) {

            Ok(_) => Ok(()),

            Err(e) if format!("{e:#}").contains("404") => Ok(()),

            Err(e) => Err(e),

        }

    }

    pub fn log(&self, job: i64, records: Vec<Record>) -> Result<()> {

        self.post(&format!("/flow/jobs/{job}/log"), &LogBatch { records })?;

        Ok(())

    }

    /// Upload artifact `name` of job `job` (spec §5.2); the hub keeps it as

    /// a blob of the job's repository and answers with its ObjectRef.

    pub fn artifact(&self, job: i64, name: &str, bytes: &[u8]) -> Result<String> {

        let url = format!("{}/api/v1/flow/jobs/{job}/artifacts/{name}", self.base);

        let r = retrying(&url, &PATIENT_DELAYS, true, || {

            self.agent

                .post(&url)

                .timeout(Duration::from_secs(600))

                .set("Authorization", &format!("Bearer {}", self.token))

                .set("Content-Type", "application/octet-stream")

                .send_bytes(bytes)

                .map_err(Box::new)

        })?;

        let v: serde_json::Value = r.into_json()?;

        v["blob"].as_str().map(String::from).context("the hub named no blob")

    }

    /// Input `index` of job `job` (oxygenflow#58): the bytes of the

    /// producing job's artifact, which the caller checks against the blob the

    /// claim names.

    pub fn input(&self, job: i64, index: usize) -> Result<Vec<u8>> {

        let url = format!("{}/api/v1/flow/jobs/{job}/inputs/{index}", self.base);

        let r = retrying(&url, &PATIENT_DELAYS, false, || {

            self.agent

                .get(&url)

                .timeout(Duration::from_secs(600))

                .set("Authorization", &format!("Bearer {}", self.token))

                .call()

                .map_err(Box::new)

        })?;

        let mut bytes = Vec::new();

        let mut reader = r.into_reader();

        std::io::Read::read_to_end(&mut std::io::Read::take(&mut reader, MAX_INPUT_BYTES as u64 + 1), &mut bytes)?;

        if bytes.len() > MAX_INPUT_BYTES {

            bail!("larger than {} MiB", MAX_INPUT_BYTES >> 20);

        }

        Ok(bytes)

    }

    /// Upload file `path` of report `name` of job `job` (oxygenhub#127);

    /// the hub keeps it as a blob and answers with its ObjectRef. The

    /// report's manifest follows through [`Hub::artifact`].

    pub fn report_file(&self, job: i64, name: &str, path: &str, bytes: &[u8]) -> Result<String> {

        let url = format!("{}/api/v1/flow/jobs/{job}/reports/{name}/files", self.base);

        let r = retrying(&url, &PATIENT_DELAYS, true, || {

            self.agent

                .post(&url)

                .timeout(Duration::from_secs(600))

                .set("Authorization", &format!("Bearer {}", self.token))

                .set("Content-Type", "application/octet-stream")

                .set("X-Flow-Report-Path", path)

                .send_bytes(bytes)

                .map_err(Box::new)

        })?;

        let v: serde_json::Value = r.into_json()?;

        v["blob"].as_str().map(String::from).context("the hub named no blob")

    }

    /// Report a job's end until the hub has it (oxygenhub#71): while the hub

    /// cannot be reached or fails (5xx), try again after `first_wait`,

    /// doubling to a minute, for up to `patience`. A refusal (4xx: the job was ended

    /// already, say) is final. A result lost here leaves the job running

    /// on the hub until its time limit.

    pub fn finish_until_accepted(&self, job: i64, f: &Finish, patience: Duration, first_wait: Duration) -> Result<()> {

        let url = format!("{}/api/v1/flow/jobs/{job}/finish", self.base);

        let deadline = std::time::Instant::now() + patience;

        let mut wait = first_wait;

        loop {

            let why = match self.agent.post(&url).set("Authorization", &format!("Bearer {}", self.token)).send_json(f) {

                Ok(_) => return Ok(()),

                Err(ureq::Error::Status(code, r)) if code < 500 => {

                    bail!("{url}: HTTP {code}: {}", r.into_string().unwrap_or_default())

                }

                Err(ureq::Error::Status(code, _)) => format!("HTTP {code}"),

                Err(e) => e.to_string(),

            };

            if std::time::Instant::now() + wait > deadline {

                bail!("{url}: gave up after {} minutes: {why}", patience.as_secs() / 60);

            }

            eprintln!("job {job}: the hub did not take the result ({why}); trying again in {} s", wait.as_secs());

            std::thread::sleep(wait);

            wait = (wait * 2).min(Duration::from_secs(60));

        }

    }

}

/// How long a runner keeps trying to report a job's end.

const FINISH_PATIENCE: Duration = Duration::from_secs(2 * 3600);

/// Where a job's result waits until the hub has it, so a runner that

/// restarts still reports it.

fn pending_path(home: &Home, job: i64) -> PathBuf {

    home.root.join("pending").join(format!("{job}.json"))

}

/// Report `f`, keeping it on disk until the hub has it.

fn report_finish(home: &Home, hub: &Hub, job: i64, f: &Finish) {

    let path = pending_path(home, job);

    let kept = std::fs::create_dir_all(path.parent().unwrap())

        .and_then(|_| std::fs::write(&path, serde_json::to_vec(f).unwrap_or_default()))

        .is_ok();

    match hub.finish_until_accepted(job, f, FINISH_PATIENCE, Duration::from_secs(5)) {

        Ok(()) => {}

        Err(e) => eprintln!("job {job}: reporting failed: {e:#}"),

    }

    if kept {

        let _ = std::fs::remove_file(&path);

    }

}

/// Results a previous run of the runner could not report, to send again.

fn report_pending(home: &Home) -> Vec<(i64, Finish)> {

    let Ok(dir) = std::fs::read_dir(home.root.join("pending")) else {

        return Vec::new();

    };

    dir.filter_map(|e| {

        let p = e.ok()?.path();

        let job: i64 = p.file_stem()?.to_str()?.parse().ok()?;

        let f: Finish = serde_json::from_slice(&std::fs::read(&p).ok()?).ok()?;

        Some((job, f))

    })

    .collect()

}

/// `flow runner register`: exchange a one-time registration token from

/// the hub for this runner's credential, binding the runner's signing key,

/// and keep hub, credential and name in runner.toml.

pub fn register(home: &Home, hub: &str, token: &str, name: &str, ca: Option<PathBuf>, labels: Vec<String>) -> Result<Registered> {

    if home.config.runner.is_some() && !home.root.join(crate::home::RUNNER_TOML).is_file() {

        bail!("flow.toml already has a [runner] section; remove it to register this runner instead");

    }

    let key = home.runner_key(name)?;

    let hub = hub.trim_end_matches('/').to_string();

    let tls = oxy_core::protocol::tls_client_config(ca.as_deref()).map_err(|e| anyhow::anyhow!("{e}"))?;

    let agent = ureq::AgentBuilder::new().tls_config(tls).timeout(Duration::from_secs(60)).build();

    let req = Register {

        token: token.trim().to_string(),

        name: name.to_string(),

        identity: hex::encode(key.identity.encode()),

        info: RunnerInfo { version: env!("OXY_VERSION_LONG").into(), ..RunnerInfo::here(labels.clone(), 1, None, None) },

    };

    let url = format!("{hub}/api/v1/flow/runners/register");

    let got: Registered = match agent.post(&url).send_json(&req) {

        Ok(r) => r.into_json()?,

        Err(ureq::Error::Status(code, r)) => bail!("{url}: HTTP {code}: {}", r.into_string().unwrap_or_default()),

        Err(e) => bail!("{url}: {e}"),

    };

    let mut text = format!(

        "# Written by `flow runner register`; the credential is secret.\nhub = {hub:?}\ntoken = {:?}\nname = {:?}\n",

        got.credential, got.name

    );

    if let Some(ca) = &ca {

        text.push_str(&format!("ca = {:?}\n", ca.display().to_string()));

    }

    if !labels.is_empty() {

        text.push_str(&format!("labels = {labels:?}\n"));

    }

    let path = home.root.join(crate::home::RUNNER_TOML);

    {

        use std::io::Write;

        use std::os::unix::fs::OpenOptionsExt;

        let tmp = path.with_extension("toml.tmp");

        let mut f = std::fs::OpenOptions::new().write(true).create(true).truncate(true).mode(0o600).open(&tmp)?;

        f.write_all(text.as_bytes())?;

        f.sync_all()?;

        std::fs::rename(&tmp, &path)?;

    }

    Ok(got)

}

/// Waits before trying a hub request again after a server error.

const RETRY_DELAYS: [Duration; 3] = [Duration::from_secs(1), Duration::from_secs(2), Duration::from_secs(4)];

/// Waits for what a job leaves behind (oxygenflow#49): three minutes in

/// all, longer than a hub takes to restart.

/// Check an input's bytes against the blob the producing job's receipt

/// names, then write them at its path under `tree` (the workspace's `/src`).

fn place_input(tree: &Path, i: &flow_core::api::ArtifactInput, bytes: &[u8]) -> Result<[u8; 32]> {

    let Some(rel) = i.path.strip_prefix("/src/").filter(|r| r.split('/').all(|p| !p.is_empty() && p != "." && p != ".." && p != ".oxy"))

    else {

        bail!("path {:?} is not a clean path under /src", i.path);

    };

    let blob = Repo::blob_ref(bytes);

    if blob.to_hex() != i.blob {

        bail!("it is not what {}'s receipt names (blob {}, not {})", i.from, &blob.to_hex()[..12], &i.blob[..i.blob.len().min(12)]);

    }

    let dest = tree.join(rel);

    if let Some(dir) = dest.parent() {

        std::fs::create_dir_all(dir)?;

    }

    std::fs::write(&dest, bytes)?;

    Ok(blob.0)

}

/// The largest input a job takes: the hub's artifact limit.

pub const MAX_INPUT_BYTES: usize = 512 << 20;

const PATIENT_DELAYS: [Duration; 15] = {

    let mut d = [Duration::from_secs(15); 15];

    d[0] = Duration::from_secs(1);

    d[1] = Duration::from_secs(2);

    d[2] = Duration::from_secs(4);

    d[3] = Duration::from_secs(8);

    d

};

/// Send a request until the hub answers, waiting `delays` in turn after a

/// server error (5xx: nothing was committed) and, when `connection`, after

/// the hub could not be reached at all. A refusal (4xx) is final, as a

/// [`Refused`] the caller can look at.

fn retrying(

    url: &str,

    delays: &[Duration],

    connection: bool,

    mut send: impl FnMut() -> Result<ureq::Response, Box<ureq::Error>>,

) -> Result<ureq::Response> {

    let mut delays = delays.iter();

    loop {

        let why = match send().map_err(|e| *e) {

            Ok(r) => return Ok(r),

            Err(ureq::Error::Status(code, r)) if code < 500 => {

                let text = r.into_string().unwrap_or_default();

                return Err(anyhow::Error::new(Refused { code, text }).context(url.to_string()));

            }

            Err(ureq::Error::Status(code, r)) => format!("HTTP {code}: {}", r.into_string().unwrap_or_default()),

            Err(e) if connection => e.to_string(),

            Err(e) => bail!("{url}: {e}"),

        };

        match delays.next() {

            Some(d) => std::thread::sleep(*d),

            None => bail!("{url}: {why}"),

        }

    }

}

/// What the runner may run at once (flow.toml `[runner]`): at most

/// `max_jobs` jobs, whose sandboxes together fit `memory_mib` and `cpus`.

#[derive(Debug, Default, PartialEq, Eq)]

pub struct Budget {

    pub max_jobs: usize,

    pub memory_mib: Option<u32>,

    pub cpus: Option<u32>,

}

/// What is in use now.

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]

pub struct Used {

    pub jobs: usize,

    pub memory_mib: u32,

    pub cpus: u32,

}

impl Budget {