unki/lib/u_lib/src/builder.rs

use crate::{
    cache::JobCache,
    executor::{FutRes, Waiter, DynFut},
    models::{Agent, AssignedJob, JobMeta, JobType},
    utils::{CombinedResult, OneOrMany},
    UError,
};
use guess_host_triple::guess_host_triple;
use std::collections::HashMap;

pub struct JobBuilder {
    jobs: Waiter,
}

impl JobBuilder {
    pub fn from_request<J: OneOrMany<AssignedJob>>(job_requests: J) -> CombinedResult<Self> {
        let job_requests = job_requests.into_vec();
        let mut prepared: Vec<DynFut> = vec![];
        let mut result = CombinedResult::new();
        for req in job_requests {
            let job_meta = JobCache::get(&req.job_id);
            if job_meta.is_none() {
                result.err(UError::NoJob(req.job_id));
                continue;
            }
            let job_meta = job_meta.unwrap();
            let built_req = (|| {
                Ok(match job_meta.exec_type {
                    JobType::Shell => {
                        let meta = JobCache::get(&req.job_id).ok_or(UError::NoJob(req.job_id))?;
                        let curr_platform = guess_host_triple().unwrap_or("unknown").to_string();
                        if meta.platform != curr_platform {
                            return Err(UError::InsuitablePlatform(
                                meta.platform.clone(),
                                curr_platform,
                            ));
                        }
                        let job = AssignedJob::new(req.job_id, Some(&req));
                        prepared.push(Box::pin(job.run()))
                    }
                    JobType::Manage => prepared.push(Box::pin(Agent::run())),
                    _ => todo!(),
                })
            })();
            if let Err(e) = built_req {
                result.err(e)
            }
        }
        result.ok(Self {
            jobs: Waiter::new(prepared),
        });
        result
    }

    pub fn from_meta<J: OneOrMany<JobMeta>>(job_metas: J) -> CombinedResult<Self> {
        let job_requests = job_metas
            .into_vec()
            .into_iter()
            .map(|jm| {
                let j_uid = jm.id;
                JobCache::insert(jm);
                AssignedJob::new(j_uid, None)
            })
            .collect::<Vec<AssignedJob>>();
        JobBuilder::from_request(job_requests)
    }

    /// Spawn jobs and pop results later
    pub async fn spawn(mut self) -> Self {
        self.jobs = self.jobs.spawn().await;
        self
    }

    /// Spawn jobs and wait for result
    pub async fn wait(self) -> Vec<FutRes> {
        self.jobs.spawn().await.wait().await
    }

    /// Spawn one job and wait for result
    pub async fn wait_one(self) -> FutRes {
        self.jobs.spawn().await.wait().await.pop().unwrap()
    }
}

/// Store jobs and get results by name
pub struct NamedJobBuilder {
    builder: Option<JobBuilder>,
    job_names: Vec<&'static str>,
    results: HashMap<&'static str, FutRes>,
}

impl NamedJobBuilder {
    pub fn from_shell<J: OneOrMany<(&'static str, &'static str)>>(
        named_jobs: J,
    ) -> CombinedResult<Self> {
        let mut result = CombinedResult::new();
        let jobs: Vec<(&'static str, JobMeta)> = named_jobs
            .into_vec()
            .into_iter()
            .filter_map(
                |(alias, cmd)| match JobMeta::builder().with_shell(cmd).build() {
                    Ok(meta) => Some((alias, meta)),
                    Err(e) => {
                        result.err(e);
                        None
                    }
                },
            )
            .collect();
        result.ok(Self::from_meta(jobs));
        result
    }

    pub fn from_meta<J: OneOrMany<(&'static str, JobMeta)>>(named_jobs: J) -> Self {
        let mut job_names = vec![];
        let job_metas: Vec<JobMeta> = named_jobs
            .into_vec()
            .into_iter()
            .map(|(alias, meta)| {
                job_names.push(alias);
                meta
            })
            .collect();
        Self {
            builder: Some(JobBuilder::from_meta(job_metas).unwrap_one()),
            job_names,
            results: HashMap::new(),
        }
    }

    pub async fn wait(mut self) -> Self {
        let results = self.builder.take().unwrap().wait().await;
        for (name, result) in self.job_names.iter().zip(results.into_iter()) {
            self.results.insert(name, result);
        }
        self
    }

    pub fn pop_opt(&mut self, name: &'static str) -> Option<FutRes> {
        self.results.remove(name)
    }

    pub fn pop(&mut self, name: &'static str) -> FutRes {
        self.pop_opt(name).unwrap()
    }
}