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

//! Type definitions for secret sharing operations
//!
//! This module contains common types, structs, and enums used throughout
//! the secret sharing implementation.

use borsh::{BorshDeserialize, BorshSerialize};
use rialo_types::PublicKey;
use serde_big_array::BigArray;

use crate::SecretSharingError;

/// Logical identifier for a Secret Key (SK) version.
/// This is used in AAD and the envelope metadata to identify which
/// global SK is being used.
/// Note: We can eventually support rotation by making this a more complex struct
/// or replacing with a more sophisticated versioning scheme, but for now a simple
/// integer is sufficient.
#[derive(
    Clone,
    Copy,
    Debug,
    PartialEq,
    Eq,
    Hash,
    serde::Serialize,
    serde::Deserialize,
    BorshSerialize,
    BorshDeserialize,
)]
pub struct KeyId(u64);

impl KeyId {
    /// Creates a new KeyId with validation.
    ///
    /// # Validation Rules
    /// - Must not be zero (reserved for uninitialized state)
    /// - Must not be u64::MAX (reserved value)
    /// - Must not exceed u64::MAX/2 (prevents overflow in calculations)
    ///
    /// # Returns
    /// - `Ok(KeyId)` for valid values
    /// - `Err(SecretSharingError::InvalidKeyId)` for invalid values
    pub fn new(value: u64) -> Result<Self, SecretSharingError> {
        // Prevent zero as it might indicate uninitialized state
        if value == 0 {
            return Err(SecretSharingError::InvalidKeyId {
                reason: "KeyId cannot be zero (reserved for uninitialized state)".to_string(),
            });
        }

        // Add reasonable upper bound to prevent overflow issues in calculations
        const MAX_REASONABLE_KEY_ID: u64 = u64::MAX / 2;
        if value > MAX_REASONABLE_KEY_ID {
            return Err(SecretSharingError::InvalidKeyId {
                reason: format!(
                    "KeyId {} exceeds maximum allowed value {} to prevent overflow issues",
                    value, MAX_REASONABLE_KEY_ID
                ),
            });
        }

        Ok(KeyId(value))
    }

    /// Safely parses a KeyId from a JSON string with validation.
    ///
    /// # Returns
    /// - `Ok(KeyId)` for valid KeyId values
    /// - `Err(SecretSharingError)` for JSON parsing errors or validation failures
    pub fn from_json_str(json_str: &str) -> Result<Self, SecretSharingError> {
        // First, try to parse as JSON
        let value: u64 =
            serde_json::from_str(json_str).map_err(|error| SecretSharingError::InvalidKeyId {
                reason: format!("Invalid JSON number format: {}", error),
            })?;

        // Then validate the parsed value
        Self::new(value)
    }

    /// Returns the inner value of the KeyId.
    pub fn value(&self) -> u64 {
        self.0
    }
}

/// Associated Authenticated Data (AAD) bound to HPKE encryption/decryption.
/// This is application-specific data that is not secret but must be
/// integrity-protected and bound into the encryption/authentication process.
///
/// # Fields:
/// - `request_id`: Unique per wrap/unwrap transaction (prevents replay).
/// - `key_id`: Which global SK is being transported (helps rotation/telemetry).
#[derive(
    Clone,
    Copy,
    Debug,
    PartialEq,
    Eq,
    serde::Serialize,
    serde::Deserialize,
    BorshSerialize,
    BorshDeserialize,
)]
pub struct Aad {
    pub request_id: u64,
    pub key_id: KeyId,
}

/// The HPKE “parcel” that carries the encrypted SK_priv to a recipient TEE.
///
/// # Fields:
/// - `sender_wrapping_pub_x25519`: The sender’s wrapping pubkey (X25519). In HPKE
///   Base mode this is NOT an authenticity proof by itself; it’s included
///   for provenance/telemetry and can be cross-checked against attestations or
///   consensus metadata by the application.
/// - `recipient_wrapping_pub_x25519`: The recipient’s wrapping pubkey (X25519).
///   We perform a hard equality check against the local TEE pubkey before
///   attempting unwrap to avoid misdelivery.
/// - `aad`: The AAD bound into the AEAD/HPKE context.
/// - `ciphertext`: 48 bytes total: `32-byte body || 16-byte AEAD tag`.
///   The body is the encrypted SK_priv; the tag authenticates both body and AAD.
///   (ChaCha20-Poly1305 has a fixed 16-byte tag.)
/// - `enc`: 32-byte HPKE “encapsulated key” (KEM output). Required to derive
///   the HPKE context on the receiver; treat as an opaque handle.
///
/// Invariants enforced by the engine:
/// - `ciphertext.len() == 48` (32 body + 16 tag)
/// - `enc.len() == 32` (for X25519 KEM)
/// - HPKE open fails if AAD, recipient key, or bytes are altered.
#[derive(
    Clone,
    Debug,
    PartialEq,
    Eq,
    serde::Serialize,
    serde::Deserialize,
    BorshSerialize,
    BorshDeserialize,
)]
pub struct WrappedKeyEnvelope {
    /// Sender's wrapping pubkey (X25519).
    pub sender_wrapping_pub_x25519: PublicKey,

    /// Recipient's wrapping pubkey (X25519).
    pub recipient_wrapping_pub_x25519: PublicKey,

    /// Associated Authenticated Data bound into the AEAD/HPKE context.
    pub aad: Aad,

    /// HPKE ciphertext (includes tag).
    #[serde(with = "BigArray")]
    pub ciphertext: [u8; 48],

    /// HPKE encapsulated key (required for open).
    #[serde(with = "BigArray")]
    pub enc: [u8; 32],
}
