//! `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, 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::PathBuf;

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

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 mut delays = delays.iter();

        loop {

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

                Ok(r) if r.status() == 204 => return Ok(None),

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

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

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

                    match delays.next() {

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

                        _ => bail!("{url}: HTTP {code}: {text}"),

                    }

                }

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

            }

        }

    }

    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),

        }

    }

    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 mut delays = RETRY_DELAYS.iter();

        loop {

            let req = self

                .agent

                .post(&url)

                .timeout(Duration::from_secs(600))

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

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

            match req.send_bytes(bytes) {

                Ok(r) => {

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

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

                }

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

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

                    match delays.next() {

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

                        _ => bail!("{url}: HTTP {code}: {text}"),

                    }

                }

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

            }

        }

    }

    /// 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)];

/// 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 {

    pub fn of(cfg: &RunnerConfig) -> Budget {

        Budget { max_jobs: cfg.max_jobs.unwrap_or(1).max(1), memory_mib: cfg.memory_mib, cpus: cfg.cpus }

    }

    /// The room to ask the hub for, or `None` when no job more can start.

    /// A lone job may use the whole budget even if it asks for more, so a

    /// large sandbox is never starved.

    pub fn room(&self, used: Used) -> Option<(Option<u32>, Option<u32>)> {

        if used.jobs >= self.max_jobs {

            return None;

        }

        if used.jobs == 0 {

            return Some((None, None));

        }

        let mem = self.memory_mib.map(|m| m.saturating_sub(used.memory_mib));

        let cpus = self.cpus.map(|c| c.saturating_sub(used.cpus));

        if mem == Some(0) || cpus == Some(0) {

            return None;

        }

        Some((mem, cpus))

    }

}

pub fn serve(home: &Home) -> Result<()> {

    let cfg = home.config.runner.clone().context("flow.toml has no [runner] section")?;

    let hub = Hub::new(&cfg)?;

    let name = cfg.name.clone().unwrap_or_else(|| "flow".into());

    let budget = Budget::of(&cfg);

    let key = home.runner_key(&name)?;

    eprintln!(

        "flow runner {name}: polling {} (up to {} jobs at once), signing receipts as {}",

        cfg.hub,

        budget.max_jobs,

        key.identity.hash().to_hex()

    );

    let used = std::sync::Mutex::new(Used::default());

    let mut idle = Duration::from_secs(2);

    let pending = report_pending(home);

    // The jobs held now, told to the hub at each claim so it can end the

    // ones this runner lost: a job is held from its claim until its result

    // is reported.

    let held = std::sync::Mutex::new(std::collections::BTreeSet::<i64>::new());

    held.lock().unwrap().extend(pending.iter().map(|(job, _)| *job));

    std::thread::scope(|scope| {

        for (job, f) in pending {

            let (hub, held) = (&hub, &held);

            scope.spawn(move || {

                eprintln!("job {job}: reporting the result kept from before the restart");

                report_finish(home, hub, job, &f);

                held.lock().unwrap().remove(&job);

            });

        }

        loop {

            let room = budget.room(*used.lock().unwrap());

            let Some((free_memory_mib, free_cpus)) = room else {

                std::thread::sleep(Duration::from_millis(500));

                continue;

            };

            let req = ClaimRequest {

                runner: name.clone(),

                runtimes: vec!["microvm".into()],

                free_memory_mib,

                free_cpus,

                info: Some(RunnerInfo {

                    version: env!("OXY_VERSION_LONG").into(),

                    ..RunnerInfo::here(cfg.labels.clone(), budget.max_jobs as u32, cfg.memory_mib, cfg.cpus)

                }),

                // Read each time: an image imported while the runner runs counts.

                images: Some(home.images()),

                image_archs: Some(home.image_archs()),

                holding: Some(held.lock().unwrap().iter().copied().collect()),

            };

            match hub.claim(&req) {

                Ok(Some(claim)) => {

                    idle = Duration::from_secs(2);

                    let (mem, cpus) = (claim.sandbox.memory_mib, claim.sandbox.cpus);

                    {

                        let mut u = used.lock().unwrap();

                        u.jobs += 1;

                        u.memory_mib += mem;

                        u.cpus += cpus;

                    }

                    held.lock().unwrap().insert(claim.job_id);

                    let (home, hub, cfg, used, key, held) = (home, &hub, &cfg, &used, &key, &held);

                    scope.spawn(move || {

                        let id = claim.job_id;

                        eprintln!(

                            "job {id}: {} {} on {}",

                            claim.repo,

                            claim.job.name,

                            claim.heads.first().map(|h| &h[..12]).unwrap_or("-")

                        );

                        let finish = run_claim(home, hub, cfg, key, &claim).unwrap_or_else(|e| Finish {

                            exit_code: None,

                            duration_ms: 0,

                            receipt: None,

                            signed_receipt: None,

                            error: Some(format!("{e:#}")),

                        });

                        eprintln!(

                            "job {id}: exit {:?}{}",

                            finish.exit_code,

                            finish.error.as_deref().map(|e| format!(" ({e})")).unwrap_or_default()

                        );

                        report_finish(home, hub, id, &finish);

                        held.lock().unwrap().remove(&id);

                        let mut u = used.lock().unwrap();

                        u.jobs -= 1;

                        u.memory_mib -= mem;

                        u.cpus -= cpus;

                    });

                }

                Ok(None) => std::thread::sleep(idle),

                Err(e) => {

                    eprintln!("claim: {e:#}");

                    std::thread::sleep(idle);

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

                }

            }

        }

    })

}

/// Batches records to the hub every second (or every 200 lines).

struct LogSink<'a> {

    hub: &'a Hub,

    job: i64,

    buf: Vec<Record>,

    last: Instant,

}

impl LogSink<'_> {

    fn push(&mut self, mut r: Record) {

        if r.text.len() > MAX_RECORD_TEXT {

            let mut cut = MAX_RECORD_TEXT;

            while !r.text.is_char_boundary(cut) {

                cut -= 1;

            }

            r.text.truncate(cut);

        }

        self.buf.push(r);

        if self.buf.len() >= 200 || self.last.elapsed() >= Duration::from_secs(1) {

            self.flush();

        }

    }

    fn flush(&mut self) {

        if !self.buf.is_empty() {

            if let Err(e) = self.hub.log(self.job, std::mem::take(&mut self.buf)) {

                eprintln!("job {}: log: {e:#}", self.job);

            }

        }

        self.last = Instant::now();

    }

}

/// Held while a job's workspace is prepared (see `run_claim`).

static PREPARE: std::sync::Mutex<()> = std::sync::Mutex::new(());

fn run_claim(home: &Home, hub: &Hub, cfg: &RunnerConfig, key: &RunnerKey, claim: &Claim) -> Result<Finish> {

    let started = Instant::now();

    let mut sink = LogSink { hub, job: claim.job_id, buf: Vec::new(), last: Instant::now() };

    let note =

        |sink: &mut LogSink, text: String| sink.push(Record { t: started.elapsed().as_millis() as u64, stream: "flow".into(), text });

    let heads: Vec<ChangeHash> = claim.heads.iter().map(|h| ChangeHash::from_hex(h).context("bad head")).collect::<Result<_>>()?;

    note(&mut sink, format!("fetching {} for state {}", claim.repo, claim.heads.iter().map(|h| &h[..12]).collect::<Vec<_>>().join("+")));

    // Preparing a workspace (the cached repositories, the export, the

    // image) is one job at a time: fetches into a cached repository must not

    // overlap, and the first clone passes its token through the process

    // environment. Only the microVMs run side by side.

    let mut preparing = Some(PREPARE.lock().unwrap_or_else(|p| p.into_inner()));

    let repo = cached_repo(home, cfg, claim)?;

    let scratch = home.root.join("work").join(format!("job-{}", claim.job_id));