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

//! Secret sharing key management for TEEs

// The zeroize derive macro with #[zeroize(skip)] generates code that triggers this lint
#![allow(unused_assignments)]

use core::fmt;

use rialo_types::{PublicKey, X25519_KEY_SIZE};
use x25519_dalek::StaticSecret as PrivateKey;
use zeroize::{Zeroize, Zeroizing};

/// Size of BLAKE3 hash output in bytes.
pub const BLAKE3_HASH_SIZE: usize = 32;

/// Domain separator for Secret Sharing Key attestation payloads.
pub const ATTEST_DOMAIN: &[u8] = b"RIALO_SSK_ATTEST_V1";

/// Key type tag for X25519 Secret Sharing Keys.
pub const ATTEST_KEY_TYPE_X25519: &[u8] = b"x25519";

/// Compute the attestation payload digest for a Secret Sharing Key public key.
///
/// Preimage = ATTEST_DOMAIN || ATTEST_KEY_TYPE_X25519 || sk_pub_bytes
/// Digest   = BLAKE3(preimage) (32 bytes)
pub fn compute_ssk_attest_payload(sk_pub_bytes: &[u8; X25519_KEY_SIZE]) -> [u8; BLAKE3_HASH_SIZE] {
    let mut hasher = blake3::Hasher::new();
    hasher.update(ATTEST_DOMAIN);
    hasher.update(ATTEST_KEY_TYPE_X25519);
    hasher.update(sk_pub_bytes);

    let hash = hasher.finalize();
    *hash.as_bytes()
}

/// SecretSharingKey: X25519 keypair used by the TEE.
///
/// Intended use:
/// - The TEE generates an SK and publishes `SK.public_bytes()` inside its attestation.
/// - Off-TEE senders encrypt to this SK using ephemeral X25519 → KDF → AEAD
/// - The TEE decrypts using `SK.private_key()` plus the sender's ephemeral public key
///   carried in the message.
///
/// We will have to add a mechanism to wipe out the key if the TEE aborts or crashes.
/// <https://linear.app/subzero-labs/issue/SUB-1094/tee-crash-safe-handling-of-secretsharingkey-zeroize-doesnt-run-on>
#[must_use] // warn if a key is created and ignored
#[derive(Clone, Zeroize)]
#[zeroize(drop)] // wipe fields implementing Zeroize when dropped
pub struct SecretSharingKey {
    #[zeroize(skip)]
    sk_pub: PublicKey,
    sk_priv: PrivateKey,
}

impl SecretSharingKey {
    /// Generate a fresh keypair using the HW CSPRNG.
    /// Note: x25519-dalek v2 expects the *older* rand_core 0.6 `RngCore` for `new(...)`.
    /// Our workspace uses rand_core 0.9.
    /// Filling bytes ourselves avoids pulling two rand_core versions; `from(...)` clamps.
    pub fn generate() -> Self {
        // HW-only policy enforced here as well to keep consistency across crates.
        let raw = Zeroizing::new(
            crate::hw_rng::bytes::<X25519_KEY_SIZE>().expect("HW RNG required in TEE"),
        );
        Self::from_private_key_bytes(raw)
    }

    /// Construct from a 32-byte private key; clamping handled by dalek.
    /// `x25519-dalek` applies "RFC 7748 clamping" on construction (clear/flip bits);
    /// this is NOT an Ed25519 key. Use only X25519 material here.
    pub fn from_private_key_bytes(sk_priv_bytes: Zeroizing<[u8; 32]>) -> Self {
        let sk_priv = PrivateKey::from(*sk_priv_bytes);
        let sk_pub = PublicKey::from(&sk_priv);
        Self { sk_priv, sk_pub }
    }

    /// Borrow the public key.
    pub fn public_key(&self) -> &PublicKey {
        &self.sk_pub
    }

    /// Internal-only: borrow the private key reference.
    /// This is used by internal code paths that need to pass the private key to another TEE.
    /// Keep this crate-scoped to reduce accidental exposure.
    pub(crate) fn private_key(&self) -> &PrivateKey {
        &self.sk_priv
    }

    /// Get raw public key bytes.
    pub fn public_key_bytes(&self) -> PublicKey {
        self.sk_pub
    }

    /// Get private key bytes for TPM sealing.
    /// This is intentionally separate from test methods to make the use case explicit.
    /// Used by TPM sealing to get the raw key material for encryption.
    pub fn to_private_key_bytes(&self) -> [u8; 32] {
        self.sk_priv.to_bytes()
    }
}

impl Default for SecretSharingKey {
    fn default() -> Self {
        Self::generate()
    }
}

/// Redact secrets in debug output; include the public key for troubleshooting.
impl fmt::Debug for SecretSharingKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SecretSharingKey")
            .field("public_key", &self.sk_pub.to_bytes())
            .field("private_key", &"<redacted>")
            .finish()
    }
}

// ---------------- TESTS ONLY ----------------
/// (tests only) Use raw private bytes instead of zeroizing.
/// This is only for test assertions. Do not use in production.
#[cfg(any(test, feature = "testing"))]
impl SecretSharingKey {
    /// Get raw private bytes as a plain copy.
    pub fn private_key_bytes_for_tests(&self) -> [u8; 32] {
        self.sk_priv.to_bytes()
    }

    /// Construct from a 32-byte private key; clamping handled by dalek.
    pub fn from_test_private_key_bytes(sk_priv_bytes: [u8; 32]) -> Self {
        let sk_priv = PrivateKey::from(sk_priv_bytes);
        let sk_pub = PublicKey::from(&sk_priv);
        Self { sk_priv, sk_pub }
    }

    /// Deterministic key from arbitrary seed using BLAKE3 (32 bytes).
    pub fn from_test_seed(seed: impl AsRef<[u8]>) -> Self {
        let h = blake3::hash(seed.as_ref());
        let mut sk = [0u8; 32];
        sk.copy_from_slice(h.as_bytes());
        Self::from_test_private_key_bytes(sk)
    }
}

/// Infallible conversion from raw 32-byte private key.
#[cfg(any(test, feature = "testing"))]
impl From<[u8; 32]> for SecretSharingKey {
    fn from(sk: [u8; 32]) -> Self {
        Self::from_test_private_key_bytes(sk)
    }
}

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

    // Compile-time trait-bound assertion (tests only).
    // We define a dummy trait that requires `Send + Sync` and implement it for
    // `SecretSharingKey`. If those bounds ever stop holding (e.g., a field changes),
    // this impl will fail to compile in `cargo test`. Zero runtime code.
    trait _AssertSendSync: Send + Sync {}
    impl _AssertSendSync for SecretSharingKey {}

    trait _AssertZeroize: zeroize::Zeroize {}
    impl _AssertZeroize for x25519_dalek::StaticSecret {}
    impl _AssertZeroize for SecretSharingKey {}

    #[test]
    fn roundtrip_key_generation() {
        let k = SecretSharingKey::generate();
        let k2 = SecretSharingKey::from_test_private_key_bytes(k.private_key_bytes_for_tests());
        assert_eq!(k.public_key_bytes(), k2.public_key_bytes());
        assert_eq!(
            k.private_key_bytes_for_tests(),
            k2.private_key_bytes_for_tests()
        );
    }

    #[test]
    fn deterministic_from_seed_is_stable() {
        let a = SecretSharingKey::from_test_seed(b"test-seed");
        let b = SecretSharingKey::from_test_seed(b"test-seed");
        assert_eq!(
            a.private_key_bytes_for_tests(),
            b.private_key_bytes_for_tests()
        );
        assert_eq!(a.public_key_bytes(), b.public_key_bytes());
    }

    #[test]
    fn different_seeds_different_keys() {
        let a = SecretSharingKey::from_test_seed(b"test-seed-1");
        let b = SecretSharingKey::from_test_seed(b"test-seed-2");
        assert_ne!(
            a.private_key_bytes_for_tests(),
            b.private_key_bytes_for_tests()
        );
        assert_ne!(a.public_key_bytes(), b.public_key_bytes());
    }

    #[test]
    fn default_is_random() {
        // Extremely high probability these differ.
        let k1 = SecretSharingKey::default();
        let k2 = SecretSharingKey::default();
        assert_ne!(
            k1.private_key_bytes_for_tests(),
            k2.private_key_bytes_for_tests()
        );
        assert_ne!(k1.public_key_bytes(), k2.public_key_bytes());
    }

    #[test]
    fn from_array_via_into() {
        let raw = [7u8; 32];
        let k: SecretSharingKey = raw.into();
        assert_eq!(k.private_key_bytes_for_tests(), raw);
    }
}
