agent-migrate-core 0.6.1

use std::collections::BTreeMap;
use std::path::PathBuf;

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use thiserror::Error;

use crate::{BundleRelativePath, FileKind, FileRecord, Manifest, SchemaVersion};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Bundle {
    manifest: Manifest,
    content: BTreeMap<BundleRelativePath, Vec<u8>>,
}

impl Bundle {
    /// Restores a Bundle from an existing manifest and content map.
    ///
    /// # Errors
    ///
    /// Returns an error when the manifest and content do not match.
    pub fn from_parts(
        manifest: Manifest,
        content: BTreeMap<BundleRelativePath, Vec<u8>>,
    ) -> Result<Self, BundleError> {
        let bundle = Self { manifest, content };
        bundle.verify()?;
        Ok(bundle)
    }

    /// Builds a verified Bundle and computes every content hash.
    ///
    /// # Errors
    ///
    /// Returns an error for unsupported schemas or an inconsistent file set.
    pub fn from_files(
        mut manifest: Manifest,
        files: BTreeMap<BundleRelativePath, (FileKind, Vec<u8>)>,
    ) -> Result<Self, BundleError> {
        SchemaVersion::new(manifest.schema.major(), manifest.schema.minor())?;
        let mut content = BTreeMap::new();
        for (path, (kind, bytes)) in files {
            manifest.files.insert(
                path.clone(),
                FileRecord {
                    kind,
                    mode: 0o600,
                    sha256: digest(&bytes),
                },
            );
            content.insert(path, bytes);
        }
        let bundle = Self { manifest, content };
        bundle.verify()?;
        Ok(bundle)
    }

    /// Builds a verified Bundle with portable file modes.
    ///
    /// # Errors
    ///
    /// Returns an error when the schema or file set is invalid.
    pub fn from_files_with_modes(
        mut manifest: Manifest,
        files: BTreeMap<BundleRelativePath, (FileKind, u32, Vec<u8>)>,
    ) -> Result<Self, BundleError> {
        SchemaVersion::new(manifest.schema.major(), manifest.schema.minor())?;
        let mut content = BTreeMap::new();
        for (path, (kind, mode, bytes)) in files {
            manifest.files.insert(
                path.clone(),
                FileRecord {
                    kind,
                    mode: mode & 0o777,
                    sha256: digest(&bytes),
                },
            );
            content.insert(path, bytes);
        }
        let bundle = Self { manifest, content };
        bundle.verify()?;
        Ok(bundle)
    }

    #[must_use]
    pub fn manifest(&self) -> &Manifest {
        &self.manifest
    }

    #[must_use]
    pub fn files(&self) -> &BTreeMap<BundleRelativePath, Vec<u8>> {
        &self.content
    }

    #[must_use]
    pub fn content_hash(&self) -> String {
        let mut canonical = Vec::new();
        for (path, content) in &self.content {
            canonical.extend_from_slice(path.to_string().as_bytes());
            canonical.push(0);
            canonical.extend_from_slice(content);
            canonical.push(0xff);
        }
        digest(&canonical)
    }

    /// Verifies the schema, declared files, and SHA-256 hashes.
    ///
    /// # Errors
    ///
    /// Returns an error when schema or file integrity checks fail.
    pub fn verify(&self) -> Result<(), BundleError> {
        SchemaVersion::new(self.manifest.schema.major(), self.manifest.schema.minor())?;
        for (path, record) in &self.manifest.files {
            let bytes = self
                .content
                .get(path)
                .ok_or_else(|| BundleError::MissingFile(path.clone()))?;
            if digest(bytes) != record.sha256 {
                return Err(BundleError::HashMismatch { path: path.clone() });
            }
        }
        if let Some(path) = self
            .content
            .keys()
            .find(|path| !self.manifest.files.contains_key(*path))
        {
            return Err(BundleError::UnexpectedFile(path.clone()));
        }
        Ok(())
    }

    /// Serializes this verified Bundle to deterministic YAML bytes.
    ///
    /// # Errors
    ///
    /// Returns an error when verification or serialization fails.
    pub fn to_yaml_bytes(&self) -> Result<Vec<u8>, BundleError> {
        self.verify()?;
        let content = self
            .content
            .iter()
            .map(|(path, bytes)| (path.clone(), hex::encode(bytes)))
            .collect();
        let wire = BundleWire {
            manifest: self.manifest.clone(),
            content,
        };
        Ok(serde_yaml::to_string(&wire)?.into_bytes())
    }

    /// Parses and verifies a Bundle from YAML bytes.
    ///
    /// # Errors
    ///
    /// Returns an error for malformed YAML, unsupported schemas, invalid hex, or hash mismatch.
    pub fn from_yaml_bytes(input: &[u8]) -> Result<Self, BundleError> {
        let wire: BundleWire = serde_yaml::from_slice(input)?;
        SchemaVersion::new(wire.manifest.schema.major(), wire.manifest.schema.minor())?;
        let content = wire
            .content
            .into_iter()
            .map(|(path, encoded)| {
                hex::decode(encoded)
                    .map(|bytes| (path, bytes))
                    .map_err(BundleError::InvalidHex)
            })
            .collect::<Result<BTreeMap<_, _>, _>>()?;
        let bundle = Self {
            manifest: wire.manifest,
            content,
        };
        bundle.verify()?;
        Ok(bundle)
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct BundleWire {
    manifest: Manifest,
    content: BTreeMap<BundleRelativePath, String>,
}

fn digest(bytes: &[u8]) -> String {
    hex::encode(Sha256::digest(bytes))
}

#[derive(Debug, Error)]
pub enum BundleError {
    #[error("unsupported bundle schema major: {0}")]
    UnsupportedMajor(u16),
    #[error("bundle file is missing: {0}")]
    MissingFile(BundleRelativePath),
    #[error("bundle contains undeclared file: {0}")]
    UnexpectedFile(BundleRelativePath),
    #[error("bundle hash mismatch: {path}")]
    HashMismatch { path: BundleRelativePath },
    #[error("bundle content is not valid hexadecimal")]
    InvalidHex(#[source] hex::FromHexError),
    #[error("bundle path is invalid")]
    Path(#[from] crate::BundlePathError),
    #[error("bundle YAML is invalid")]
    Yaml(#[from] serde_yaml::Error),
    #[error("bundle directory already exists: {0}")]
    OutputExists(PathBuf),
    #[error("bundle directory contains an unsafe symbolic link: {0}")]
    DirectorySymlink(PathBuf),
    #[error("bundle directory contains an undeclared file: {0}")]
    UnexpectedDiskFile(PathBuf),
    #[error("bundle secret reference file does not match the manifest")]
    SecretReferenceMismatch,
    #[error("bundle file mode does not match the manifest: {0}")]
    ModeMismatch(PathBuf),
    #[error("bundle filesystem operation failed at {path}")]
    Io {
        path: PathBuf,
        #[source]
        source: std::io::Error,
    },
}