//! Images by hash (oxygenflow#39, OxygenFlow docs/jetstream.md §8): rootfs

//! images kept site-wide for runners to fetch, so nobody copies an image

//! between machines. An image arrives as chunks (Oxygen's FastCDC, sent

//! only when this hub lacks them) and then its manifest, which is accepted

//! only when every chunk is here and the image they make has the hash it

//! is named by. Chunks are kept zstd-compressed under

//! `<data>/images/chunks`, shared by every image that has them.

use crate::db::{now_ms, Pool};

use jet::images::{is_hash, Manifest};

use std::io::Write;

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

/// Where this hub keeps images.

#[derive(Clone, Debug)]

pub struct Store {

    root: PathBuf,

}

impl Store {

    pub fn new(data: &Path) -> Store {

        Store { root: data.join("images") }

    }

    fn chunk_path(&self, hash: &str) -> PathBuf {

        self.root.join("chunks").join(&hash[..2]).join(hash)

    }

    fn manifest_path(&self, hash: &str) -> PathBuf {

        self.root.join("manifests").join(format!("{hash}.json"))

    }

    pub fn has_chunk(&self, hash: &str) -> bool {

        is_hash(hash) && self.chunk_path(hash).exists()

    }

    /// Keep chunk `hash`, whose bytes the caller has checked.

    pub fn put_chunk(&self, hash: &str, bytes: &[u8]) -> std::io::Result<()> {

        let path = self.chunk_path(hash);

        if path.exists() {

            return Ok(());

        }

        std::fs::create_dir_all(path.parent().expect("a chunk's directory"))?;

        let tmp = path.with_extension(format!("tmp{}", std::process::id()));

        let mut f = std::fs::File::create(&tmp)?;

        f.write_all(&zstd::encode_all(bytes, 3)?)?;

        f.sync_all()?;

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

    }

    pub fn chunk(&self, hash: &str) -> std::io::Result<Vec<u8>> {

        if !is_hash(hash) {

            return Err(std::io::Error::new(std::io::ErrorKind::NotFound, "not a chunk name"));

        }

        zstd::decode_all(std::fs::File::open(self.chunk_path(hash))?)

    }

    pub fn manifest(&self, hash: &str) -> Option<Manifest> {

        if !is_hash(hash) {

            return None;

        }

        serde_json::from_slice(&std::fs::read(self.manifest_path(hash)).ok()?).ok()

    }

    /// Accept manifest `m`: well formed, every chunk here, and the image

    /// they make hashing to its name. Reads the whole image once.

    pub fn verify_and_keep(&self, m: &Manifest) -> Result<(), String> {

        m.check().map_err(|e| e.to_string())?;

        let mut whole = blake3::Hasher::new();

        for c in &m.chunks {

            if c.is_zeros() {

                whole.update(&vec![0u8; c.len()]);

                continue;

            }

            let bytes = self.chunk(c.hash()).map_err(|_| format!("chunk {} was not sent; send it first", &c.hash()[..12]))?;

            if bytes.len() != c.len() {

                return Err(format!("chunk {} is {} bytes, not {}", &c.hash()[..12], bytes.len(), c.len()));

            }

            whole.update(&bytes);

        }

        let got = whole.finalize().to_hex().to_string();

        if got != m.hash {

            return Err(format!("the chunks make an image named {}, not {}", &got[..12], &m.hash[..12]));

        }

        self.keep_manifest(m)

    }

    /// Keep manifest `m` of an image only runners hold (oxygenflow#94): well

    /// formed, its chunks not here, so not checked against its hash; a

    /// fetcher checks the image it assembles against the pinned hash.

    /// Never over a manifest already kept.

    pub fn keep_peer_manifest(&self, m: &Manifest) -> Result<(), String> {

        m.check().map_err(|e| e.to_string())?;

        if self.manifest_path(&m.hash).exists() {

            return Ok(());

        }

        self.keep_manifest(m)

    }

    /// Write manifest `m`, whose image is here.

    fn keep_manifest(&self, m: &Manifest) -> Result<(), String> {

        let path = self.manifest_path(&m.hash);

        std::fs::create_dir_all(path.parent().expect("the manifests directory")).map_err(|e| e.to_string())?;

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

        std::fs::write(&tmp, serde_json::to_vec(m).map_err(|e| e.to_string())?).map_err(|e| e.to_string())?;

        std::fs::rename(&tmp, &path).map_err(|e| e.to_string())

    }

    /// What image `m` is for, read from its own files (OxygenFlow spec

    /// §11.2, oxygenhub#173): its shell's ELF header, found through its ext4

    /// out of the kept chunks, a few of them, not the whole image. `None`

    /// when it is not ext4 or has no shell to read.

    pub fn read_arch(&self, m: &Manifest) -> Option<&'static str> {

        let mut starts = Vec::with_capacity(m.chunks.len());

        let mut end = 0u64;

        for c in &m.chunks {

            starts.push(end);

            end += c.len() as u64;

        }

        let mut cache: std::collections::HashMap<usize, Vec<u8>> = std::collections::HashMap::new();

        let short = || std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "past the image's end");

        jet::rootfs::arch(&mut |at: u64, buf: &mut [u8]| {

            let mut done = 0usize;

            while done < buf.len() {

                let pos = at + done as u64;

                let i = starts.partition_point(|&s| s <= pos).checked_sub(1).ok_or_else(short)?;

                let c = &m.chunks[i];

                let within = (pos - starts[i]) as usize;

                if within >= c.len() {

                    return Err(short());

                }

                let n = (c.len() - within).min(buf.len() - done);

                if c.is_zeros() {

                    buf[done..done + n].fill(0);

                } else {

                    if !cache.contains_key(&i) {

                        if cache.len() >= 32 {

                            cache.clear();

                        }

                        cache.insert(i, self.chunk(c.hash())?);

                    }

                    let bytes = &cache[&i];

                    if bytes.len() != c.len() {

                        return Err(short());

                    }

                    buf[done..done + n].copy_from_slice(&bytes[within..within + n]);

                }

                done += n;

            }

            Ok(())

        })

    }

    /// Keep image `m`, checked as [`Store::verify_and_keep`], with the

    /// architecture its own files name: the manifest's word is taken only

    /// when the image cannot be read. The manifest as kept.

    pub fn keep_reading_arch(&self, m: &Manifest) -> Result<Manifest, String> {

        self.verify_and_keep(m)?;

        match self.read_arch(m) {

            Some(a) if m.arch.as_deref() != Some(a) => {

                let mut kept = m.clone();

                kept.arch = Some(a.to_string());

                self.keep_manifest(&kept)?;

                Ok(kept)

            }

            _ => Ok(m.clone()),

        }

    }

    /// Record `arch` in kept image `hash`'s manifest.

    fn set_arch(&self, hash: &str, arch: &str) -> Result<(), String> {

        let mut m = self.manifest(hash).ok_or("no such image")?;

        m.arch = Some(arch.to_string());

        self.keep_manifest(&m)

    }

}

/// Record an image this hub now keeps, its chunks and all: one only

/// runners held before becomes the hub's.

pub async fn record(pool: &Pool, m: &Manifest, uploader: &str) -> sqlx::Result<()> {

    sqlx::query(

        "INSERT INTO jet_images (hash, size, arch, chunks, uploader, created_ms) VALUES ($1, $2, $3, $4, $5, $6) \

         ON CONFLICT (hash) DO UPDATE SET source = 'hub', holder = '', \

         arch = CASE WHEN excluded.arch != '' THEN excluded.arch ELSE jet_images.arch END",

    )

    .bind(&m.hash)

    .bind(m.size as i64)

    .bind(m.arch.clone().unwrap_or_default())

    .bind(m.chunks.len() as i64)

    .bind(uploader)

    .bind(now_ms())

    .execute(pool)

    .await?;

    if let Some(a) = m.arch.as_ref().filter(|a| !a.is_empty()) {

        crate::db::flow::record_image_archs(pool, &[(m.hash.clone(), a.clone())].into_iter().collect()).await?;

    }

    Ok(())

}

/// Learn what every image is for (oxygenhub#173), once at start: read

/// each kept image that has no architecture from its own files, then give

/// any image still unknown the architecture its jobs passed on. A runner's

/// own report, read from the image it holds, wins over what was learned.

/// How many images were read.

pub async fn learn_archs(app: &crate::Shared) -> sqlx::Result<usize> {

    let mut read = 0;

    for k in kept(&app.db).await?.into_iter().filter(|k| k.arch.is_empty()) {

        let store = app.images();

        let hash = k.hash.clone();

        let arch = tokio::task::spawn_blocking(move || {

            let a = store.read_arch(&store.manifest(&hash)?)?;

            store.set_arch(&hash, a).ok()?;

            Some(a)

        })

        .await

        .ok()

        .flatten();

        if let Some(a) = arch {

            sqlx::query("UPDATE jet_images SET arch = $1 WHERE hash = $2").bind(a).bind(&k.hash).execute(&app.db).await?;

            crate::db::flow::record_image_archs(&app.db, &[(k.hash.clone(), a.to_string())].into_iter().collect()).await?;

            read += 1;

        }

    }

    crate::db::flow::learn_archs_from_runs(&app.db).await?;

    Ok(read)

}

/// Record an image only runners hold, whose manifest `holder` sent

/// (oxygenflow#94). Never over one the hub keeps.

pub async fn record_peer(pool: &Pool, m: &Manifest, holder: &str) -> sqlx::Result<()> {

    sqlx::query(

        "INSERT INTO jet_images (hash, size, arch, chunks, uploader, created_ms, source, holder) VALUES ($1, $2, $3, $4, $5, $6, 'peer', $7) \

         ON CONFLICT (hash) DO NOTHING",

    )

    .bind(&m.hash)

    .bind(m.size as i64)

    .bind(m.arch.clone().unwrap_or_default())

    .bind(m.chunks.len() as i64)

    .bind(holder)

    .bind(now_ms())

    .bind(holder)

    .execute(pool)

    .await?;

    if let Some(a) = m.arch.as_ref().filter(|a| !a.is_empty()) {

        crate::db::flow::record_image_archs(pool, &[(m.hash.clone(), a.clone())].into_iter().collect()).await?;

    }

    Ok(())

}

/// Every image this hub keeps a manifest of, by hash: its own and those

/// only runners hold.

pub async fn held(pool: &Pool) -> sqlx::Result<std::collections::HashSet<String>> {

    Ok(sqlx::query_scalar::<_, String>("SELECT hash FROM jet_images").fetch_all(pool).await?.into_iter().collect())

}

/// Where the images jobs need can come from (oxygenflow#94): the hub's

/// own, pushed and verified, and those only runners hold, each with the

/// swarm sites of the online runners holding it now.

#[derive(Debug, Default)]

pub struct Reach {

    pub hub: std::collections::HashSet<String>,

    pub peers: std::collections::HashMap<String, Vec<String>>,

    /// Each runner in a swarm now, and its site.

    pub sites: std::collections::HashMap<i64, String>,

}

impl Reach {

    /// Whether a runner that fetches images can get `image`: from the hub,

    /// or from a runner of its own site.

    pub fn fetchable(&self, runner: i64, image: &str) -> bool {

        self.hub.contains(image) || self.sites.get(&runner).is_some_and(|s| self.peers.get(image).is_some_and(|v| v.contains(s)))

    }

    /// Every image `runner` can fetch.

    pub fn for_runner(&self, runner: i64) -> std::collections::HashSet<String> {

        let mut out = self.hub.clone();

        out.extend(self.peers.keys().filter(|i| self.fetchable(runner, i)).cloned());

        out

    }

    /// Whether only runners hold `image`.

    pub fn peer_only(&self, image: &str) -> bool {

        !self.hub.contains(image) && self.peers.contains_key(image)

    }

}

/// [`Reach`] now.

pub async fn reach(app: &crate::Shared) -> sqlx::Result<Reach> {

    let rows: Vec<(String, String)> = sqlx::query_as("SELECT hash, source FROM jet_images").fetch_all(&app.db).await?;

    let mut reach = Reach::default();

    for (hash, source) in rows {

        if source == "peer" {

            reach.peers.insert(hash, Vec::new());

        } else {

            reach.hub.insert(hash);

        }

    }

    for r in crate::db::runners::list(&app.db).await? {

        if r.state == "disabled" || !r.online() {

            continue;

        }

        let Some(addr) = app.swarm.addr(r.id) else { continue };

        for image in r.image_list().unwrap_or_default() {

            if let Some(sites) = reach.peers.get_mut(&image.to_ascii_lowercase()) {

                if !sites.contains(&addr.site) {

                    sites.push(addr.site.clone());

                }

            }

        }

        reach.sites.insert(r.id, addr.site);

    }

    Ok(reach)

}

/// An image this hub keeps (`jet_images`).

#[derive(sqlx::FromRow)]

pub struct Kept {

    pub hash: String,

    pub size: i64,

    pub arch: String,

    pub chunks: i64,

    pub uploader: String,

    pub created_ms: i64,

    /// 'hub', or 'peer' for one only runners hold (oxygenflow#94).

    pub source: String,

    pub holder: String,

}

pub async fn kept(pool: &Pool) -> sqlx::Result<Vec<Kept>> {

    sqlx::query_as("SELECT hash, size, arch, chunks, uploader, created_ms, source, holder FROM jet_images").fetch_all(pool).await

}

/// One repository's use of one image: its jobs since `since_ms`, and when

/// it last ran one.

#[derive(sqlx::FromRow)]

pub struct Use {

    pub image: String,

    pub owner: String,

    pub name: String,

    pub jobs: i64,

    pub last_ms: i64,

}

/// What each repository's jobs ran on since `since_ms`, newest first.

pub async fn uses(pool: &Pool, since_ms: i64) -> sqlx::Result<Vec<Use>> {

    let rootfs = crate::db::json_text("j.sandbox", "rootfs");

    sqlx::query_as(&format!(

        "SELECT lower({rootfs}) AS image, p.owner, p.name, COUNT(*) AS jobs,

                MAX(COALESCE(j.started_ms, r.created_ms)) AS last_ms

         FROM flow_jobs j JOIN flow_runs r ON r.id = j.run_id JOIN repositories p ON p.id = r.repo_id

         WHERE r.created_ms >= $1 AND COALESCE({rootfs}, '') != ''

         GROUP BY lower({rootfs}), p.id, p.owner, p.name ORDER BY last_ms DESC"

    ))

    .bind(since_ms)

    .fetch_all(pool)

    .await

}

/// What one repository's jobs ran on since `since_ms`, newest first

/// (its Images page, oxygenhub#254).

pub async fn uses_of(pool: &Pool, repo_id: i64, since_ms: i64) -> sqlx::Result<Vec<Use>> {

    let rootfs = crate::db::json_text("j.sandbox", "rootfs");

    sqlx::query_as(&format!(

        "SELECT lower({rootfs}) AS image, p.owner, p.name, COUNT(*) AS jobs,

                MAX(COALESCE(j.started_ms, r.created_ms)) AS last_ms

         FROM flow_jobs j JOIN flow_runs r ON r.id = j.run_id JOIN repositories p ON p.id = r.repo_id

         WHERE r.created_ms >= $1 AND p.id = $2 AND COALESCE({rootfs}, '') != ''

         GROUP BY lower({rootfs}), p.id, p.owner, p.name ORDER BY last_ms DESC"

    ))

    .bind(since_ms)

    .bind(repo_id)

    .fetch_all(pool)

    .await

}

/// The architecture runners reported for each image (`flow_images`).

pub async fn reported_archs(pool: &Pool) -> sqlx::Result<std::collections::HashMap<String, String>> {

    Ok(sqlx::query_as::<_, (String, String)>("SELECT hash, arch FROM flow_images").fetch_all(pool).await?.into_iter().collect())

}

#[cfg(test)]

mod reach_tests {

    use super::Reach;

    #[test]

    fn an_image_only_runners_hold_is_fetchable_in_their_site_alone() {

        let mut r = Reach::default();

        r.hub.insert("hub-image".into());

        r.peers.insert("peer-image".into(), vec!["home".into()]);

        r.sites.insert(1, "home".into());

        r.sites.insert(2, "cloud".into());

        assert!(r.fetchable(1, "hub-image") && r.fetchable(2, "hub-image") && r.fetchable(3, "hub-image"));

        assert!(r.fetchable(1, "peer-image"), "a runner of the holder's site");

        assert!(!r.fetchable(2, "peer-image"), "another site");

        assert!(!r.fetchable(3, "peer-image"), "a runner in no swarm");

        assert!(r.peer_only("peer-image") && !r.peer_only("hub-image"));

        assert_eq!(r.for_runner(1).len(), 2);

        assert_eq!(r.for_runner(2).len(), 1);

    }

}