use ed25519_dalek::{Signature, Verifier, VerifyingKey};

#[cfg(any(test, feature = "signing"))]
use ed25519_dalek::{Signer, SigningKey};

use sha2::{Digest, Sha384};

#[cfg(any(test, feature = "signing"))]
use sha2::{
    digest::{generic_array::GenericArray, OutputSizeUser},
    Sha256,
};

use serde::{Deserialize, Serialize};

#[derive(thiserror::Error, Debug)]
pub enum SignatureVerificationError {
    #[error("Failed to verify signature: {0}")]
    SignatureError(#[from] ed25519_dalek::SignatureError),
    #[error("The specified key is unknown")]
    UnknownKeyId,
}

/// Launch measurement carried in a [`PublishedVersionEntry`].
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(
    untagged,
    expecting = "a base64-encoded string or a launch measurement object"
)]
pub enum LaunchMeasurement {
    /// Generation 1: flat SHA-384 digest including VMSA pages.
    V1(#[serde(with = "base64_compat")] Vec<u8>),
    /// Generation 2: pre-VMSA digest and raw VMSA bytes stored separately so
    /// the vCPU count can change without invalidating the signed system version.
    V2 {
        /// Raw [`oak_sev_guest::vmsa::Vmsa`] bytes shared by all vCPUs.
        #[serde(with = "base64_compat")]
        vmsa: Vec<u8>,
        /// SHA-384 digest of all measured pages except VMSA pages.
        #[serde(rename = "value", with = "base64_compat")]
        launch_measurement: Vec<u8>,
    },
}

impl LaunchMeasurement {
    /// Resolves the full 48-byte launch measurement for comparison with an
    /// attestation report.
    ///
    /// For V1, returns the stored digest. For V2, reconstructs the full
    /// measurement by applying the VMSA page update once per vCPU.
    #[must_use]
    pub fn resolve(&self, vcpus: usize) -> [u8; 48] {
        match self {
            Self::V1(bytes) => bytes
                .as_slice()
                .try_into()
                .expect("V1 measurement must be 48 bytes"),
            Self::V2 {
                vmsa,
                launch_measurement,
            } => resolve_v2(launch_measurement, vmsa, vcpus),
        }
    }
}

/// Reconstructs the full launch measurement from a V2 pre-VMSA digest by
/// applying the VMSA page update for each vCPU.
///
/// This reimplements the `PAGE_INFO` digest chain (AMD SEV-SNP Firmware ABI
/// Specification, publication #56860, revision 1.58, May 2025, table 70) from
/// the `snp_measurement` crate because that crate would unnecessarily bloat
/// `clerk-report`.
fn resolve_v2(pre_vmsa_digest: &[u8], vmsa: &[u8], vcpus: usize) -> [u8; 48] {
    const PAGE_INFO_SIZE: u16 = 112;
    const VMSA_GPA: u64 = 0xffff_ffff_f000;

    // SHA-384 of VMSA bytes zero-padded to 4 KiB (same for every vCPU)
    let mut padded = [0u8; 4096];
    padded[..vmsa.len()].copy_from_slice(vmsa);
    let contents: [u8; 48] = Sha384::digest(padded).into();

    let mut digest: [u8; 48] = pre_vmsa_digest
        .try_into()
        .expect("V2 pre-VMSA digest must be 48 bytes");

    for _ in 0..vcpus {
        let mut page_info = [0u8; PAGE_INFO_SIZE as usize];
        page_info[0..48].copy_from_slice(&digest); // DIGEST_CUR
        page_info[48..96].copy_from_slice(&contents); // CONTENTS
        page_info[96..98].copy_from_slice(&PAGE_INFO_SIZE.to_le_bytes()); // LENGTH
        page_info[98] = 2; // PAGE_TYPE (VMSA)
        page_info[104..112].copy_from_slice(&VMSA_GPA.to_le_bytes()); // GPA
        digest = Sha384::digest(page_info).into();
    }

    digest
}

#[derive(Clone, Debug, Deserialize, Serialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(rename_all = "camelCase")]
pub struct PublishedVersionEntry {
    /// The kind of entry, in case we want to include other types in the future (e.g., key revocations, ...)
    pub kind: String,
    /// How [`Self::launch_measurement`] was calculated (breaking changes increase the version)
    pub generation: usize,
    #[serde(serialize_with = "rfc_3339_no_z")]
    pub created_at: chrono::DateTime<chrono::offset::FixedOffset>,
    /// Reference to the source repository as Nix flake URL
    pub ref_: String,
    /// Launch measurement for the system version. Generation 1 stores a flat
    /// SHA-384 digest; generation 2 stores the pre-VMSA digest and VMSA bytes
    /// separately to decouple the measurement from the vCPU count.
    #[cfg_attr(feature = "utoipa", schema(value_type = serde_json::Value))]
    pub launch_measurement: LaunchMeasurement,
    /// YAXI signature
    pub signature: PublishedVersionEntrySignature,
}

impl PublishedVersionEntry {
    /// Verifies whether this entry is signed with a well-known key.
    ///
    /// # Errors
    ///
    /// Returns an error if the key id is unknown. Also forwards errors from
    /// [`ed25519_dalek`](ed25519_dalek::SignatureError)
    pub fn verify_signature(&self) -> Result<(), SignatureVerificationError> {
        let Some(verifying_key) = verifying_key_by_id(&self.signature.key_id) else {
            return Err(SignatureVerificationError::UnknownKeyId);
        };

        self.verify_signature_with_key(&verifying_key)
    }

    fn verify_signature_with_key(
        &self,
        verifying_key: &VerifyingKey,
    ) -> Result<(), SignatureVerificationError> {
        let mut input = Vec::new();
        input.extend_from_slice(self.kind.as_bytes());
        input.extend_from_slice(self.generation.to_string().as_bytes());
        input.extend_from_slice(self.created_at.to_rfc3339().as_bytes());
        input.extend_from_slice(self.ref_.as_bytes());
        match &self.launch_measurement {
            LaunchMeasurement::V1(bytes) => input.extend_from_slice(bytes),
            LaunchMeasurement::V2 {
                vmsa,
                launch_measurement,
            } => {
                input.extend_from_slice(vmsa);
                input.extend_from_slice(launch_measurement);
            }
        }

        let signature = Signature::from_slice(&self.signature.value)
            .map_err(SignatureVerificationError::SignatureError)?;

        verifying_key
            .verify(&input, &signature)
            .map_err(SignatureVerificationError::SignatureError)
    }
}

#[serde_with::serde_as]
#[derive(Clone, Debug, Deserialize, Serialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(rename_all = "camelCase")]
pub struct PublishedVersionEntrySignature {
    /// SHA-256 hash of Ed25519 public key
    #[serde(with = "base64_compat")]
    #[cfg_attr(feature = "utoipa", schema(value_type = String, format = "byte", example = "dYa685dhHap8RSUtB4DDy1l4UcycsGhklBnV5a/4HSg="))]
    pub key_id: Vec<u8>,
    /// Signing input: [`PublishedVersionEntry::kind`] || `[PublishedVersionEntry::generation` || [`PublishedVersionEntry::created_at`] || [`PublishedVersionEntry::ref_`] || [`PublishedVersionEntry::launch_measurement`]
    #[serde(with = "base64_compat")]
    #[cfg_attr(feature = "utoipa", schema(value_type = String, format = "byte", example = "V89vbJZCP3ZNskERuNN1MkzaL5of68dPZ36DigmXmgBUn2fthx6M57Iyh5USOIR5GZyJi119yf3erp7Zi1mZCQ=="))]
    pub value: Vec<u8>,
}

#[cfg(not(any(test, feature = "test-signing-keys")))]
fn verifying_key_by_id(key_id: &[u8]) -> Option<VerifyingKey> {
    if key_id == b"\x02\x1a\xd4^\xc5\x7f\\\x0b\xc8\x13n\x11C\xf5m\xff5\xe8.\x87o\x8dz\x16\x0fg\xe1,k\x913\xb5" {
        // In veehaitch's possession
        Some(VerifyingKey::from_bytes(b"e\xbcL!\xf5\x9b*O\xd6\xb5}1D\x85W\xe8{[\xd1\x15\xef$\xc3\xd7R\x8c\xa0\xc3\xeco\xb4\xe2").unwrap())
    } else if key_id == b"\x0f}3E\x87\xbcR\x0f{\xdff\xf8=\xed\x81\x03\x05\x80^\x7f\xd3\xe4\xdb,\xd8rN\xa7y\x0b\x0b\\" {
        // In trundle's possession
        Some(VerifyingKey::from_bytes(b"\x08\xcbka\x11\x8c\x10\xf0\x1d\xc5h\x1a\xb3sV\x02\xd2Ik\xd3.\"Y\xaeu\xc1~\xd7\x0c\x85j%").unwrap())
    } else {
        None
    }
}

#[cfg(any(test, feature = "test-signing-keys"))]
fn verifying_key_by_id(key_id: &[u8]) -> Option<VerifyingKey> {
    // qzLgDnRegbiQzY416i9/MClrmMp24jcHzaWCWSWSutA=
    if key_id == b"u\x86\xba\xf3\x97a\x1d\xaa|E%-\x07\x80\xc3\xcbYxQ\xcc\x9c\xb0hd\x94\x19\xd5\xe5\xaf\xf8\x1d(" {
        // Signing key: KWd9W5LR/7jNTS26CvuCRcJ/rgXD/tS9S6KqA/pDczY=
        Some(VerifyingKey::from_bytes(b"\xab2\xe0\x0et^\x81\xb8\x90\xcd\x8e5\xea/\x7f0)k\x98\xcav\xe27\x07\xcd\xa5\x82Y%\x92\xba\xd0").unwrap())
    } else {
        None
    }
}

#[cfg(any(test, feature = "signing"))]
#[must_use]
pub fn key_id(
    verifying_key: &VerifyingKey,
) -> GenericArray<u8, <Sha256 as OutputSizeUser>::OutputSize> {
    let mut hasher = Sha256::new();
    hasher.update(verifying_key.to_bytes());
    hasher.finalize()
}

/// Creates a signed system version entry for the given launch measurement.
///
/// The generation is derived from the [`LaunchMeasurement`] variant.
#[cfg(any(test, feature = "signing"))]
#[must_use]
pub fn signed_version_entry(
    launch_measurement: LaunchMeasurement,
    flake_url: &str,
    signing_key: &SigningKey,
) -> PublishedVersionEntry {
    let kind = "release".to_string();
    let generation = match &launch_measurement {
        LaunchMeasurement::V1(_) => 1,
        LaunchMeasurement::V2 { .. } => 2,
    };
    let created_at: chrono::DateTime<chrono::offset::FixedOffset> =
        chrono::offset::Utc::now().into();

    let mut input = Vec::new();
    input.extend_from_slice(kind.as_bytes());
    input.extend_from_slice(generation.to_string().as_bytes());
    input.extend_from_slice(created_at.to_rfc3339().as_bytes());
    input.extend_from_slice(flake_url.as_bytes());
    match &launch_measurement {
        LaunchMeasurement::V1(bytes) => input.extend_from_slice(bytes),
        LaunchMeasurement::V2 {
            vmsa,
            launch_measurement,
        } => {
            input.extend_from_slice(vmsa);
            input.extend_from_slice(launch_measurement);
        }
    }

    let signature = {
        let signature = signing_key.sign(&input);
        PublishedVersionEntrySignature {
            key_id: key_id(&signing_key.verifying_key()).to_vec(),
            value: signature.to_vec(),
        }
    };

    PublishedVersionEntry {
        kind,
        generation,
        created_at,
        ref_: flake_url.into(),
        launch_measurement,
        signature,
    }
}

/// Used as serialize function to make sure the string value of `PublishedVersionEntry`'s
/// `created_at` field matches the value used to build the signature.
fn rfc_3339_no_z<S: serde::Serializer>(
    source: &chrono::DateTime<chrono::offset::FixedOffset>,
    serializer: S,
) -> Result<S::Ok, S::Error> {
    source.to_rfc3339().serialize(serializer)
}

mod base64_compat {
    use base64::prelude::{Engine, BASE64_STANDARD, BASE64_URL_SAFE_NO_PAD};
    use serde::{de, Deserialize, Serialize, Serializer};

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let value: String = Deserialize::deserialize(deserializer)?;
        let url_safe = value.replace('+', "-").replace('/', "_");
        let unpadded = url_safe.trim_end_matches('=');
        BASE64_URL_SAFE_NO_PAD
            .decode(unpadded)
            .map_err(de::Error::custom)
    }

    pub fn serialize<S: Serializer, T: AsRef<[u8]>>(
        source: &T,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        // Serialization still happens with base64 standard because existing consumers don't expect
        // base64url yet
        BASE64_STANDARD.encode(source).serialize(serializer)
    }
}

#[cfg(test)]
mod tests {
    use base64::Engine;
    use ed25519_dalek::Signature;

    use crate::versions::verifying_key_by_id;

    use super::{signed_version_entry, LaunchMeasurement, PublishedVersionEntry};

    #[test]
    fn deserialize_published_version_entry() {
        let entry: PublishedVersionEntry =
            serde_json::from_str(published_version_entry_json()).unwrap();
        assert_published_version_entry_fields(&entry);
    }

    #[test]
    fn serialization_round_trip() {
        let entry: PublishedVersionEntry =
            serde_json::from_str(published_version_entry_json()).unwrap();
        let serialized = serde_json::to_string_pretty(&entry).unwrap();
        let entry: PublishedVersionEntry = serde_json::from_str(&serialized).unwrap();
        assert_published_version_entry_fields(&entry);
    }

    #[test]
    fn test_signature() {
        let entry = published_version_entry();
        assert!(entry.verify_signature().is_ok());
    }

    #[test]
    fn unknown_key_id_fails() {
        let mut entry = published_version_entry();
        entry.signature.key_id[0] = entry.signature.key_id[0].wrapping_add(1);
        assert!(entry.verify_signature().is_err());
    }

    #[test]
    fn wrong_signature_fails() {
        let mut entry = published_version_entry();
        entry.signature.value[5] = entry.signature.value[5].wrapping_add(1);
        assert!(entry.verify_signature().is_err());
    }

    #[test]
    fn deserializes_standard_base64_for_compatibility() {
        let entry: PublishedVersionEntry =
            serde_json::from_str(published_version_entry_standard_base64_json()).unwrap();
        assert_published_version_entry_fields(&entry);
    }

    #[test]
    fn serializes_to_standard_base64_for_compatibility() {
        use base64::prelude::BASE64_STANDARD;

        let serialized = serde_json::to_string(&published_version_entry()).unwrap();
        let deserialized: serde_json::Value = serde_json::from_str(&serialized).unwrap();

        let launch_measurement = deserialized
            .as_object()
            .and_then(|obj| {
                obj.get("launchMeasurement")
                    .and_then(|value| value.as_str())
            })
            .unwrap();
        match BASE64_STANDARD.decode(launch_measurement) {
            Ok(measurement) => assert_launch_measurement(&measurement),
            Err(err) => panic!("Could not base64 deserialize launch measurement: {err:?}"),
        }
    }

    #[test]
    fn test_signed_version_entry() {
        use ed25519_dalek::SigningKey;

        let signing_key = SigningKey::generate(&mut rand_core::OsRng);

        let entry = signed_version_entry(
            published_version_entry().launch_measurement,
            "github:yaxitech/nixos/a24091a8d5f8dc74f64435c029816c7c2c0c7a70",
            &signing_key,
        );

        assert_eq!(entry.kind, "release");
        assert_eq!(entry.generation, 1);
        assert_eq!(
            entry.ref_,
            "github:yaxitech/nixos/a24091a8d5f8dc74f64435c029816c7c2c0c7a70".to_string()
        );
        assert_v1_launch_measurement(&entry);
        assert!(entry
            .verify_signature_with_key(&signing_key.verifying_key())
            .is_ok());
    }

    #[test]
    fn test_deserializes_entry_with_z_offset() {
        let entry: Result<PublishedVersionEntry, _> =
            serde_json::from_str(published_version_entry_json_z_offset());
        assert!(entry.is_ok());
    }

    #[test]
    fn test_signature_uses_json_time_value() {
        // Like `PublishedVersionEntry`, but with unparsed string values
        #[derive(serde::Deserialize)]
        #[serde(rename_all = "camelCase")]
        #[allow(unused)]
        struct RawVersionEntry {
            pub kind: String,
            pub generation: usize,
            pub created_at: String,
            pub ref_: String,
            #[serde(with = "super::base64_compat")]
            pub launch_measurement: Vec<u8>,
            pub signature: super::PublishedVersionEntrySignature,
        }
        let version_entry: RawVersionEntry =
            serde_json::from_str(published_version_entry_json()).unwrap();

        let mut input = Vec::new();
        input.extend_from_slice(version_entry.kind.as_bytes());
        input.extend_from_slice(version_entry.generation.to_string().as_bytes());
        input.extend_from_slice(version_entry.created_at.as_bytes());
        input.extend_from_slice(version_entry.ref_.as_bytes());
        input.extend_from_slice(&version_entry.launch_measurement);

        let signature = Signature::from_slice(&version_entry.signature.value).unwrap();
        let key = verifying_key_by_id(&version_entry.signature.key_id).unwrap();
        key.verify_strict(&input, &signature).unwrap();
    }

    #[test]
    fn v2_signed_entry_is_valid() {
        use ed25519_dalek::SigningKey;

        let signing_key = SigningKey::generate(&mut rand_core::OsRng);

        let vmsa = vec![0xabu8; 1992];
        let measurement = vec![0x42u8; 48];

        let entry = signed_version_entry(
            LaunchMeasurement::V2 {
                vmsa: vmsa.clone(),
                launch_measurement: measurement.clone(),
            },
            "github:yaxitech/cvms/v0.1.0",
            &signing_key,
        );

        assert_eq!(entry.kind, "release");
        assert_eq!(entry.generation, 2);
        assert_eq!(entry.ref_, "github:yaxitech/cvms/v0.1.0".to_string());

        match &entry.launch_measurement {
            LaunchMeasurement::V1(_) => panic!("Expected V2 launch measurement"),
            LaunchMeasurement::V2 {
                vmsa: v,
                launch_measurement: m,
            } => {
                assert_eq!(v, &vmsa);
                assert_eq!(m, &measurement);
            }
        }

        assert!(entry
            .verify_signature_with_key(&signing_key.verifying_key())
            .is_ok());
    }

    #[test]
    fn v2_serialization_round_trip() {
        use ed25519_dalek::SigningKey;

        let signing_key = SigningKey::generate(&mut rand_core::OsRng);
        let entry = signed_version_entry(
            LaunchMeasurement::V2 {
                vmsa: vec![0x01u8; 1992],
                launch_measurement: vec![0x02u8; 48],
            },
            "github:yaxitech/cvms/v1.0.0",
            &signing_key,
        );

        let serialized = serde_json::to_string_pretty(&entry).unwrap();
        let deserialized: serde_json::Value = serde_json::from_str(&serialized).unwrap();

        // Verify launchMeasurement is an object (not a string) in gen 2
        assert!(
            deserialized["launchMeasurement"].is_object(),
            "gen 2 launchMeasurement should serialize as object"
        );
        assert!(deserialized["launchMeasurement"]["vmsa"].is_string());
        assert!(deserialized["launchMeasurement"]["value"].is_string());

        // Round-trip
        let restored: PublishedVersionEntry = serde_json::from_str(&serialized).unwrap();
        assert!(restored
            .verify_signature_with_key(&signing_key.verifying_key())
            .is_ok());
    }

    #[test]
    fn v2_wrong_signature_fails() {
        use ed25519_dalek::SigningKey;

        let signing_key = SigningKey::generate(&mut rand_core::OsRng);
        let mut entry = signed_version_entry(
            LaunchMeasurement::V2 {
                vmsa: vec![0xabu8; 1992],
                launch_measurement: vec![0x42u8; 48],
            },
            "github:yaxitech/cvms/v0.1.0",
            &signing_key,
        );

        entry.signature.value[5] = entry.signature.value[5].wrapping_add(1);
        assert!(entry
            .verify_signature_with_key(&signing_key.verifying_key())
            .is_err());
    }

    #[test]
    fn launch_measurement_deserializes_both_variants() {
        // V1: plain base64 string
        let v1: LaunchMeasurement = serde_json::from_str(r#""WUFYSQ==""#).unwrap();
        assert!(matches!(v1, LaunchMeasurement::V1(ref b) if b == b"YAXI"));

        // V2: object with vmsa and value (url-safe base64)
        let v2: LaunchMeasurement =
            serde_json::from_str(r#"{"vmsa": "_wD_", "value": "AP__"}"#).unwrap();
        assert!(
            matches!(v2, LaunchMeasurement::V2 { ref vmsa, ref launch_measurement }
                if vmsa == &[0xff, 0x00, 0xff] && launch_measurement == &[0x00, 0xff, 0xff])
        );
    }

    fn published_version_entry() -> PublishedVersionEntry {
        serde_json::from_str(published_version_entry_json()).unwrap()
    }

    fn published_version_entry_standard_base64_json() -> &'static str {
        r#"
            {
              "kind": "release",
              "generation": 1,
              "createdAt": "2025-03-03T21:53:07.650800415+00:00",
              "ref": "github:yaxitech/nixos/a24091a8d5f8dc74f64435c029816c7c2c0c7a70",
              "launchMeasurement": "QZPtHQhKa5bVGeNnPOBwzH1ycRDNYFy1/yyQl2Tq+T+VOHelRODRYI5kPd6u5dgV",
              "signature": {
                "keyId": "dYa685dhHap8RSUtB4DDy1l4UcycsGhklBnV5a/4HSg=",
                "value": "NJpXrwzxi570HwBMGrlDndesoFiOPyaMXzAYjwzr+LGHnlr+gGrnX6pIOAXKdH9XVOwziAX7hYs44Ii9Gxc2CA=="
              }
            }
        "#
    }

    fn published_version_entry_json() -> &'static str {
        r#"
            {
              "kind": "release",
              "generation": 1,
              "createdAt": "2025-03-03T21:53:07.650800415+00:00",
              "ref": "github:yaxitech/nixos/a24091a8d5f8dc74f64435c029816c7c2c0c7a70",
              "launchMeasurement": "QZPtHQhKa5bVGeNnPOBwzH1ycRDNYFy1_yyQl2Tq-T-VOHelRODRYI5kPd6u5dgV",
              "signature": {
                "keyId": "dYa685dhHap8RSUtB4DDy1l4UcycsGhklBnV5a_4HSg=",
                "value": "NJpXrwzxi570HwBMGrlDndesoFiOPyaMXzAYjwzr-LGHnlr-gGrnX6pIOAXKdH9XVOwziAX7hYs44Ii9Gxc2CA=="
              }
            }
        "#
    }

    fn published_version_entry_json_z_offset() -> &'static str {
        r#"
            {
              "kind": "routex-release",
              "generation": 1,
              "createdAt": "2024-03-26T19:16:40.269247621Z",
              "ref": "github:yaxitech/nixos/a24091a8d5f8dc74f64435c029816c7c2c0c7a70",
              "launchMeasurement": "QZPtHQhKa5bVGeNnPOBwzH1ycRDNYFy1_yyQl2Tq-T-VOHelRODRYI5kPd6u5dgV",
              "signature": {
                "keyId": "dYa685dhHap8RSUtB4DDy1l4UcycsGhklBnV5a_4HSg=",
                "value": "w9kKh0xKFNKm5Mcn1-VSiMgCWs_jAZIAkJfGjiTkFQnj4CYD5qskb2LnQeE4CJa9CrgmT21l5pkIJR7gUP6UDA=="
              }
            }
        "#
    }

    fn assert_published_version_entry_fields(entry: &PublishedVersionEntry) {
        assert_eq!(entry.kind, "release");
        assert_eq!(entry.generation, 1);
        assert_eq!(
            entry.ref_,
            "github:yaxitech/nixos/a24091a8d5f8dc74f64435c029816c7c2c0c7a70".to_string()
        );
        assert_v1_launch_measurement(entry);
        assert_eq!(
            entry.signature.key_id,
            &[
                117, 134, 186, 243, 151, 97, 29, 170, 124, 69, 37, 45, 7, 128, 195, 203, 89, 120,
                81, 204, 156, 176, 104, 100, 148, 25, 213, 229, 175, 248, 29, 40
            ],
        );
    }

    fn assert_v1_launch_measurement(entry: &PublishedVersionEntry) {
        match &entry.launch_measurement {
            LaunchMeasurement::V1(bytes) => assert_launch_measurement(bytes),
            LaunchMeasurement::V2 { .. } => panic!("Expected V1 launch measurement"),
        }
    }

    fn assert_launch_measurement(actual: &[u8]) {
        assert_eq!(
            actual,
            &[
                65, 147, 237, 29, 8, 74, 107, 150, 213, 25, 227, 103, 60, 224, 112, 204, 125, 114,
                113, 16, 205, 96, 92, 181, 255, 44, 144, 151, 100, 234, 249, 63, 149, 56, 119, 165,
                68, 224, 209, 96, 142, 100, 61, 222, 174, 229, 216, 21
            ]
        );
    }

    #[test]
    fn v1_resolve_returns_stored_digest() {
        let bytes = [0x42u8; 48];
        let measurement = LaunchMeasurement::V1(bytes.to_vec());
        assert_eq!(measurement.resolve(4), bytes);
    }

    #[test]
    fn v2_resolve_with_zero_vcpus_returns_pre_vmsa_digest() {
        let digest = [0x42u8; 48];
        let measurement = LaunchMeasurement::V2 {
            vmsa: vec![0xab; 1992],
            launch_measurement: digest.to_vec(),
        };
        assert_eq!(measurement.resolve(0), digest);
    }

    #[test]
    fn v2_resolve_changes_with_vcpu_count() {
        let measurement = LaunchMeasurement::V2 {
            vmsa: vec![0xab; 1992],
            launch_measurement: vec![0x42; 48],
        };
        assert_ne!(measurement.resolve(1), measurement.resolve(2));
    }
}
