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

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

    }

    /// 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(())

    }

    pub fn finish(&self, job: i64, f: &Finish) -> Result<()> {

        self.post(&format!("/flow/jobs/{job}/finish"), f)?;

        Ok(())

    }

}

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

    eprintln!("flow runner {name}: polling {} (up to {} jobs at once)", cfg.hub, budget.max_jobs);

    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) = (home, &hub, &cfg, &used);

                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, &claim).unwrap_or_else(|e| Finish {

                        exit_code: None,

                        duration_ms: 0,

                        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, 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 files = workspace::export_state(&repo, &heads, &scratch.join("tree"))?;

        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 outcome = microvm::run(

            home,

            &claim.sandbox,

            &claim.job,

            &scratch.join("workspace.ext4"),

            &scratch,

            Duration::from_secs(claim.timeout_s),

            &mut |r| sink.push(r),

        )?;

        let receipt = outcome.exit_code.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: vec![],

            exit_code: code as u32,

            duration_ms: outcome.duration_ms,

        });

        Ok(Finish {

            exit_code: outcome.exit_code,

            duration_ms: started.elapsed().as_millis() as u64,

            receipt,

            error: outcome.timed_out.then(|| "stopped at the time limit".to_string()),

        })

    })();

    let _ = std::fs::remove_dir_all(&scratch);

    sink.flush();

    result

}

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

    let dir: PathBuf = home.root.join("repos").join(claim.repo.replace('/', "__"));

    let set_remote = |repo: &mut Repo| -> Result<()> {

        repo.config.set("remote.origin.url", &claim.clone_url).map_err(|e| anyhow::anyhow!("{e}"))?;

        repo.config.set("remote.origin.token", &claim.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 {}: {e}", claim.clone_url))?;

        return Ok(repo);

    }

    // First job for this repository: clone with the job's token.

    std::env::set_var("OXY_TOKEN", &claim.token);

    if let Some(ca) = &cfg.ca {

        std::env::set_var("OXY_CA", ca);

    }

    let cloned = Repo::clone(&claim.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}", claim.clone_url)

    })?;

    set_remote(&mut repo)?;

    Ok(repo)

}

#[cfg(test)]

mod tests {

    use super::*;

    /// A hub that answers each request with the next of `codes`, counting them.

    fn fake_hub(codes: Vec<u16>) -> (String, std::sync::Arc<std::sync::atomic::AtomicUsize>) {

        use std::io::{Read, Write};

        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();

        let addr = listener.local_addr().unwrap();

        let seen = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));

        let count = seen.clone();

        std::thread::spawn(move || {

            for code in codes {

                let (mut s, _) = listener.accept().unwrap();

                // Read the whole request (headers, then Content-Length bytes

                // of body) before answering: closing early would reset the

                // connection while the client is still sending.

                let mut req = Vec::new();

                let mut buf = [0u8; 8192];

                loop {

                    let n = s.read(&mut buf).unwrap_or(0);

                    if n == 0 {

                        break;

                    }

                    req.extend_from_slice(&buf[..n]);

                    let text = String::from_utf8_lossy(&req);

                    if let Some(end) = text.find("\r\n\r\n") {

                        let len = text[..end]

                            .lines()

                            .find_map(|l| {

                                l.to_ascii_lowercase().strip_prefix("content-length:").map(|v| v.trim().parse::<usize>().unwrap_or(0))

                            })

                            .unwrap_or(0);

                        if req.len() >= end + 4 + len {

                            break;

                        }

                    }

                }

                count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);

                let reason = if code == 204 { "No Content" } else { "Error" };

                let _ = write!(s, "HTTP/1.1 {code} {reason}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");

            }

        });

        (format!("http://{addr}"), seen)

    }

    fn hub_at(base: String) -> Hub {

        Hub { base, token: "t".into(), agent: ureq::AgentBuilder::new().timeout(Duration::from_secs(5)).build() }

    }

    #[test]

    fn server_errors_are_tried_again_and_client_errors_are_not() {

        let quick = [Duration::from_millis(10); 3];

        let (base, seen) = fake_hub(vec![500, 503, 204]);

        let r = hub_at(base).post_with("/flow/jobs/1/log", &serde_json::json!({}), &quick).unwrap();

        assert!(r.is_none());

        assert_eq!(seen.load(std::sync::atomic::Ordering::SeqCst), 3, "two retries, then accepted");

        let (base, seen) = fake_hub(vec![500, 500, 500, 500]);

        let e = hub_at(base).post_with("/flow/claim", &serde_json::json!({}), &quick).unwrap_err().to_string();

        assert!(e.contains("HTTP 500"), "{e}");

        assert_eq!(seen.load(std::sync::atomic::Ordering::SeqCst), 4, "gives up after the last delay");

        let (base, seen) = fake_hub(vec![401]);

        assert!(hub_at(base).post_with("/flow/claim", &serde_json::json!({}), &quick).is_err());

        assert_eq!(seen.load(std::sync::atomic::Ordering::SeqCst), 1, "a refusal is final");

    }

    #[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");