//! `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, 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 })
    }

    fn post(&self, path: &str, body: &impl serde::Serialize) -> Result<Option<ureq::Response>> {
        let url = format!("{}/api/v1{path}", self.base);
        match self.agent.post(&url).set("Authorization", &format!("Bearer {}", self.token)).send_json(body) {
            Ok(r) if r.status() == 204 => Ok(None),
            Ok(r) => Ok(Some(r)),
            Err(ureq::Error::Status(code, r)) => bail!("{url}: HTTP {code}: {}", r.into_string().unwrap_or_default()),
            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}"),
            }
        }
    }

    pub fn finish(&self, job: i64, f: &Finish) -> Result<()> {
        self.post(&format!("/flow/jobs/{job}/finish"), f)?;
        Ok(())
    }
}

/// 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);
    std::thread::scope(|scope| 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 };
        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;
                }
                let (home, hub, cfg, used, key) = (home, &hub, &cfg, &used, &key);
                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()
                    );
                    if let Err(e) = hub.finish(id, &finish) {
                        eprintln!("job {id}: reporting failed: {e:#}");
                    }
                    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));
    let _ = std::fs::remove_dir_all(&scratch);
    std::fs::create_dir_all(scratch.join("tree"))?;
    let result = (|| -> Result<Finish> {
        let mut files = workspace::export_state(&repo, &heads, &scratch.join("tree"))?;
        // Pinned subrepos (Oxygen spec §7.8), at their paths, as part of the
        // workspace and of the scope hash.
        for sub in &claim.subrepos {
            if sub.path.split('/').any(|p| p.is_empty() || p == "." || p == ".." || p == ".oxy") {
                bail!("subrepo path {:?} is not a clean relative path", sub.path);
            }
            note(
                &mut sink,
                format!(
                    "fetching subrepo {} at {} for {}",
                    sub.repo,
                    sub.heads.iter().map(|h| &h[..12]).collect::<Vec<_>>().join("+"),
                    sub.path
                ),
            );
            let sub_repo = cached_repo_at(home, cfg, &sub.repo, &sub.clone_url, &claim.token)?;
            let sub_heads: Vec<ChangeHash> =
                sub.heads.iter().map(|h| ChangeHash::from_hex(h).context("bad subrepo head")).collect::<Result<_>>()?;
            let dir = scratch.join("tree").join(&sub.path);
            std::fs::create_dir_all(&dir)?;
            let sub_files = workspace::export_state(&sub_repo, &sub_heads, &dir)?;
            files.extend(sub_files.into_iter().map(|(p, h)| (format!("{}/{p}", sub.path), h)));
        }
        let scope_hash = scope_tree_hash(&claim.job.scope, files.iter().map(|(p, h)| (p.as_str(), *h)));
        workspace::image(&scratch.join("tree"), &scratch.join("workspace.ext4"))?;
        std::fs::remove_dir_all(scratch.join("tree"))?;
        preparing.take();
        note(
            &mut sink,
            format!(
                "booting a microVM: {} vCPU, {} MiB, rootfs {}",
                claim.sandbox.cpus,
                claim.sandbox.memory_mib,
                &claim.sandbox.rootfs[..12]
            ),
        );
        let job = claim.job.in_sandbox(&claim.sandbox).with_flow_env(&claim.flow_env());
        let outcome = microvm::run(
            home,
            &claim.sandbox,
            &job,
            &scratch.join("workspace.ext4"),
            &scratch,
            Duration::from_secs(claim.timeout_s),
            &mut |r| sink.push(r),
        )?;
        // A passing job's artifacts, read from its disks and uploaded; a
        // declared artifact that cannot be kept fails the job.
        let mut outputs = vec![];
        let mut output_error = None;
        if outcome.exit_code == Some(0) {
            for o in &claim.job.outputs {
                let (image, inner) = crate::outputs::locate(&claim.job, &scratch.join("workspace.ext4"), &scratch, &o.path);
                let kept = crate::outputs::read(&image, &inner, &scratch).and_then(|bytes| {
                    let blob = Repo::blob_ref(&bytes);
                    let named = hub.artifact(claim.job_id, &o.name, &bytes)?;
                    if named != blob.to_hex() {
                        bail!("the hub kept it as {named}, not {}", blob.to_hex());
                    }
                    note(&mut sink, format!("artifact {}: {} bytes, blob {}", o.name, bytes.len(), &named[..12]));
                    Ok(blob.0)
                });
                match kept {
                    Ok(b) => outputs.push(b),
                    Err(e) => {
                        output_error = Some(format!("artifact {}: {e:#}", o.name));
                        break;
                    }
                }
            }
        }
        if let Some(e) = &output_error {
            note(&mut sink, e.clone());
        }
        let receipt = outcome.exit_code.filter(|_| output_error.is_none()).map(|code| Receipt {
            job_name: claim.job.name.clone(),
            pipeline_hash: hex32(&claim.pipeline_hash),
            scope_tree_hash: scope_hash,
            inputs: vec![],
            environment: outcome.environment.hash(),
            outputs,
            exit_code: code as u32,
            duration_ms: outcome.duration_ms,
        });
        let signed_receipt = match &receipt {
            Some(r) => Some(hex::encode(sign_receipt(&repo, &heads, cfg, key, claim, r)?.encode())),
            None => None,
        };
        Ok(Finish {
            exit_code: outcome.exit_code,
            duration_ms: started.elapsed().as_millis() as u64,
            receipt,
            signed_receipt,
            error: output_error.or_else(|| outcome.timed_out.then(|| "stopped at the time limit".to_string())),
        })
    })();
    let _ = std::fs::remove_dir_all(&scratch);
    sink.flush();
    result
}

/// Bind `receipt` to the state of `heads` in `repo` and sign it as the
/// runner (Oxygen spec §14.6). The state is computed here, not taken from
/// the claim, so the runner vouches only for what it fetched.
fn sign_receipt(
    repo: &Repo,
    heads: &[ChangeHash],
    cfg: &RunnerConfig,
    key: &RunnerKey,
    claim: &Claim,
    receipt: &Receipt,
) -> Result<oxy_core::receipt::SignedReceipt> {
    let state = state_of(heads);
    let run = format!("{}/{}/flow/{}", cfg.hub.trim_end_matches('/'), claim.repo, claim.run_number);
    Ok(oxy_core::receipt::SignedReceipt::sign(receipt.to_oxy(), repo.id()?, state, &run, &key.identity, &key.secret))
}

/// The ObjectRef of the state whose heads are `heads` (Oxygen spec §10.1).
pub fn state_of(heads: &[ChangeHash]) -> oxy_core::hash::ObjectRef {
    use oxy_core::hash::{hash_object, ObjectType};
    hash_object(ObjectType::State, &oxy_core::objects::StateObject::new(heads.to_vec()).encode())
}

fn hex32(s: &str) -> [u8; 32] {
    let mut out = [0; 32];
    if let Ok(v) = hex::decode(s) {
        if v.len() == 32 {
            out.copy_from_slice(&v);
        }
    }
    out
}

/// The runner's copy of a repository, fetched up to date with the job's token.
fn cached_repo(home: &Home, cfg: &RunnerConfig, claim: &Claim) -> Result<Repo> {
    cached_repo_at(home, cfg, &claim.repo, &claim.clone_url, &claim.token)
}

/// The runner's copy of `repo`, cloned or fetched with `token`.
fn cached_repo_at(home: &Home, cfg: &RunnerConfig, repo_name: &str, clone_url: &str, token: &str) -> Result<Repo> {
    let dir: PathBuf = home.root.join("repos").join(repo_name.replace('/', "__"));
    let set_remote = |repo: &mut Repo| -> Result<()> {
        repo.config.set("remote.origin.url", clone_url).map_err(|e| anyhow::anyhow!("{e}"))?;
        repo.config.set("remote.origin.token", token).map_err(|e| anyhow::anyhow!("{e}"))?;
        if let Some(ca) = &cfg.ca {
            repo.config.set("remote.origin.ca", &ca.to_string_lossy()).map_err(|e| anyhow::anyhow!("{e}"))?;
        }
        repo.save_config().map_err(|e| anyhow::anyhow!("{e}"))
    };
    if dir.join(".oxy").is_dir() {
        let mut repo = Repo::open(&dir).map_err(|e| anyhow::anyhow!("{e}"))?;
        set_remote(&mut repo)?;
        repo.fetch("origin").map_err(|e| anyhow::anyhow!("fetch {clone_url}: {e}"))?;
        return Ok(repo);
    }
    // First job for this repository: clone with the job's token.
    std::env::set_var("OXY_TOKEN", token);
    if let Some(ca) = &cfg.ca {
        std::env::set_var("OXY_CA", ca);
    }
    let cloned = Repo::clone(clone_url, &dir);
    std::env::remove_var("OXY_TOKEN");
    std::env::remove_var("OXY_CA");
    let mut repo = cloned.map_err(|e| {
        let _ = std::fs::remove_dir_all(&dir);
        anyhow::anyhow!("clone {}: {e}", clone_url)
    })?;
    set_remote(&mut repo)?;
    Ok(repo)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_budget_admits_jobs_that_fit() {
        let b = Budget { max_jobs: 3, memory_mib: Some(24_576), cpus: Some(8) };
        let used = |jobs, memory_mib, cpus| Used { jobs, memory_mib, cpus };
        assert_eq!(b.room(used(0, 0, 0)), Some((None, None)), "a lone job may take everything");
        assert_eq!(b.room(used(1, 8192, 4)), Some((Some(16_384), Some(4))));
        assert_eq!(b.room(used(2, 16_384, 8)), None, "no vCPU left");
        assert_eq!(b.room(used(3, 1024, 1)), None, "at max_jobs");
        let one = Budget { max_jobs: 1, memory_mib: None, cpus: None };
        assert_eq!(one.room(used(0, 0, 0)), Some((None, None)));
        assert_eq!(one.room(used(1, 0, 0)), None, "the default runs one job at a time");
        let open = Budget { max_jobs: 2, memory_mib: None, cpus: None };
        assert_eq!(open.room(used(1, 8192, 4)), Some((None, None)), "no limits configured");
    }
}