//! Resolving an image reference, pulling its manifest and layers, and unpacking them.

use std::fs::{self, Permissions};
use std::io::{self, Read};
use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;

use oci_client::Reference;
use oci_client::client::{Client, ClientConfig};
use oci_client::manifest::OciDescriptor;
use oci_client::secrets::RegistryAuth as OciAuth;
use serde::Deserialize;
use sha2::{Digest, Sha256};
use tokio::io::AsyncWriteExt;
use tokio::sync::mpsc;

use super::helper::{UnpackRequest, run_unpack_helper};
use super::store::{ImageMeta, ImageStore, lock_file};
use super::unpack::{Ownership, UnpackStats, unpack_layers};
use super::{ImageError, PreparedImage, RegistryAuth};

impl ImageStore {
    /// Makes `image` available as an unpacked rootfs and returns it with its config.
    ///
    /// The manifest is always fetched, so the digest is current. When the digest already has an
    /// unpacked rootfs, nothing else is downloaded. Missing layer blobs are downloaded and checked
    /// against their digest. When the registry cannot be reached and the tag was pulled before,
    /// the cached image is used and a warning is logged. Progress lines go to `log`.
    pub async fn ensure(
        &self,
        image: &str,
        auth: Option<&RegistryAuth>,
        log: &mpsc::Sender<Vec<u8>>,
    ) -> Result<PreparedImage, ImageError> {
        let reference = parse_reference(image)?;
        let whole = reference.whole();
        log_line(log, format!("pulling {whole}")).await;

        let oci_auth = match auth {
            Some(auth) => OciAuth::Basic(auth.username.clone(), auth.password.clone()),
            None => OciAuth::Anonymous,
        };
        let client = Client::new(ClientConfig::default());

        let (manifest, digest) = match client.pull_image_manifest(&reference, &oci_auth).await {
            Ok(found) => found,
            Err(source) => {
                if let Some(prepared) = self.cached_for_tag(&whole)? {
                    log_line(
                        log,
                        format!(
                            "warning: registry unavailable ({source}), using cached image {}",
                            prepared.digest
                        ),
                    )
                    .await;
                    self.touch(hex_of(&prepared.digest)?)?;
                    return Ok(prepared);
                }
                return Err(ImageError::Registry {
                    image: whole,
                    source,
                });
            }
        };

        let hex = hex_of(&digest)?.to_owned();
        let lock_path = self.lock_path(&hex);
        let _lock = tokio::task::spawn_blocking(move || lock_file(&lock_path)).await??;

        if let Some(prepared) = self.cached(&hex)? {
            self.write_tag(&whole, &prepared.digest)?;
            self.touch(&hex)?;
            return Ok(prepared);
        }

        let mut config_bytes = Vec::new();
        client
            .pull_blob(&reference, &manifest.config, &mut config_bytes)
            .await
            .map_err(|source| ImageError::Registry {
                image: whole.clone(),
                source,
            })?;
        verify_bytes(&config_bytes, &manifest.config.digest)?;
        let config: ConfigFile = serde_json::from_slice(&config_bytes)?;
        let container = config.config.unwrap_or_default();

        let mut layer_hexes = Vec::with_capacity(manifest.layers.len());
        for layer in &manifest.layers {
            let layer_hex = hex_of(&layer.digest)?.to_owned();
            if !self.blob_path(&layer_hex).is_file() {
                log_line(log, format!("fetching layer {}", layer.digest)).await;
                self.download_blob(&client, &reference, layer, &layer_hex)
                    .await?;
            }
            layer_hexes.push(layer_hex);
        }

        log_line(log, format!("unpacking {} layers", layer_hexes.len())).await;
        let rootfs_parent = self.root_dir().join("rootfs");
        fs::create_dir_all(&rootfs_parent)?;
        let staging = tempfile::Builder::new()
            .prefix(".unpack-")
            .tempdir_in(&rootfs_parent)?;
        let staging_path = staging.path().to_path_buf();
        let layer_paths: Vec<PathBuf> = layer_hexes.iter().map(|h| self.blob_path(h)).collect();
        let max_bytes = self.max_unpacked_bytes();
        let stats = match self
            .unpack_into(staging_path.clone(), layer_paths, max_bytes)
            .await
        {
            Ok(stats) => stats,
            Err(err) => {
                // Subordinate-owned files stay behind if the runner removes them directly.
                let _ = self.remove_entry(&staging_path);
                return Err(err);
            }
        };
        if stats.rejected > 0 {
            log_line(
                log,
                format!(
                    "warning: {} layer entries were outside the rootfs rules and were skipped",
                    stats.rejected
                ),
            )
            .await;
        }
        if stats.skipped > 0 {
            log_line(
                log,
                format!("skipped {} device or special entries", stats.skipped),
            )
            .await;
        }

        // The rootfs must be readable by the container user, whatever the tempdir default is.
        fs::set_permissions(staging.path(), Permissions::from_mode(0o755))?;
        let unpacked = staging.keep();
        let final_dir = self.rootfs_path(&hex);
        if let Err(err) = fs::rename(&unpacked, &final_dir) {
            let _ = self.remove_entry(&unpacked);
            return Err(err.into());
        }

        let meta = ImageMeta {
            digest: digest.clone(),
            env: container.env.unwrap_or_default(),
            working_dir: container.working_dir.filter(|dir| !dir.is_empty()),
            user: container.user.filter(|user| !user.is_empty()),
            layers: layer_hexes,
        };
        self.write_meta(&hex, &meta)?;
        self.write_tag(&whole, &digest)?;
        self.touch(&hex)?;
        Ok(self.prepared(&hex, meta))
    }

    /// Unpacks `layers` into the staging dir `staging`. A subuid map unpacks in the helper process,
    /// so the files get subordinate owners. A single map unpacks in this process as before.
    async fn unpack_into(
        &self,
        staging: PathBuf,
        layers: Vec<PathBuf>,
        max_bytes: u64,
    ) -> Result<UnpackStats, ImageError> {
        let map = self.idmap.clone();
        let helper_exe = self.helper_exe.clone();
        let stats = tokio::task::spawn_blocking(move || {
            if map.is_subuid() {
                let req = UnpackRequest {
                    map,
                    rootfs: staging,
                    layers,
                    max_bytes,
                };
                run_unpack_helper(&helper_exe, &req)
            } else {
                unpack_layers(&staging, &layers, max_bytes, Ownership::Flatten)
            }
        })
        .await??;
        Ok(stats)
    }

    /// Downloads one layer blob into a temp file, checks its digest, and renames it into place.
    async fn download_blob(
        &self,
        client: &Client,
        reference: &Reference,
        layer: &OciDescriptor,
        hex: &str,
    ) -> Result<(), ImageError> {
        let blob_dir = self.root_dir().join("blobs/sha256");
        fs::create_dir_all(&blob_dir)?;
        let tmp = tempfile::NamedTempFile::new_in(&blob_dir)?;
        {
            let mut out = tokio::fs::File::from_std(tmp.reopen()?);
            client
                .pull_blob(reference, layer, &mut out)
                .await
                .map_err(|source| ImageError::Registry {
                    image: reference.whole(),
                    source,
                })?;
            out.flush().await?;
        }
        let actual = sha256_file(tmp.path().to_path_buf()).await?;
        if actual != hex {
            return Err(ImageError::DigestMismatch {
                expected: format!("sha256:{hex}"),
                actual: format!("sha256:{actual}"),
            });
        }
        tmp.persist(self.blob_path(hex)).map_err(|err| err.error)?;
        Ok(())
    }
}

/// The parts of an image config that the runner uses.
#[derive(Deserialize)]
struct ConfigFile {
    #[serde(default)]
    config: Option<ContainerConfig>,
}

#[derive(Default, Deserialize)]
struct ContainerConfig {
    #[serde(rename = "Env", default)]
    env: Option<Vec<String>>,
    #[serde(rename = "WorkingDir", default)]
    working_dir: Option<String>,
    #[serde(rename = "User", default)]
    user: Option<String>,
}

/// Parses an image reference, applying the Docker Hub defaults: `alpine` is
/// `docker.io/library/alpine:latest`.
pub(super) fn parse_reference(image: &str) -> Result<Reference, ImageError> {
    image
        .parse::<Reference>()
        .map_err(|err| ImageError::InvalidReference {
            image: image.to_owned(),
            reason: err.to_string(),
        })
}

/// Returns the hex part of a `sha256:<hex>` digest. The digest comes from a registry and becomes
/// a file name, so anything other than 64 lowercase hex digits is refused.
pub(super) fn hex_of(digest: &str) -> Result<&str, ImageError> {
    let hex = digest
        .strip_prefix("sha256:")
        .ok_or_else(|| ImageError::UnsupportedDigest(digest.to_owned()))?;
    let valid = hex.len() == 64
        && hex
            .bytes()
            .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b));
    if valid {
        Ok(hex)
    } else {
        Err(ImageError::UnsupportedDigest(digest.to_owned()))
    }
}

fn verify_bytes(bytes: &[u8], digest: &str) -> Result<(), ImageError> {
    let expected = hex_of(digest)?;
    let actual = hex::encode(Sha256::digest(bytes));
    if actual == expected {
        Ok(())
    } else {
        Err(ImageError::DigestMismatch {
            expected: digest.to_owned(),
            actual: format!("sha256:{actual}"),
        })
    }
}

/// Hashes a file on a blocking thread and returns the hex digest.
async fn sha256_file(path: PathBuf) -> Result<String, ImageError> {
    let hash = tokio::task::spawn_blocking(move || -> io::Result<String> {
        let mut file = fs::File::open(&path)?;
        let mut hasher = Sha256::new();
        let mut buf = vec![0u8; 64 * 1024];
        loop {
            let n = file.read(&mut buf)?;
            if n == 0 {
                break;
            }
            hasher.update(&buf[..n]);
        }
        Ok(hex::encode(hasher.finalize()))
    })
    .await??;
    Ok(hash)
}

/// Sends one line to the job log. A closed log channel is not an error for the pull.
async fn log_line(log: &mpsc::Sender<Vec<u8>>, line: String) {
    let mut bytes = line.into_bytes();
    bytes.push(b'\n');
    let _ = log.send(bytes).await;
}

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

    const DIGEST: &str = "sha256:0000000000000000000000000000000000000000000000000000000000000001";

    #[test]
    fn short_names_normalize_to_docker_hub() {
        assert_eq!(
            parse_reference("alpine").expect("parse").whole(),
            "docker.io/library/alpine:latest"
        );
        assert_eq!(
            parse_reference("rust:1").expect("parse").whole(),
            "docker.io/library/rust:1"
        );
        assert_eq!(
            parse_reference("ghcr.io/o/i:2").expect("parse").whole(),
            "ghcr.io/o/i:2"
        );
    }

    #[test]
    fn digest_references_keep_their_digest() {
        let image = format!("ghcr.io/o/i@{DIGEST}");
        let reference = parse_reference(&image).expect("parse");
        assert_eq!(reference.whole(), image);
        assert_eq!(reference.digest(), Some(DIGEST));
    }

    #[test]
    fn invalid_references_are_errors() {
        assert!(parse_reference("Bad Name!!").is_err());
    }

    #[test]
    fn only_clean_sha256_digests_are_accepted() {
        assert_eq!(hex_of(DIGEST).expect("ok").len(), 64);
        assert!(hex_of("sha512:00").is_err());
        assert!(hex_of("sha256:../../etc/passwd").is_err());
        assert!(
            hex_of("sha256:ABCDEF0000000000000000000000000000000000000000000000000000000000")
                .is_err()
        );
        assert!(hex_of("sha256:abc").is_err());
    }

    #[test]
    fn verify_bytes_checks_the_digest() {
        let digest = format!("sha256:{}", hex::encode(Sha256::digest(b"layer")));
        verify_bytes(b"layer", &digest).expect("matches");
        assert!(matches!(
            verify_bytes(b"other", &digest),
            Err(ImageError::DigestMismatch { .. })
        ));
    }

    #[tokio::test]
    async fn unreachable_registry_uses_the_cached_image() {
        let dir = tempfile::tempdir().expect("tempdir");
        let store = ImageStore::new(dir.path().join("oci"));
        let reference = "127.0.0.1:1/forge-test/cached:1";
        let hex = "1".repeat(64);
        let digest = format!("sha256:{hex}");

        fs::create_dir_all(store.rootfs_path(&hex)).expect("rootfs");
        fs::write(store.rootfs_path(&hex).join("marker"), b"x").expect("marker");
        store
            .write_meta(
                &hex,
                &ImageMeta {
                    digest: digest.clone(),
                    env: vec!["PATH=/bin".to_owned()],
                    working_dir: Some("/work".to_owned()),
                    user: None,
                    layers: Vec::new(),
                },
            )
            .expect("meta");
        store.write_tag(reference, &digest).expect("tag");

        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(64);
        let drain = tokio::spawn(async move {
            let mut lines = Vec::new();
            while let Some(line) = rx.recv().await {
                lines.push(String::from_utf8_lossy(&line).into_owned());
            }
            lines
        });
        let prepared = store.ensure(reference, None, &tx).await.expect("cached");
        drop(tx);
        let lines = drain.await.expect("drain");

        assert_eq!(prepared.digest, digest);
        assert_eq!(prepared.env, vec!["PATH=/bin".to_owned()]);
        assert_eq!(prepared.working_dir.as_deref(), Some("/work"));
        assert!(prepared.rootfs.join("marker").is_file());
        assert!(lines.iter().any(|l| l.contains("using cached image")));
    }

    #[tokio::test]
    async fn unreachable_registry_without_cache_is_an_error() {
        let dir = tempfile::tempdir().expect("tempdir");
        let store = ImageStore::new(dir.path().join("oci"));
        let (tx, _rx) = mpsc::channel::<Vec<u8>>(64);
        let result = store
            .ensure("127.0.0.1:1/forge-test/none:1", None, &tx)
            .await;
        assert!(matches!(result, Err(ImageError::Registry { .. })));
    }

    /// Pulls the real alpine image. Needs network access to Docker Hub.
    /// Run with: cargo test -p forge-runner oci -- --ignored
    #[tokio::test]
    #[ignore]
    async fn pulls_alpine_for_real() {
        let dir = tempfile::tempdir().expect("tempdir");
        let store = ImageStore::new(dir.path().join("oci"));
        let (tx, mut rx) = mpsc::channel::<Vec<u8>>(256);
        let drain = tokio::spawn(async move { while rx.recv().await.is_some() {} });

        let first = store
            .ensure("docker.io/library/alpine:latest", None, &tx)
            .await
            .expect("pull alpine");
        assert!(first.digest.starts_with("sha256:"));
        assert!(first.rootfs.join("bin/sh").symlink_metadata().is_ok());
        assert!(first.env.iter().any(|e| e.starts_with("PATH=")));

        // A second ensure with the same digest reuses the unpacked rootfs.
        let second = store
            .ensure("alpine", None, &tx)
            .await
            .expect("reuse alpine");
        assert_eq!(second.digest, first.digest);
        assert_eq!(second.rootfs, first.rootfs);

        drop(tx);
        let _ = drain.await;
    }
}