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

//! # Attestation Module
//!
//! Provides attestation functionality for REX services running in Trusted Execution Environments.
//!
//! This module handles the generation and verification of attestation reports that prove:
//! - The REX service is running in a genuine TEE
//! - The REX service's code hasn't been tampered with
//! - The REX service's signing keys are protected by the TEE
//!
//! ## Overview
//!
//! Attestation reports contain:
//! - The REX service's public key (base64 encoded)
//! - A cryptographic signature over this field
//! - Optional platform-specific attestation evidence (when running in a TEE)
//!
//! ## Usage
//!
//! To generate attestation reports, use the `AttestationReportExt` trait from `rialo_tee_attestation`.
//! To verify attestation reports, use the `AttestationReportVerifiable` trait from `rialo_tee_verification`:
//!
//! ```ignore
//! use rialo_types::{AttestationReport, PublicKey};
//! use rialo_tee_attestation::AttestationReportExt;
//! use rialo_tee_verification::AttestationReportVerifiable;
//! use ring::signature::Ed25519KeyPair;
//! use rcgen::PKCS_ED25519;
//!
//! // Replace with your actual key generation logic
//! let keypair_der = rcgen::KeyPair::generate_for(&PKCS_ED25519)
//!         .unwrap()
//!         .serialize_der();
//! let keypair = Ed25519KeyPair::from_pkcs8(&keypair_der).unwrap();
//!
//! // Create a wrapping public key for secret sharing
//! let wrapping_key = PublicKey::from_bytes([42u8; 32]);
//!
//! // Generate the attestation report
//! let report = AttestationReport::generate(&keypair, wrapping_key).unwrap();
//! let verified = report.verify(None).unwrap();
//! let verification_key = verified.verification_key();
//! ```
//!
//! ## Platform Support
//!
//! Currently supported TEE platforms:
//! - AMD SEV-SNP Bare Metal (planned)
//! - AWS SEV-SNP
//! - Azure CVM (planned)
//! - Intel TDX (planned)

#[cfg(feature = "non-pdk")]
use anyhow::Result;
use borsh::{BorshDeserialize, BorshSerialize};
#[cfg(feature = "non-pdk")]
use fastcrypto::{hash, hash::HashFunction};
#[cfg(feature = "non-pdk")]
use ring::signature::{Ed25519KeyPair, KeyPair, Signature};
use serde::{Deserialize, Serialize};
use serde_big_array::BigArray;

use crate::{tee_types::AttestationEvidence, PublicKey};

/// Attestation report that will be sent to the caller
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct AttestationReport {
    pub inner: AttestationReportInner,
    #[serde(with = "BigArray")]
    pub signature: [u8; 64],
    pub platform_attestation: Option<AttestationEvidence>,
}

#[cfg(feature = "non-pdk")]
impl AttestationReport {
    /// Computes the hash of the attestation report fields
    pub fn compute_report_hash(inner: &AttestationReportInner) -> [u8; 32] {
        let mut hash = hash::Sha256::new();
        hash.update(inner.verification_key);
        hash.update::<&[u8]>(inner.wrapping_key.as_ref());
        hash.finalize().digest
    }

    /// Verifies the report signature
    pub fn verify_signature(
        &self,
        digest: &[u8],
    ) -> Result<VerifiedAttestation, VerificationError> {
        ring::signature::UnparsedPublicKey::new(
            &ring::signature::ED25519,
            &self.inner.verification_key,
        )
        .verify(digest, &self.signature)
        .map_err(|_| VerificationError::SignatureVerification)?;

        Ok(VerifiedAttestation::new_verified_dangerous(
            self.inner.clone(),
        ))
    }
}

impl Default for AttestationReport {
    fn default() -> Self {
        Self {
            inner: AttestationReportInner::default(),
            signature: [0; 64],
            platform_attestation: None,
        }
    }
}

/// Inner structure of the attestation report containing all the fields that need to be verified
#[derive(
    Debug, Clone, Eq, PartialEq, Default, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
)]
pub struct AttestationReportInner {
    pub verification_key: [u8; 32],
    pub wrapping_key: PublicKey,
}

/// Result of attestation verification containing necessary public data
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct VerifiedAttestation {
    inner: AttestationReportInner,
}

impl VerifiedAttestation {
    /// Creates a new VerifiedAttestation from the inner report
    ///
    /// # Safety / Warning
    /// This constructor is marked "dangerous" because it bypasses all verification checks.
    /// It must ONLY be called after proper cryptographic verification has been performed
    /// by the verification trait in rialo-tee-verification.
    ///
    /// Misuse of this method can lead to severe security vulnerabilities as it allows
    /// creating a "verified" attestation without actually verifying anything.
    pub fn new_verified_dangerous(inner: AttestationReportInner) -> Self {
        Self { inner }
    }

    /// Get the public verification key bytes from the verified attestation
    pub fn verification_key(&self) -> [u8; 32] {
        self.inner.verification_key
    }

    /// Returns the public wrapping key from the verified attestation.
    pub fn wrapping_key(&self) -> PublicKey {
        self.inner.wrapping_key
    }
}

/// Errors that can happen while verifying an [`AttestationReport`].
#[derive(Debug, thiserror::Error)]
pub enum VerificationError {
    #[error("Failed to verify attestation report's signature")]
    SignatureVerification,

    #[error("Invalid platform attestation evidence")]
    InvalidPlatformAttestation,

    #[error("Platform attestation verification failed: {0}")]
    PlatformAttestationVerificationFailed(String),
}

/// Converts an Ed25519 signature to a fixed-size byte array.
///
/// # Arguments
///
/// * `signature` - The Ed25519 signature to convert.
///
/// # Returns
///
/// * `[u8; 64]` - The signature as a fixed-size byte slice.
#[cfg(feature = "non-pdk")]
pub fn ed25519_signature_to_slice(signature: &Signature) -> [u8; 64] {
    let bytes = signature.as_ref();
    if bytes.len() != 64 {
        unreachable!(
            "Ed25519 signatures are always 64 bytes, but got {}",
            bytes.len()
        );
    }
    bytes.try_into().unwrap()
}

/// Converts am Ed25519 keypair pubkey to a fixed-size byte array.
///
/// # Arguments
///
/// * `pubkey` - The Ed25519 keypair with the public key to convert.
///
/// # Returns
///
/// * `[u8; 32]` - The public key as a fixed-size byte array.
#[cfg(feature = "non-pdk")]
pub fn ed25519_keypair_to_pubkey_slice(keypair: &Ed25519KeyPair) -> [u8; 32] {
    let pubkey = keypair.public_key().as_ref();
    if pubkey.len() != 32 {
        unreachable!(
            "Ed25519 public keys are always 32 bytes, but got {}",
            pubkey.len()
        );
    }
    pubkey.try_into().unwrap()
}

#[cfg(all(test, feature = "non-pdk"))]
mod tests {
    use rcgen::PKCS_ED25519;
    use ring::signature::Ed25519KeyPair;

    use super::*;
    use crate::tee_types::TeeType;

    // Test constant for wrapping public key
    const TEST_WRAPPING_KEY_BYTES: [u8; 32] = [42u8; 32];

    fn test_wrapping_key() -> PublicKey {
        PublicKey::from_bytes(TEST_WRAPPING_KEY_BYTES)
    }

    // Helper function to generate a test keypair
    fn generate_test_keypair() -> Ed25519KeyPair {
        let key_pair_der = rcgen::KeyPair::generate_for(&PKCS_ED25519)
            .unwrap()
            .serialize_der();
        Ed25519KeyPair::from_pkcs8(&key_pair_der).unwrap()
    }

    // NOTE: Tests that use AttestationReport::generate have been moved to rialo-tee-attestation
    // to avoid circular dependencies during test builds. The generate() method is now provided
    // by the AttestationReportExt trait in rialo-tee-attestation.

    #[test]
    /// This test verifies that the constructor works correctly.
    /// It tests the manual creation of attestation reports with platform evidence.
    fn test_attestation_report_with_platform_attestation_constructor() {
        let key_pair = generate_test_keypair();
        let public_key_bytes = ed25519_keypair_to_pubkey_slice(&key_pair);

        // Create test platform attestation evidence
        let platform_evidence = AttestationEvidence {
            tee_type: TeeType::SevSnp,
            evidence_data: vec![1, 2, 3, 4, 5], // Mock evidence data
        };

        // Create the inner attestation report
        let inner = AttestationReportInner {
            verification_key: public_key_bytes,
            wrapping_key: test_wrapping_key(),
        };

        // Compute hash and sign
        let digest = AttestationReport::compute_report_hash(&inner);
        let signature = ed25519_signature_to_slice(&key_pair.sign(&digest));

        // Create attestation report with platform attestation
        let attestation_report = AttestationReport {
            inner,
            signature,
            platform_attestation: Some(platform_evidence.clone()),
        };

        // Verify that the platform attestation is included
        assert!(attestation_report.platform_attestation.is_some());
        let platform_attestation = attestation_report.platform_attestation.unwrap();
        assert_eq!(platform_attestation.tee_type, platform_evidence.tee_type);
        assert_eq!(
            platform_attestation.evidence_data,
            platform_evidence.evidence_data
        );
    }
}
