//! # ECASM - Encrypted CASM
//!
//! Encrypted bytecode format for confidential execution.
//!
//! ## Format Structure
//!
//! ```text
//! ECASM v1 Binary Layout
//! +-----------------------------------------------------------+
//! | Header (plaintext, 88 bytes)                              |
//! +-----------------------------------------------------------+
//! |  magic: [u8; 4]        = "ECSM" (0x4543534D)             |
//! |  version: u16          = 1                                |
//! |  flags: u16            = EcasmFlags                       |
//! |  vm_version_min: u32   = minimum VM version required      |
//! |  vm_version_max: u32   = maximum VM version supported     |
//! |  crypto_suite: u8      = 0x01 (ChaCha20-Poly1305)        |
//! |  reserved: [u8; 3]                                        |
//! |  capsule_id: [u8; 32]  = SHA256(capsule manifest)         |
//! |  caps_hash: [u8; 32]   = SHA256(required capabilities)    |
//! |  page_size: u32        = bytes per page (default 4096)    |
//! |  page_count: u32       = number of encrypted pages        |
//! +-----------------------------------------------------------+
//! | Page Table (plaintext, variable)                          |
//! +-----------------------------------------------------------+
//! |  For each page:                                           |
//! |    offset: u64         = byte offset in payload           |
//! |    size: u32           = encrypted size (with tag)        |
//! |    hash: [u8; 32]      = SHA256(ciphertext)              |
//! +-----------------------------------------------------------+
//! | Encrypted Payload                                         |
//! +-----------------------------------------------------------+
//! |  Page 0: [ciphertext || auth_tag (16 bytes)]             |
//! |  Page 1: [ciphertext || auth_tag (16 bytes)]             |
//! |  ...                                                      |
//! +-----------------------------------------------------------+
//! | Footer                                                    |
//! +-----------------------------------------------------------+
//! |  header_hash: [u8; 32] = SHA256(header || page_table)    |
//! |  signature: [u8; 64]   = Ed25519 (if SIGNED flag)        |
//! +-----------------------------------------------------------+
//! ```

use bitflags::bitflags;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use sha2::{Digest, Sha256};
use std::io::{self, Read, Write};

use crate::{Format, Program};
use crush_errors::CrushResult;
use crush_errors::convert::casm::CasmError;

/// ECASM file magic bytes: "ECSM"
pub const ECASM_MAGIC: [u8; 4] = [0x45, 0x43, 0x53, 0x4D];

/// Current ECASM format version
pub const ECASM_VERSION: u16 = 1;

/// Default page size (4KB)
pub const DEFAULT_PAGE_SIZE: u32 = 4096;

/// ChaCha20-Poly1305 crypto suite identifier
pub const CRYPTO_SUITE_CHACHA20_POLY1305: u8 = 0x01;

/// Authentication tag size (16 bytes for Poly1305)
pub const AUTH_TAG_SIZE: usize = 16;

/// Header size in bytes (fixed)
/// 4 (magic) + 2 (version) + 2 (flags) + 4 (vm_min) + 4 (vm_max) +
/// 1 (crypto_suite) + 3 (reserved) + 32 (capsule_id) + 32 (caps_hash) +
/// 4 (page_size) + 4 (page_count) = 92 bytes
pub const HEADER_SIZE: usize = 92;

/// Page table entry size in bytes
pub const PAGE_TABLE_ENTRY_SIZE: usize = 44; // 8 + 4 + 32

bitflags! {
    /// ECASM format flags
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub struct EcasmFlags: u16 {
        /// Ed25519 signature present in footer
        const SIGNED = 0b0000_0001;
        /// Pages are zstd-compressed before encryption
        const COMPRESSED = 0b0000_0010;
        /// Encrypted debug symbols in separate section
        const DEBUG_INFO = 0b0000_0100;
        /// Supports partial/streaming load
        const STREAMING = 0b0000_1000;
    }
}

/// ECASM file header (plaintext)
#[derive(Debug, Clone)]
pub struct EcasmHeader {
    /// Format version
    pub version: u16,
    /// Format flags
    pub flags: EcasmFlags,
    /// Minimum VM version required
    pub vm_version_min: u32,
    /// Maximum VM version supported
    pub vm_version_max: u32,
    /// Crypto suite identifier
    pub crypto_suite: u8,
    /// Capsule ID (SHA256 of manifest)
    pub capsule_id: [u8; 32],
    /// Hash of required capabilities
    pub caps_hash: [u8; 32],
    /// Size of each plaintext page in bytes
    pub page_size: u32,
    /// Number of encrypted pages
    pub page_count: u32,
}

impl Default for EcasmHeader {
    fn default() -> Self {
        Self {
            version: ECASM_VERSION,
            flags: EcasmFlags::empty(),
            vm_version_min: 1,
            vm_version_max: u32::MAX,
            crypto_suite: CRYPTO_SUITE_CHACHA20_POLY1305,
            capsule_id: [0u8; 32],
            caps_hash: [0u8; 32],
            page_size: DEFAULT_PAGE_SIZE,
            page_count: 0,
        }
    }
}

impl EcasmHeader {
    /// Create a new header with default values
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the capsule ID from a manifest hash
    pub fn with_capsule_id(mut self, id: [u8; 32]) -> Self {
        self.capsule_id = id;
        self
    }

    /// Set capabilities hash
    pub fn with_caps_hash(mut self, hash: [u8; 32]) -> Self {
        self.caps_hash = hash;
        self
    }

    /// Set page size
    pub fn with_page_size(mut self, size: u32) -> Self {
        self.page_size = size;
        self
    }

    /// Set flags
    pub fn with_flags(mut self, flags: EcasmFlags) -> Self {
        self.flags = flags;
        self
    }

    /// Serialize header to bytes (88 bytes)
    pub fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        // Magic
        writer.write_all(&ECASM_MAGIC)?;
        // Version
        writer.write_u16::<LittleEndian>(self.version)?;
        // Flags
        writer.write_u16::<LittleEndian>(self.flags.bits())?;
        // VM versions
        writer.write_u32::<LittleEndian>(self.vm_version_min)?;
        writer.write_u32::<LittleEndian>(self.vm_version_max)?;
        // Crypto suite + reserved
        writer.write_u8(self.crypto_suite)?;
        writer.write_all(&[0u8; 3])?; // reserved
        // Capsule ID
        writer.write_all(&self.capsule_id)?;
        // Caps hash
        writer.write_all(&self.caps_hash)?;
        // Page info
        writer.write_u32::<LittleEndian>(self.page_size)?;
        writer.write_u32::<LittleEndian>(self.page_count)?;

        Ok(())
    }

    /// Deserialize header from bytes
    pub fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
        // Magic
        let mut magic = [0u8; 4];
        reader.read_exact(&mut magic)?;
        if magic != ECASM_MAGIC {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!(
                    "invalid ECASM magic: expected {:?}, got {:?}",
                    ECASM_MAGIC, magic
                ),
            ));
        }

        // Version
        let version = reader.read_u16::<LittleEndian>()?;
        if version != ECASM_VERSION {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("unsupported ECASM version: {}", version),
            ));
        }

        // Flags
        let flags_bits = reader.read_u16::<LittleEndian>()?;
        let flags = EcasmFlags::from_bits_truncate(flags_bits);

        // VM versions
        let vm_version_min = reader.read_u32::<LittleEndian>()?;
        let vm_version_max = reader.read_u32::<LittleEndian>()?;

        // Crypto suite + reserved
        let crypto_suite = reader.read_u8()?;
        let mut reserved = [0u8; 3];
        reader.read_exact(&mut reserved)?;

        // Capsule ID
        let mut capsule_id = [0u8; 32];
        reader.read_exact(&mut capsule_id)?;

        // Caps hash
        let mut caps_hash = [0u8; 32];
        reader.read_exact(&mut caps_hash)?;

        // Page info
        let page_size = reader.read_u32::<LittleEndian>()?;
        let page_count = reader.read_u32::<LittleEndian>()?;

        Ok(Self {
            version,
            flags,
            vm_version_min,
            vm_version_max,
            crypto_suite,
            capsule_id,
            caps_hash,
            page_size,
            page_count,
        })
    }

    /// Validate header against current VM version
    pub fn validate(&self, vm_version: u32) -> Result<(), EcasmError> {
        if vm_version < self.vm_version_min {
            return Err(EcasmError::VmVersionTooOld {
                required: self.vm_version_min,
                actual: vm_version,
            });
        }
        if vm_version > self.vm_version_max {
            return Err(EcasmError::VmVersionTooNew {
                max_supported: self.vm_version_max,
                actual: vm_version,
            });
        }
        if self.crypto_suite != CRYPTO_SUITE_CHACHA20_POLY1305 {
            return Err(EcasmError::UnsupportedCryptoSuite(self.crypto_suite));
        }
        Ok(())
    }
}

/// Page table entry
#[derive(Debug, Clone)]
pub struct PageTableEntry {
    /// Byte offset of this page in the payload section
    pub offset: u64,
    /// Size of encrypted data (including auth tag)
    pub size: u32,
    /// SHA256 hash of the ciphertext
    pub hash: [u8; 32],
}

impl PageTableEntry {
    /// Serialize entry to bytes
    pub fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        writer.write_u64::<LittleEndian>(self.offset)?;
        writer.write_u32::<LittleEndian>(self.size)?;
        writer.write_all(&self.hash)?;
        Ok(())
    }

    /// Deserialize entry from bytes
    pub fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
        let offset = reader.read_u64::<LittleEndian>()?;
        let size = reader.read_u32::<LittleEndian>()?;
        let mut hash = [0u8; 32];
        reader.read_exact(&mut hash)?;
        Ok(Self { offset, size, hash })
    }
}

/// Complete ECASM file structure
#[derive(Debug, Clone)]
pub struct EcasmFile {
    /// File header
    pub header: EcasmHeader,
    /// Page table entries
    pub page_table: Vec<PageTableEntry>,
    /// Encrypted pages (ciphertext + auth tag)
    pub pages: Vec<Vec<u8>>,
    /// Hash of header + page table (for integrity)
    pub header_hash: [u8; 32],
    /// Optional Ed25519 signature
    pub signature: Option<[u8; 64]>,
}

impl EcasmFile {
    /// Create a new ECASM file from plaintext CASM
    ///
    /// # Arguments
    /// * `program` - The CASM program to encrypt
    /// * `capsule_id` - SHA256 of the capsule manifest
    /// * `caps_hash` - SHA256 of required capabilities
    /// * `encrypt_page` - Closure to encrypt a page: (page_index, plaintext) -> ciphertext
    pub fn encrypt<F>(
        program: &Program,
        capsule_id: [u8; 32],
        caps_hash: [u8; 32],
        page_size: u32,
        mut encrypt_page: F,
    ) -> CrushResult<Self>
    where
        F: FnMut(u32, &[u8]) -> CrushResult<Vec<u8>>,
    {
        // Serialize program to binary format
        let plaintext = program.serialize(Format::Binary)?;

        // Split into pages
        let mut pages = Vec::new();
        let mut page_table = Vec::new();
        let mut offset = 0u64;

        for (page_index, chunk) in plaintext.chunks(page_size as usize).enumerate() {
            // Encrypt the page
            let ciphertext = encrypt_page(page_index as u32, chunk)?;

            // Compute hash of ciphertext
            let mut hasher = Sha256::new();
            hasher.update(&ciphertext);
            let hash: [u8; 32] = hasher.finalize().into();

            // Create page table entry
            page_table.push(PageTableEntry {
                offset,
                size: ciphertext.len() as u32,
                hash,
            });

            offset += ciphertext.len() as u64;
            pages.push(ciphertext);
        }

        // Create header
        let header = EcasmHeader {
            version: ECASM_VERSION,
            flags: EcasmFlags::empty(),
            vm_version_min: 1,
            vm_version_max: u32::MAX,
            crypto_suite: CRYPTO_SUITE_CHACHA20_POLY1305,
            capsule_id,
            caps_hash,
            page_size,
            page_count: pages.len() as u32,
        };

        // Compute header hash
        let header_hash = Self::compute_header_hash(&header, &page_table)?;

        Ok(Self {
            header,
            page_table,
            pages,
            header_hash,
            signature: None,
        })
    }

    /// Compute SHA256 hash of header + page table
    fn compute_header_hash(
        header: &EcasmHeader,
        page_table: &[PageTableEntry],
    ) -> CrushResult<[u8; 32]> {
        let mut hasher = Sha256::new();

        // Hash header
        let mut header_bytes = Vec::with_capacity(HEADER_SIZE);
        header
            .serialize(&mut header_bytes)
            .map_err(|e| CasmError::IoError(e.to_string()))?;
        hasher.update(&header_bytes);

        // Hash page table
        for entry in page_table {
            let mut entry_bytes = Vec::with_capacity(PAGE_TABLE_ENTRY_SIZE);
            entry
                .serialize(&mut entry_bytes)
                .map_err(|e| CasmError::IoError(e.to_string()))?;
            hasher.update(&entry_bytes);
        }

        Ok(hasher.finalize().into())
    }

    /// Serialize the complete ECASM file to bytes
    pub fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
        // Write header
        self.header.serialize(writer)?;

        // Write page table
        for entry in &self.page_table {
            entry.serialize(writer)?;
        }

        // Write encrypted pages
        for page in &self.pages {
            writer.write_all(page)?;
        }

        // Write footer
        writer.write_all(&self.header_hash)?;
        if self.header.flags.contains(EcasmFlags::SIGNED)
            && let Some(sig) = &self.signature
        {
            writer.write_all(sig)?;
        }

        Ok(())
    }

    /// Deserialize ECASM file from bytes (does not decrypt)
    pub fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
        // Read header
        let header = EcasmHeader::deserialize(reader)?;

        // Read page table
        let mut page_table = Vec::with_capacity(header.page_count as usize);
        for _ in 0..header.page_count {
            page_table.push(PageTableEntry::deserialize(reader)?);
        }

        // Read encrypted pages
        let mut pages = Vec::with_capacity(header.page_count as usize);
        for entry in &page_table {
            let mut page = vec![0u8; entry.size as usize];
            reader.read_exact(&mut page)?;
            pages.push(page);
        }

        // Read footer
        let mut header_hash = [0u8; 32];
        reader.read_exact(&mut header_hash)?;

        let signature = if header.flags.contains(EcasmFlags::SIGNED) {
            let mut sig = [0u8; 64];
            reader.read_exact(&mut sig)?;
            Some(sig)
        } else {
            None
        };

        Ok(Self {
            header,
            page_table,
            pages,
            header_hash,
            signature,
        })
    }

    /// Verify integrity of the file (does not decrypt)
    pub fn verify_integrity(&self) -> CrushResult<()> {
        // Verify header hash
        let computed_hash = Self::compute_header_hash(&self.header, &self.page_table)?;
        if computed_hash != self.header_hash {
            return Err(CasmError::IntegrityError("header hash mismatch".to_string()).into());
        }

        // Verify page hashes
        for (i, (entry, page)) in self.page_table.iter().zip(&self.pages).enumerate() {
            let mut hasher = Sha256::new();
            hasher.update(page);
            let hash: [u8; 32] = hasher.finalize().into();

            if hash != entry.hash {
                return Err(CasmError::IntegrityError(format!("page {} hash mismatch", i)).into());
            }
        }

        Ok(())
    }

    /// Decrypt a single page
    ///
    /// # Arguments
    /// * `page_index` - Index of the page to decrypt
    /// * `decrypt` - Closure to decrypt: (page_index, ciphertext) -> plaintext
    pub fn decrypt_page<F>(&self, page_index: u32, decrypt: F) -> CrushResult<Vec<u8>>
    where
        F: FnOnce(u32, &[u8]) -> CrushResult<Vec<u8>>,
    {
        let page = self
            .pages
            .get(page_index as usize)
            .ok_or(CasmError::InvalidPageIndex {
                index: page_index,
                max: self.header.page_count,
            })?;

        decrypt(page_index, page)
    }

    /// Decrypt all pages and reconstruct the CASM program
    ///
    /// # Arguments
    /// * `decrypt_page` - Closure to decrypt: (page_index, ciphertext) -> plaintext
    pub fn decrypt<F>(&self, mut decrypt_page: F) -> CrushResult<Program>
    where
        F: FnMut(u32, &[u8]) -> CrushResult<Vec<u8>>,
    {
        // Decrypt all pages
        let mut plaintext = Vec::new();
        for (i, page) in self.pages.iter().enumerate() {
            let decrypted = decrypt_page(i as u32, page)?;
            plaintext.extend_from_slice(&decrypted);
        }

        // Deserialize program
        Program::deserialize(&plaintext, Format::Binary)
    }

    /// Get total file size in bytes
    pub fn file_size(&self) -> usize {
        let mut size = HEADER_SIZE;
        size += self.page_table.len() * PAGE_TABLE_ENTRY_SIZE;
        size += self.pages.iter().map(|p| p.len()).sum::<usize>();
        size += 32; // header_hash
        if self.header.flags.contains(EcasmFlags::SIGNED) {
            size += 64; // signature
        }
        size
    }
}

/// ECASM-specific errors
#[derive(Debug, Clone)]
pub enum EcasmError {
    /// VM version is too old
    VmVersionTooOld { required: u32, actual: u32 },
    /// VM version is too new
    VmVersionTooNew { max_supported: u32, actual: u32 },
    /// Unsupported crypto suite
    UnsupportedCryptoSuite(u8),
    /// Page index out of bounds
    InvalidPageIndex { index: u32, max: u32 },
    /// Decryption failed
    DecryptionFailed(String),
}

impl std::fmt::Display for EcasmError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::VmVersionTooOld { required, actual } => {
                write!(f, "VM version {} is too old, requires {}", actual, required)
            }
            Self::VmVersionTooNew {
                max_supported,
                actual,
            } => {
                write!(
                    f,
                    "VM version {} is too new, max supported is {}",
                    actual, max_supported
                )
            }
            Self::UnsupportedCryptoSuite(id) => {
                write!(f, "unsupported crypto suite: 0x{:02x}", id)
            }
            Self::InvalidPageIndex { index, max } => {
                write!(f, "page index {} out of bounds (max {})", index, max)
            }
            Self::DecryptionFailed(msg) => write!(f, "decryption failed: {}", msg),
        }
    }
}

impl std::error::Error for EcasmError {}

/// Compute SHA256 hash of capability list
pub fn hash_capabilities(caps: &[String]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    // Sort for deterministic hashing
    let mut sorted: Vec<_> = caps.iter().collect();
    sorted.sort();
    for cap in sorted {
        hasher.update(cap.as_bytes());
        hasher.update(b"\0"); // null separator
    }
    hasher.finalize().into()
}

/// Compute SHA256 hash of arbitrary data (for capsule manifest, etc.)
pub fn sha256_hash(data: &[u8]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(data);
    hasher.finalize().into()
}

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

    fn sample_program() -> Program {
        let mut functions = HashMap::new();
        functions.insert(
            "main".to_string(),
            crate::Function {
                params: vec![],
                locals: vec![],
                body: vec![
                    crate::Instruction {
                        op: "push_str".to_string(),
                        lang: None,
                        meta: None,
                        args: serde_json::json!({"value": "Hello, ECASM!"}),
                    },
                    crate::Instruction {
                        op: "cap_call".to_string(),
                        lang: None,
                        meta: None,
                        args: serde_json::json!({"name": "io.print", "argc": 1}),
                    },
                    crate::Instruction {
                        op: "ret".to_string(),
                        lang: None,
                        meta: None,
                        args: serde_json::json!({}),
                    },
                ],
            },
        );

        Program {
            version: "1.0".to_string(),
            functions,
            lang: None,
            manifest: crate::Manifest {
                permissions: vec!["io.print".to_string()],
            },
        }
    }

    #[test]
    fn test_header_roundtrip() {
        let header = EcasmHeader {
            version: ECASM_VERSION,
            flags: EcasmFlags::SIGNED | EcasmFlags::COMPRESSED,
            vm_version_min: 1,
            vm_version_max: 100,
            crypto_suite: CRYPTO_SUITE_CHACHA20_POLY1305,
            capsule_id: [0x42; 32],
            caps_hash: [0x99; 32],
            page_size: 4096,
            page_count: 10,
        };

        let mut buf = Vec::new();
        header.serialize(&mut buf).unwrap();
        assert_eq!(buf.len(), HEADER_SIZE);

        let parsed = EcasmHeader::deserialize(&mut buf.as_slice()).unwrap();
        assert_eq!(parsed.version, header.version);
        assert_eq!(parsed.flags, header.flags);
        assert_eq!(parsed.vm_version_min, header.vm_version_min);
        assert_eq!(parsed.vm_version_max, header.vm_version_max);
        assert_eq!(parsed.crypto_suite, header.crypto_suite);
        assert_eq!(parsed.capsule_id, header.capsule_id);
        assert_eq!(parsed.caps_hash, header.caps_hash);
        assert_eq!(parsed.page_size, header.page_size);
        assert_eq!(parsed.page_count, header.page_count);
    }

    #[test]
    fn test_page_table_entry_roundtrip() {
        let entry = PageTableEntry {
            offset: 12345,
            size: 4112, // 4096 + 16 tag
            hash: [0xAB; 32],
        };

        let mut buf = Vec::new();
        entry.serialize(&mut buf).unwrap();
        assert_eq!(buf.len(), PAGE_TABLE_ENTRY_SIZE);

        let parsed = PageTableEntry::deserialize(&mut buf.as_slice()).unwrap();
        assert_eq!(parsed.offset, entry.offset);
        assert_eq!(parsed.size, entry.size);
        assert_eq!(parsed.hash, entry.hash);
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let program = sample_program();
        let capsule_id = [0x11; 32];
        let caps_hash = hash_capabilities(&["io.print".to_string()]);

        // Simple XOR "encryption" for testing (not secure!)
        let key = [0x42u8; 32];

        let ecasm = EcasmFile::encrypt(
            &program,
            capsule_id,
            caps_hash,
            DEFAULT_PAGE_SIZE,
            |_idx, plaintext| {
                let mut ciphertext: Vec<u8> = plaintext
                    .iter()
                    .enumerate()
                    .map(|(i, b)| b ^ key[i % 32])
                    .collect();
                // Add fake auth tag
                ciphertext.extend_from_slice(&[0u8; AUTH_TAG_SIZE]);
                Ok(ciphertext)
            },
        )
        .unwrap();

        // Verify structure
        assert_eq!(ecasm.header.capsule_id, capsule_id);
        assert_eq!(ecasm.header.caps_hash, caps_hash);
        assert!(ecasm.header.page_count > 0);
        assert_eq!(ecasm.pages.len(), ecasm.header.page_count as usize);

        // Verify integrity
        ecasm.verify_integrity().unwrap();

        // Decrypt
        let decrypted = ecasm
            .decrypt(|_idx, ciphertext| {
                // Strip auth tag
                let data = &ciphertext[..ciphertext.len() - AUTH_TAG_SIZE];
                let plaintext: Vec<u8> = data
                    .iter()
                    .enumerate()
                    .map(|(i, b)| b ^ key[i % 32])
                    .collect();
                Ok(plaintext)
            })
            .unwrap();

        // Compare programs
        assert_eq!(decrypted.version, program.version);
        assert_eq!(decrypted.functions.len(), program.functions.len());
        assert!(decrypted.functions.contains_key("main"));
    }

    #[test]
    fn test_file_serialize_deserialize() {
        let program = sample_program();
        let key = [0x42u8; 32];

        let ecasm = EcasmFile::encrypt(
            &program,
            [0x11; 32],
            [0x22; 32],
            DEFAULT_PAGE_SIZE,
            |_idx, plaintext| {
                let mut ct: Vec<u8> = plaintext
                    .iter()
                    .enumerate()
                    .map(|(i, b)| b ^ key[i % 32])
                    .collect();
                ct.extend_from_slice(&[0u8; AUTH_TAG_SIZE]);
                Ok(ct)
            },
        )
        .unwrap();

        // Serialize
        let mut buf = Vec::new();
        ecasm.serialize(&mut buf).unwrap();

        // Deserialize
        let parsed = EcasmFile::deserialize(&mut buf.as_slice()).unwrap();

        // Verify
        assert_eq!(parsed.header.version, ecasm.header.version);
        assert_eq!(parsed.header.page_count, ecasm.header.page_count);
        assert_eq!(parsed.pages.len(), ecasm.pages.len());
        assert_eq!(parsed.header_hash, ecasm.header_hash);

        // Verify integrity
        parsed.verify_integrity().unwrap();
    }

    #[test]
    fn test_integrity_failure_on_tamper() {
        let program = sample_program();
        let key = [0x42u8; 32];

        let mut ecasm = EcasmFile::encrypt(
            &program,
            [0x11; 32],
            [0x22; 32],
            DEFAULT_PAGE_SIZE,
            |_idx, plaintext| {
                let mut ct: Vec<u8> = plaintext
                    .iter()
                    .enumerate()
                    .map(|(i, b)| b ^ key[i % 32])
                    .collect();
                ct.extend_from_slice(&[0u8; AUTH_TAG_SIZE]);
                Ok(ct)
            },
        )
        .unwrap();

        // Tamper with page
        if let Some(page) = ecasm.pages.first_mut() {
            page[0] ^= 0xFF;
        }

        // Integrity check should fail
        assert!(ecasm.verify_integrity().is_err());
    }

    #[test]
    fn test_header_validation() {
        let header = EcasmHeader::default();

        // Valid VM version
        assert!(header.validate(1).is_ok());
        assert!(header.validate(100).is_ok());

        // Too old
        let old_header = EcasmHeader {
            vm_version_min: 10,
            ..Default::default()
        };
        assert!(matches!(
            old_header.validate(5),
            Err(EcasmError::VmVersionTooOld { .. })
        ));

        // Too new
        let new_header = EcasmHeader {
            vm_version_max: 10,
            ..Default::default()
        };
        assert!(matches!(
            new_header.validate(15),
            Err(EcasmError::VmVersionTooNew { .. })
        ));
    }

    #[test]
    fn test_hash_capabilities() {
        let caps1 = vec!["io.print".to_string(), "fs.read".to_string()];
        let caps2 = vec!["fs.read".to_string(), "io.print".to_string()];

        // Order shouldn't matter (sorted internally)
        let hash1 = hash_capabilities(&caps1);
        let hash2 = hash_capabilities(&caps2);
        assert_eq!(hash1, hash2);

        // Different caps should produce different hash
        let caps3 = vec!["io.print".to_string()];
        let hash3 = hash_capabilities(&caps3);
        assert_ne!(hash1, hash3);
    }

    /// A self-contained, deterministic authenticated cipher used to exercise the
    /// `EcasmFile` encrypt/decrypt *closure contract* without reaching up into the
    /// `exo-crypto` layer. casm is crypto-agnostic by design — the caller injects the
    /// cipher via closures — so casm's own unit tests must respect that dependency
    /// boundary (a `core/vm` crate must not depend on the `exo/` platform layer). Real
    /// ChaCha20-Poly1305 is the concern of the `exo-crypto` crate (which has its own
    /// cipher tests) and the `exo-core`/`exo-cli` layers that wire it into ecasm. This
    /// stub mirrors AEAD semantics (keystream XOR + a key-bound tag) so the wrong-key
    /// authentication path is genuinely exercised.
    struct TestCipher {
        key: [u8; 32],
    }

    impl TestCipher {
        fn new(key: [u8; 32]) -> Self {
            Self { key }
        }

        fn keystream(&self, page_idx: u32, len: usize) -> Vec<u8> {
            let mut seed = Vec::with_capacity(self.key.len() + 4);
            seed.extend_from_slice(&self.key);
            seed.extend_from_slice(&page_idx.to_le_bytes());
            let mut out = Vec::with_capacity(len);
            let mut block = sha256_hash(&seed);
            while out.len() < len {
                out.extend_from_slice(&block);
                block = sha256_hash(&block);
            }
            out.truncate(len);
            out
        }

        fn tag(&self, page_idx: u32, plaintext: &[u8]) -> [u8; 32] {
            let mut buf = Vec::with_capacity(self.key.len() + 4 + plaintext.len());
            buf.extend_from_slice(&self.key);
            buf.extend_from_slice(&page_idx.to_le_bytes());
            buf.extend_from_slice(plaintext);
            sha256_hash(&buf)
        }

        /// `ciphertext = (plaintext XOR keystream) || tag` — tag binds key + page + plaintext.
        fn encrypt_page(&self, page_idx: u32, plaintext: &[u8]) -> Vec<u8> {
            let ks = self.keystream(page_idx, plaintext.len());
            let mut ct: Vec<u8> = plaintext.iter().zip(&ks).map(|(p, k)| p ^ k).collect();
            ct.extend_from_slice(&self.tag(page_idx, plaintext));
            ct
        }

        /// Decrypt and authenticate; returns `Err` on a tag mismatch (e.g. a wrong key).
        fn decrypt_page(&self, page_idx: u32, ciphertext: &[u8]) -> Result<Vec<u8>, String> {
            if ciphertext.len() < 32 {
                return Err("ciphertext shorter than auth tag".to_string());
            }
            let (body, tag) = ciphertext.split_at(ciphertext.len() - 32);
            let ks = self.keystream(page_idx, body.len());
            let plaintext: Vec<u8> = body.iter().zip(&ks).map(|(c, k)| c ^ k).collect();
            if self.tag(page_idx, &plaintext) != tag {
                return Err("authentication failed".to_string());
            }
            Ok(plaintext)
        }
    }

    /// Round-trip an ecasm file through the encrypt/decrypt closure contract.
    #[test]
    fn test_ecasm_encrypt_decrypt_roundtrip() {
        let program = sample_program();
        let cipher = TestCipher::new([0x11; 32]);

        let capsule_id = sha256_hash(b"test.capsule.v1");
        let caps_hash = hash_capabilities(&["io.print".to_string()]);

        let ecasm = EcasmFile::encrypt(
            &program,
            capsule_id,
            caps_hash,
            DEFAULT_PAGE_SIZE,
            |page_idx, plaintext| Ok(cipher.encrypt_page(page_idx, plaintext)),
        )
        .unwrap();

        // Verify integrity
        ecasm.verify_integrity().unwrap();

        // Serialize then deserialize
        let mut bytes = Vec::new();
        ecasm.serialize(&mut bytes).unwrap();
        let loaded = EcasmFile::deserialize(&mut bytes.as_slice()).unwrap();
        loaded.verify_integrity().unwrap();

        let decrypted = loaded
            .decrypt(|page_idx, ciphertext| {
                cipher.decrypt_page(page_idx, ciphertext).map_err(|e| {
                    crush_errors::CrushError::new(crush_errors::ErrorKind::Internal, e)
                })
            })
            .unwrap();

        // Verify decrypted program matches original
        assert_eq!(decrypted.version, program.version);
        assert_eq!(decrypted.functions.len(), program.functions.len());
        assert!(decrypted.functions.contains_key("main"));

        // Verify main function body
        let orig_main = program.functions.get("main").unwrap();
        let decr_main = decrypted.functions.get("main").unwrap();
        assert_eq!(orig_main.body.len(), decr_main.body.len());
    }

    /// Decryption with a different key must fail authentication.
    #[test]
    fn test_ecasm_wrong_key_fails() {
        let program = sample_program();
        let cipher1 = TestCipher::new([0x11; 32]);

        let ecasm = EcasmFile::encrypt(
            &program,
            [0u8; 32],
            [0u8; 32],
            DEFAULT_PAGE_SIZE,
            |page_idx, plaintext| Ok(cipher1.encrypt_page(page_idx, plaintext)),
        )
        .unwrap();

        // Try to decrypt with a different key
        let cipher2 = TestCipher::new([0x22; 32]);
        let result = ecasm.decrypt(|page_idx, ciphertext| {
            cipher2
                .decrypt_page(page_idx, ciphertext)
                .map_err(|e| crush_errors::CrushError::new(crush_errors::ErrorKind::Internal, e))
        });

        // Should fail with authentication error
        assert!(result.is_err());
    }
}
