//! Core seal/unseal operations.

use aes_gcm::{
    aead::{Aead, AeadCore, KeyInit, OsRng},
    Aes256Gcm, Key, Nonce,
};
use rand::RngCore;

use crate::artifact::{DeploymentBinding, SealAlgorithm, SealConfig, SealedArtifact};
use crate::error::{SealError, SealResult};

// ── Public API ────────────────────────────────────────────────────────────────

/// Seal `bytecode` into a [`SealedArtifact`] using the given [`SealConfig`].
///
/// # Sealing steps
///
/// 1. Resolve the deployment binding material.
/// 2. Generate a random 256-bit symmetric key.
/// 3. Encrypt bytecode with AES-256-GCM.
/// 4. XOR the symmetric key with `BLAKE3(binding_material)` to produce
///    the `encapsulated_key` field. Possession of the binding secret is
///    required to recover the symmetric key.
/// 5. MAC the header + ciphertext with the symmetric key.
/// 6. Serialize into [`SealedArtifact`].
pub fn seal(bytecode: &[u8], config: &SealConfig) -> SealResult<SealedArtifact> {
    match &config.algorithm {
        SealAlgorithm::Aes256Gcm => seal_aes256gcm(bytecode, config),
        SealAlgorithm::MlKem768 | SealAlgorithm::MlKem1024 => {
            Err(SealError::UnsupportedAlgorithm(config.algorithm.id().to_string()))
        }
    }
}

/// Unseal a [`SealedArtifact`], recovering the original bytecode.
///
/// The `binding_key` must match the key that was used during sealing.
/// For [`DeploymentBinding::EnvironmentKey`], this is the raw env var bytes.
/// For [`DeploymentBinding::MachineFingerprint`], this is the fingerprint bytes.
/// For [`DeploymentBinding::None`], pass `&[]`.
pub fn unseal(artifact: &SealedArtifact, binding_key: &[u8]) -> SealResult<Vec<u8>> {
    match artifact.algorithm_id.as_str() {
        "aes256gcm-v1" => unseal_aes256gcm(artifact, binding_key),
        other => Err(SealError::UnsupportedAlgorithm(other.to_string())),
    }
}

/// Verify the MAC/signature on a [`SealedArtifact`] without decrypting.
///
/// Returns `true` if the artifact is intact. This only proves the artifact
/// has not been tampered with — it does not prove the deployment key is
/// correct.
///
/// Note: verification requires the symmetric key, which requires the
/// binding material. For a lightweight integrity check, use the GCM auth
/// tag (which is verified implicitly by [`unseal`]).
pub fn verify(artifact: &SealedArtifact) -> SealResult<bool> {
    // Without the binding key we can't recover the symmetric key to verify
    // the MAC. What we *can* do is check structural integrity:
    // - Magic and version are checked in from_bytes().
    // - Nonce must be 12 bytes (AES-GCM).
    // - Ciphertext must be non-empty.
    let structural_ok = artifact.nonce.len() == 12
        && !artifact.ciphertext.is_empty()
        && !artifact.encapsulated_key.is_empty()
        && !artifact.signature.is_empty();
    Ok(structural_ok)
}

// ── AES-256-GCM sealing ───────────────────────────────────────────────────────

fn seal_aes256gcm(bytecode: &[u8], config: &SealConfig) -> SealResult<SealedArtifact> {
    let algorithm_id = SealAlgorithm::Aes256Gcm.id().to_string();

    // 1. Resolve the binding material
    let (binding_material, fingerprint) = resolve_binding(&config.deployment_binding)?;

    // 2. Generate a random 256-bit symmetric key
    let mut sym_key = [0u8; 32];
    OsRng.fill_bytes(&mut sym_key);

    // 3. Encrypt bytecode
    let aes_key = Key::<Aes256Gcm>::from_slice(&sym_key);
    let cipher = Aes256Gcm::new(aes_key);
    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);

    let ciphertext = cipher
        .encrypt(&nonce, bytecode)
        .map_err(|e| SealError::EncryptionFailed(e.to_string()))?;

    // 4. Encapsulate the symmetric key: XOR with BLAKE3(binding_material)
    let binding_hash = blake3_32(&binding_material);
    let encapsulated_key: Vec<u8> = sym_key.iter().zip(binding_hash.iter()).map(|(a, b)| a ^ b).collect();

    // 5. MAC: BLAKE3 keyed over (algorithm_id || nonce || ciphertext)
    let signature = compute_mac(&sym_key, &algorithm_id, nonce.as_slice(), &ciphertext);

    Ok(SealedArtifact {
        algorithm_id,
        signature,
        encapsulated_key,
        nonce: nonce.to_vec(),
        ciphertext,
        deployment_fingerprint: fingerprint,
    })
}

fn unseal_aes256gcm(artifact: &SealedArtifact, binding_key: &[u8]) -> SealResult<Vec<u8>> {
    // 1. Derive binding hash from the provided key.
    // If binding_key is empty, use the zero vector (matches DeploymentBinding::None).
    let effective_key = if binding_key.is_empty() {
        vec![0u8; 32]
    } else {
        binding_key.to_vec()
    };
    let binding_hash = blake3_32(&effective_key);

    // 1b. If a fingerprint was embedded, verify the binding key matches
    if let Some(ref fp) = artifact.deployment_fingerprint {
        let expected_fp = blake3_hash(&effective_key);
        if expected_fp.as_slice() != fp.as_slice() {
            return Err(SealError::BindingMismatch);
        }
    }

    // 2. Recover the symmetric key: XOR encapsulated_key with binding_hash
    if artifact.encapsulated_key.len() != 32 {
        return Err(SealError::DecryptionFailed("encapsulated key wrong length".into()));
    }
    let sym_key: Vec<u8> = artifact.encapsulated_key.iter().zip(binding_hash.iter()).map(|(a, b)| a ^ b).collect();
    let sym_key_arr: [u8; 32] = sym_key.try_into().unwrap();

    // 3. Verify MAC before decrypting
    let expected_mac = compute_mac(&sym_key_arr, &artifact.algorithm_id, &artifact.nonce, &artifact.ciphertext);
    if expected_mac != artifact.signature {
        return Err(SealError::SignatureInvalid);
    }

    // 4. Decrypt
    if artifact.nonce.len() != 12 {
        return Err(SealError::DecryptionFailed("invalid nonce length".into()));
    }
    let nonce = Nonce::from_slice(&artifact.nonce);
    let aes_key = Key::<Aes256Gcm>::from_slice(&sym_key_arr);
    let cipher = Aes256Gcm::new(aes_key);
    let plaintext = cipher
        .decrypt(nonce, artifact.ciphertext.as_slice())
        .map_err(|e| SealError::DecryptionFailed(e.to_string()))?;

    Ok(plaintext)
}

// ── Binding resolution ────────────────────────────────────────────────────────

/// Resolve a deployment binding to raw bytes and an optional fingerprint.
///
/// Returns `(binding_material, deployment_fingerprint)`.
/// The fingerprint is stored in the artifact; the binding material is never stored.
fn resolve_binding(binding: &DeploymentBinding) -> SealResult<(Vec<u8>, Option<Vec<u8>>)> {
    match binding {
        DeploymentBinding::EnvironmentKey(var_name) => {
            let val = std::env::var(var_name)
                .map_err(|_| SealError::MissingEnvKey(var_name.clone()))?;
            let material = val.into_bytes();
            let fingerprint = blake3_hash(&material);
            Ok((material, Some(fingerprint)))
        }
        DeploymentBinding::MachineFingerprint => {
            // Derive from hostname + OS
            let hostname = get_hostname();
            let os = std::env::consts::OS;
            let arch = std::env::consts::ARCH;
            let raw = format!("{hostname}::{os}::{arch}");
            let material = raw.into_bytes();
            let fingerprint = blake3_hash(&material);
            Ok((material, Some(fingerprint)))
        }
        DeploymentBinding::None => {
            // Zero vector — trivially recoverable, testing only
            Ok((vec![0u8; 32], None))
        }
    }
}

fn get_hostname() -> String {
    std::env::var("HOSTNAME")
        .or_else(|_| std::env::var("COMPUTERNAME"))
        .unwrap_or_else(|_| "unknown-host".into())
}

// ── Crypto helpers ────────────────────────────────────────────────────────────

fn blake3_32(data: &[u8]) -> [u8; 32] {
    *blake3::hash(data).as_bytes()
}

fn blake3_hash(data: &[u8]) -> Vec<u8> {
    blake3::hash(data).as_bytes().to_vec()
}

fn compute_mac(key: &[u8; 32], algorithm_id: &str, nonce: &[u8], ciphertext: &[u8]) -> Vec<u8> {
    let mut hasher = blake3::Hasher::new_keyed(key);
    hasher.update(algorithm_id.as_bytes());
    hasher.update(nonce);
    hasher.update(ciphertext);
    hasher.finalize().as_bytes().to_vec()
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::artifact::{DeploymentBinding, SealAlgorithm, SealConfig};

    fn no_binding_config() -> SealConfig {
        SealConfig {
            algorithm: SealAlgorithm::Aes256Gcm,
            deployment_binding: DeploymentBinding::None,
        }
    }

    fn env_binding_config(var: &str) -> SealConfig {
        SealConfig {
            algorithm: SealAlgorithm::Aes256Gcm,
            deployment_binding: DeploymentBinding::EnvironmentKey(var.to_string()),
        }
    }

    #[test]
    fn test_seal_unseal_roundtrip_no_binding() {
        let bytecode = b"PUSH 42\nCALL print\nRETURN";
        let config = no_binding_config();
        let artifact = seal(bytecode, &config).unwrap();
        let recovered = unseal(&artifact, &[]).unwrap();
        assert_eq!(recovered, bytecode);
    }

    #[test]
    fn test_seal_unseal_roundtrip_env_key() {
        std::env::set_var("_EL_TEST_SEAL_KEY", "super-secret-deployment-key");
        let bytecode = b"sealed bytecode payload";
        let config = env_binding_config("_EL_TEST_SEAL_KEY");
        let artifact = seal(bytecode, &config).unwrap();
        let recovered = unseal(&artifact, b"super-secret-deployment-key").unwrap();
        assert_eq!(recovered, bytecode);
        std::env::remove_var("_EL_TEST_SEAL_KEY");
    }

    #[test]
    fn test_wrong_binding_key_rejected() {
        let bytecode = b"secret bytecode";
        let config = no_binding_config();
        let artifact = seal(bytecode, &config).unwrap();
        // Use wrong key — MAC should fail
        let mut bad_artifact = artifact.clone();
        bad_artifact.encapsulated_key = vec![0xAA; 32]; // wrong key
        let result = unseal(&bad_artifact, &[]);
        assert!(result.is_err());
    }

    #[test]
    fn test_tampered_ciphertext_rejected() {
        let bytecode = b"important bytecode";
        let config = no_binding_config();
        let mut artifact = seal(bytecode, &config).unwrap();
        // Flip a byte in the ciphertext
        if let Some(b) = artifact.ciphertext.first_mut() {
            *b ^= 0xFF;
        }
        let result = unseal(&artifact, &[]);
        assert!(result.is_err());
    }

    #[test]
    fn test_tampered_mac_rejected() {
        let bytecode = b"important bytecode";
        let config = no_binding_config();
        let mut artifact = seal(bytecode, &config).unwrap();
        // Flip the first byte of the MAC
        if let Some(b) = artifact.signature.first_mut() {
            *b ^= 0xFF;
        }
        let result = unseal(&artifact, &[]);
        assert!(result.is_err());
    }

    #[test]
    fn test_serialization_roundtrip() {
        let bytecode = b"fn main() { return 42 }";
        let config = no_binding_config();
        let artifact = seal(bytecode, &config).unwrap();
        let bytes = artifact.to_bytes().unwrap();
        let restored = SealedArtifact::from_bytes(&bytes).unwrap();
        let recovered = unseal(&restored, &[]).unwrap();
        assert_eq!(recovered, bytecode);
    }

    #[test]
    fn test_magic_header_present() {
        let artifact = seal(b"test", &no_binding_config()).unwrap();
        let bytes = artifact.to_bytes().unwrap();
        assert_eq!(&bytes[..8], b"ENGRAM01");
    }

    #[test]
    fn test_wrong_magic_rejected() {
        let mut bytes = seal(b"test", &no_binding_config()).unwrap().to_bytes().unwrap();
        bytes[0] = 0xFF; // corrupt magic
        let result = SealedArtifact::from_bytes(&bytes);
        assert!(result.is_err());
    }

    #[test]
    fn test_verify_structural_ok() {
        let artifact = seal(b"bytecode", &no_binding_config()).unwrap();
        assert!(verify(&artifact).unwrap());
    }

    #[test]
    fn test_empty_bytecode_sealable() {
        let artifact = seal(b"", &no_binding_config()).unwrap();
        let recovered = unseal(&artifact, &[]).unwrap();
        assert_eq!(recovered, b"");
    }

    #[test]
    fn test_large_bytecode_sealable() {
        let bytecode = vec![0x42u8; 100_000];
        let artifact = seal(&bytecode, &no_binding_config()).unwrap();
        let recovered = unseal(&artifact, &[]).unwrap();
        assert_eq!(recovered, bytecode);
    }

    #[test]
    fn test_algorithm_id_stored() {
        let artifact = seal(b"test", &no_binding_config()).unwrap();
        assert_eq!(artifact.algorithm_id, "aes256gcm-v1");
    }

    #[test]
    fn test_nonce_is_12_bytes() {
        let artifact = seal(b"test", &no_binding_config()).unwrap();
        assert_eq!(artifact.nonce.len(), 12);
    }

    #[test]
    fn test_encapsulated_key_is_32_bytes() {
        let artifact = seal(b"test", &no_binding_config()).unwrap();
        assert_eq!(artifact.encapsulated_key.len(), 32);
    }

    #[test]
    fn test_different_seals_produce_different_ciphertexts() {
        let bytecode = b"same input";
        let config = no_binding_config();
        let a1 = seal(bytecode, &config).unwrap();
        let a2 = seal(bytecode, &config).unwrap();
        // Random nonce means ciphertexts differ
        assert_ne!(a1.ciphertext, a2.ciphertext);
        assert_ne!(a1.nonce, a2.nonce);
    }

    #[test]
    fn test_missing_env_key_returns_error() {
        std::env::remove_var("_EL_NONEXISTENT_KEY");
        let result = seal(b"test", &env_binding_config("_EL_NONEXISTENT_KEY"));
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), SealError::MissingEnvKey(_)));
    }
}