use bex_types::{BexError, Manifest};
use sha2::{Sha256, Digest};

pub const MAGIC: &[u8; 4] = b"BEX\x01";
pub const CONTAINER_VERSION: u16 = 1;
pub const HEADER_LEN: usize = 96;

pub struct BexPackage {
    pub manifest: Manifest,
    pub wasm: Vec<u8>,
}

fn sha256(data: &[u8]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(data);
    hasher.finalize().into()
}

pub fn read_manifest(data: &[u8]) -> Result<Manifest, BexError> {
    if data.len() < HEADER_LEN { return Err(BexError::ManifestInvalid("too short".into())); }
    if &data[0..4] != MAGIC { return Err(BexError::ManifestInvalid("bad magic".into())); }
    let expected_crc = u32::from_le_bytes(data[92..96].try_into().unwrap());
    let actual_crc = crc32fast::hash(&data[0..92]);
    if expected_crc != actual_crc { return Err(BexError::ManifestInvalid("header CRC mismatch".into())); }
    let manifest_len = u32::from_le_bytes(data[8..12].try_into().unwrap()) as usize;
    if HEADER_LEN.checked_add(manifest_len).map_or(true, |end| end > data.len()) {
        return Err(BexError::ManifestInvalid(
            format!("manifest_len {} exceeds package size {}", manifest_len, data.len())
        ));
    }
    let manifest_bytes = &data[HEADER_LEN..HEADER_LEN + manifest_len];
    let expected_hash = &data[60..92];
    let actual_hash = sha256(manifest_bytes);
    if expected_hash != actual_hash { return Err(BexError::ManifestInvalid("manifest hash mismatch".into())); }
    serde_yaml::from_slice(manifest_bytes).map_err(|e| BexError::ManifestInvalid(format!("yaml: {e}")))
}

pub fn unpack(data: &[u8]) -> Result<BexPackage, BexError> {
    let manifest = read_manifest(data)?;
    let flags = u16::from_le_bytes(data[6..8].try_into().unwrap());
    let manifest_len = u32::from_le_bytes(data[8..12].try_into().unwrap()) as usize;
    if HEADER_LEN.checked_add(manifest_len).map_or(true, |end| end > data.len()) {
        return Err(BexError::ManifestInvalid(
            format!("manifest_len {} exceeds package size {}", manifest_len, data.len())
        ));
    }
    let wasm_original_len = u64::from_le_bytes(data[20..28].try_into().unwrap()) as usize;
    let wasm_start = HEADER_LEN + manifest_len;
    let wasm = if flags & 1 != 0 {
        let mut out = Vec::with_capacity(wasm_original_len);
        zstd::stream::copy_decode(&data[wasm_start..], &mut out)
            .map_err(|e| BexError::Internal(format!("zstd: {e}")))?;
        out
    } else {
        data[wasm_start..].to_vec()
    };
    let expected = &data[28..60];
    let actual = sha256(&wasm);
    if expected != actual { return Err(BexError::HashMismatch { plugin_id: manifest.id.clone() }); }
    Ok(BexPackage { manifest, wasm })
}

pub fn pack(manifest: &Manifest, wasm: &[u8]) -> Result<Vec<u8>, BexError> {
    let yaml = serde_yaml::to_string(manifest).map_err(|e| BexError::Internal(e.to_string()))?;
    let yaml_bytes = yaml.as_bytes();
    let compressed = zstd::bulk::compress(wasm, 9).map_err(|e| BexError::Internal(e.to_string()))?;
    let mut out = Vec::with_capacity(HEADER_LEN + yaml_bytes.len() + compressed.len());
    out.extend_from_slice(MAGIC);
    out.extend_from_slice(&CONTAINER_VERSION.to_le_bytes());
    out.extend_from_slice(&1u16.to_le_bytes());
    out.extend_from_slice(&(yaml_bytes.len() as u32).to_le_bytes());
    out.extend_from_slice(&(compressed.len() as u64).to_le_bytes());
    out.extend_from_slice(&(wasm.len() as u64).to_le_bytes());
    out.extend_from_slice(&sha256(wasm));
    out.extend_from_slice(&sha256(yaml_bytes));
    let crc = crc32fast::hash(&out[0..92]);
    out.extend_from_slice(&crc.to_le_bytes());
    out.extend_from_slice(yaml_bytes);
    out.extend_from_slice(&compressed);
    Ok(out)
}

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

    fn test_manifest() -> Manifest {
        Manifest {
            schema: 1, id: "com.test.plugin".into(), name: "Test".into(),
            version: "1.0.0".into(), authors: vec!["dev".into()], abi: ">=1.0.0,<2.0.0".into(),
            provides: ProvidesSpec { search: true, info: true, ..Default::default() },
            network: NetworkSpec { hosts: vec!["*".into()], concurrent: 8 },
            storage: false, secrets: vec![],
            allow_js: false, allow_js_fetch: false,
            display: DisplaySpec { description: Some("test".into()), ..Default::default() },
        }
    }

    #[test]
    fn round_trip() {
        let m = test_manifest();
        let wasm = b"\x00asm\x01\x00\x00\x00";
        let packed = pack(&m, wasm).unwrap();
        let unpacked = unpack(&packed).unwrap();
        assert_eq!(unpacked.manifest.id, "com.test.plugin");
        assert_eq!(unpacked.wasm, wasm.to_vec());
    }

    #[test]
    fn read_manifest_from_packed() {
        let m = test_manifest();
        let packed = pack(&m, b"test").unwrap();
        let read = read_manifest(&packed).unwrap();
        assert_eq!(read.id, "com.test.plugin");
    }

    #[test]
    fn bad_magic() {
        let data = vec![0u8; 96];
        let r = read_manifest(&data);
        assert!(r.is_err());
    }
}