// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Core TEE types used in REX attestation

use borsh::{BorshDeserialize, BorshSerialize};
pub use rialo_tee_types::TeeType;
use serde::{Deserialize, Serialize};

// Re-export crypto types for backward compatibility
pub use crate::crypto::{PublicKey, TeeUserData, TEE_USER_DATA_SIZE, X25519_KEY_SIZE};

/// TEE platform attestation evidence
#[derive(
    Clone, Default, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
)]
pub struct AttestationEvidence {
    pub tee_type: TeeType,
    pub evidence_data: Vec<u8>,
}

impl std::fmt::Debug for AttestationEvidence {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let truncated_data = if self.evidence_data.len() <= 16 {
            format!("{:?}", self.evidence_data)
        } else {
            format!(
                "{:?}... (truncated, total: {} bytes)",
                &self.evidence_data[..16],
                self.evidence_data.len()
            )
        };

        f.debug_struct("AttestationEvidence")
            .field("tee_type", &self.tee_type)
            .field("evidence_data", &truncated_data)
            .finish()
    }
}

/// SEV-SNP specific evidence data structure
///
/// This structure contains the raw SEV-SNP attestation data that gets
/// serialized/deserialized from the `evidence_data` field of `AttestationEvidence`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SevSnpEvidenceData {
    /// The raw SEV-SNP attestation report
    pub report: Vec<u8>,
    /// the VLEK or VCEK certificate
    pub vek_certificate: Vec<u8>,
}

/// Azure SEV-SNP specific evidence data structure
///
/// This structure contains the raw SEV-SNP attestation data that gets
/// serialized/deserialized from the `evidence_data` field of `AttestationEvidence`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AzureSevSnpEvidenceData {
    /// SEV-SNP attestation report containing user_data
    pub snp_evidence: SevSnpEvidenceData,
}

/// AWS Nitro TPM COSE Evidence structure used for attestation verification
///
/// This structure contains the AWS Nitro TPM attestation data in COSE format along with
/// supporting metadata required for verification. The COSE report includes TPM-based
/// measurements and signatures that prove the identity and state of the AWS Nitro enclave.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AwsNitroCoseEvidence {
    /// AWS COSE signed report containing TPM attestation data
    ///
    /// This report includes Platform Configuration Registers (PCRs)
    /// and digital signatures required to verify the authenticity and integrity of
    /// the enclave.
    pub report: Vec<u8>,
    /// Cryptographic nonce value used to prevent replay attacks
    ///
    /// This nonce is included in the attestation request and must match
    /// the value embedded in the COSE report to ensure freshness.
    pub nonce: Vec<u8>,
    /// Unix timestamp (in seconds) when the COSE report was generated
    ///
    /// Used to verify the attestation report's recency and detect
    /// potential replay attacks.
    pub time: u64,
}

/// AWS SEV-SNP specific evidence data structure optimized for AWS Nitro
///
/// This simplified structure leverages AWS Nitro's built-in attestation capabilities
/// and prepares for future AWS COSE signed reports that will include all TPM data.
/// The user_data contains the REX digest (32 bytes) that gets embedded in both
/// the SEV-SNP report and the AWS COSE report for binding.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AwsSevSnpEvidenceData {
    /// SEV-SNP attestation report containing user_data
    pub snp_evidence: SevSnpEvidenceData,
    /// AWS Nitro COSE attestation report
    pub aws_nitro_cose_evidence: AwsNitroCoseEvidence,
}

/// Error type for TEE evidence parsing
#[derive(Debug, thiserror::Error)]
pub enum TeeEvidenceError {
    #[error("Invalid evidence: {0}")]
    InvalidEvidence(String),
}

impl TryFrom<&AttestationEvidence> for SevSnpEvidenceData {
    type Error = TeeEvidenceError;
    fn try_from(evidence: &AttestationEvidence) -> Result<Self, Self::Error> {
        parse_evidence(evidence, &[TeeType::SevSnp])
    }
}

impl TryFrom<&AttestationEvidence> for AwsSevSnpEvidenceData {
    type Error = TeeEvidenceError;
    fn try_from(evidence: &AttestationEvidence) -> Result<Self, Self::Error> {
        parse_evidence(evidence, &[TeeType::AwsSevSnp])
    }
}

impl TryFrom<&AttestationEvidence> for AzureSevSnpEvidenceData {
    type Error = TeeEvidenceError;
    fn try_from(evidence: &AttestationEvidence) -> Result<Self, Self::Error> {
        parse_evidence(evidence, &[TeeType::AzureSevSnp])
    }
}

/// Generic evidence parsing function that validates TEE type and deserializes evidence data
///
/// This function provides a unified way to parse evidence data across different TEE types,
/// validating that the evidence matches one of the expected TEE types before attempting
/// to deserialize the data into the target type.
pub fn parse_evidence<T>(
    evidence: &AttestationEvidence,
    expected_tee_types: &[TeeType],
) -> Result<T, TeeEvidenceError>
where
    T: for<'de> Deserialize<'de>,
{
    // Validate TEE type
    if !expected_tee_types.contains(&evidence.tee_type) {
        return Err(TeeEvidenceError::InvalidEvidence(format!(
            "Expected TEE types {:?}, got {:?}",
            expected_tee_types, evidence.tee_type
        )));
    }

    // Parse evidence data
    serde_json::from_slice(&evidence.evidence_data)
        .map_err(|e| TeeEvidenceError::InvalidEvidence(format!("Failed to parse evidence: {e}")))
}
