// Copyright 2026 Thomas Zuyev

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[cfg(feature = "serde")]
use serde::{Deserialize, Deserializer, Serialize, Serializer};

/// Module containing XChaCha20-Poly1305 encryption/decryption logic.
///
/// XChaCha20-Poly1305:
/// - 256-bit key
/// - 192-bit (24-byte) nonce
/// - Authenticated encryption (AEAD)
/// - Safe for randomly generated nonces due to extended nonce size
pub(crate) mod xchacha20poly1305 {
    use aead::Error;
    use chacha20poly1305::{
        XChaCha20Poly1305,
        aead::{Aead, AeadCore, KeyInit, OsRng},
    };

    pub const NONCE_LEN: usize = 24;

    /// Encrypts `plaintext` using the provided 32-byte key.
    ///
    /// Output format:
    ///     [ nonce (24 bytes) | ciphertext + authentication tag ]
    ///
    /// The authentication tag is appended automatically by the AEAD implementation.
    pub fn encrypt(plaintext: &[u8], key: &[u8; 32]) -> Result<Vec<u8>, Error> {
        let cipher = XChaCha20Poly1305::new(key.into());
        let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng);
        let ciphertext = cipher.encrypt(&nonce, plaintext)?;

        // Output layout: [nonce (24 bytes)] + [ciphertext]
        let mut res = Vec::with_capacity(NONCE_LEN + ciphertext.len());
        res.extend_from_slice(&nonce);
        res.extend_from_slice(&ciphertext);

        Ok(res)
    }

    /// Decrypts ciphertext previously produced by `encrypt`.
    ///
    /// Expects layout:
    ///     [ nonce (24 bytes) | ciphertext + tag ]
    ///
    /// Returns an AEAD error if:
    /// - authentication fails
    /// - nonce length is invalid
    /// - ciphertext is malformed
    pub fn decrypt(ciphertext: &[u8], key: &[u8; 32]) -> Result<Vec<u8>, Error> {
        let Some((head, tail)) = ciphertext.split_at_checked(NONCE_LEN) else {
            return Err(Error);
        };
        let cipher = XChaCha20Poly1305::new(key.into());
        let plaintext = cipher.decrypt(head.try_into().map_err(|_| Error)?, tail)?;
        Ok(plaintext)
    }
}

/// Module containing AES-256-GCM-SIV encryption/decryption logic.
///
/// AES-256-GCM-SIV:
/// - 256-bit key
/// - 96-bit (12-byte) nonce
/// - Misuse-resistant AEAD (safer than AES-GCM if nonce reuse happens)
/// - Deterministic per nonce
pub(crate) mod aes256_gcm_siv {
    use aead::Error;
    use aes_gcm_siv::{
        AeadCore, Aes256GcmSiv,
        aead::{Aead, KeyInit, OsRng},
    };

    pub const NONCE_LEN: usize = 12;

    /// Encrypts plaintext with AES-256-GCM-SIV.
    ///
    /// Output format:
    ///     [ nonce (12 bytes) | ciphertext + authentication tag ]
    pub fn encrypt(plaintext: &[u8], key: &[u8; 32]) -> Result<Vec<u8>, Error> {
        let cipher = Aes256GcmSiv::new(key.into());
        let nonce = Aes256GcmSiv::generate_nonce(&mut OsRng);
        let ciphertext = cipher.encrypt(&nonce, plaintext)?;

        // Output layout: [nonce (12 bytes)] + [ciphertext]
        let mut res = Vec::with_capacity(NONCE_LEN + ciphertext.len());
        res.extend_from_slice(&nonce);
        res.extend_from_slice(&ciphertext);

        Ok(res)
    }

    /// Decrypts data produced by `encrypt`.
    ///
    /// Fails if:
    /// - authentication tag mismatch
    /// - wrong key
    /// - corrupted data
    pub fn decrypt(ciphertext: &[u8], key: &[u8; 32]) -> Result<Vec<u8>, Error> {
        let Some((head, tail)) = ciphertext.split_at_checked(NONCE_LEN) else {
            return Err(Error);
        };
        let cipher = Aes256GcmSiv::new(key.into());
        let plaintext = cipher.decrypt(head.try_into().map_err(|_| Error)?, tail)?;
        Ok(plaintext)
    }
}

/// Unified error type for this module.
///
/// This abstracts:
/// - Unknown cipher identifiers
/// - Key length validation failures
/// - AEAD cryptographic failures
#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Unknown algorithm specified: {0}")]
    UnknownAlgorithm(u8),
    #[error("Algorithm is not specified")]
    MissingAlgorithm,
    #[error("Invalid key length")]
    InvalidKeyLength,
    #[error("AEAD error")]
    Aead(aead::Error),
}

impl From<aead::Error> for Error {
    fn from(value: aead::Error) -> Self {
        Self::Aead(value)
    }
}

pub type Result<T> = std::result::Result<T, Error>;

/// Supported cipher suites.
///
/// `repr(u8)` ensures stable binary representation,
/// allowing it to be serialized as a single byte prefix.
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[repr(u8)]
pub enum CipherSuite {
    XChaCha20Poly1305 = 0,
    Aes256GcmSiv = 1,
}

impl std::fmt::Display for CipherSuite {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::XChaCha20Poly1305 => {
                write!(f, "XChaCha20-Poly1305")
            }
            Self::Aes256GcmSiv => {
                write!(f, "AES-256-GCM-SIV")
            }
        }
    }
}

impl TryFrom<u8> for CipherSuite {
    type Error = Error;
    fn try_from(value: u8) -> std::result::Result<Self, Self::Error> {
        match value {
            0 => Ok(Self::XChaCha20Poly1305),
            1 => Ok(Self::Aes256GcmSiv),
            algo => Err(Error::UnknownAlgorithm(algo)),
        }
    }
}

/// Opaque encrypted value container.
///
/// Memory safety considerations:
/// - `payload` is zeroized on drop
/// - `algo` is intentionally NOT zeroized (non-sensitive metadata)
/// - Debug/Display never expose raw ciphertext
#[derive(zeroize::ZeroizeOnDrop, Clone)]
pub struct EncryptedSecret {
    /// Cipher suite used for encryption.
    /// Marked `skip` so it is not zeroized (not secret material).
    #[zeroize(skip)]
    algo: CipherSuite,

    /// Raw encrypted payload:
    ///     [ nonce | ciphertext | tag ]
    ///
    /// This buffer is zeroized on drop.
    payload: Vec<u8>,
}

impl EncryptedSecret {
    /// Construct from serialized format:
    ///     [ algo (1 byte) | payload... ]
    pub fn from_encrypted_bytes(input: &[u8]) -> Result<Self> {
        let (algo, payload) = input.split_first().ok_or(Error::MissingAlgorithm)?;
        let algo: CipherSuite = (*algo).try_into()?;

        Ok(Self {
            algo,
            payload: payload.to_vec(),
        })
    }

    /// Serialize into format:
    ///     [ algo (1 byte) | payload... ]
    pub fn to_encrypted_bytes(&self) -> Vec<u8> {
        let mut res = Vec::with_capacity(1 + self.payload.len());

        res.push(self.algo as u8);
        res.extend_from_slice(&self.payload);
        res
    }

    /// Encrypt input data using selected algorithm.
    ///
    /// Validates key length at runtime.
    pub fn encrypt(input: &[u8], key: &[u8], algo: CipherSuite) -> Result<Self> {
        let payload = match algo {
            CipherSuite::XChaCha20Poly1305 => {
                xchacha20poly1305::encrypt(input, Self::into_key(key)?)?
            }
            CipherSuite::Aes256GcmSiv => aes256_gcm_siv::encrypt(input, Self::into_key(key)?)?,
        };

        Ok(Self { algo, payload })
    }

    /// Decrypt using stored algorithm.
    ///
    /// Fails if:
    /// - wrong key
    /// - tampered ciphertext
    /// - truncated payload
    pub fn decrypt(&self, key: &[u8]) -> Result<Vec<u8>> {
        match self.algo {
            CipherSuite::XChaCha20Poly1305 => Ok(xchacha20poly1305::decrypt(
                &self.payload,
                Self::into_key(key)?,
            )?),
            CipherSuite::Aes256GcmSiv => Ok(aes256_gcm_siv::decrypt(
                &self.payload,
                Self::into_key(key)?,
            )?),
        }
    }

    /// Returns cipher suite used
    pub fn algorithm(&self) -> CipherSuite {
        self.algo
    }

    /// Convert raw key slice into required fixed-size type.
    ///
    /// This allows generic conversion into `[u8; 32]`
    /// and ensures key length validation is centralized.
    fn into_key<'a, T: TryFrom<&'a [u8]>>(key: &'a [u8]) -> Result<T> {
        key.try_into().map_err(|_| Error::InvalidKeyLength)
    }
}

impl std::fmt::Debug for EncryptedSecret {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "<EncryptedSecret algo={} redacted>", self.algo)
    }
}

impl std::fmt::Display for EncryptedSecret {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "<EncryptedSecret algo={} redacted>", self.algo)
    }
}

#[cfg(feature = "serde")]
impl Serialize for EncryptedSecret {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_bytes(&self.to_encrypted_bytes())
    }
}

#[cfg(feature = "serde")]
impl<'de> Deserialize<'de> for EncryptedSecret {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let bytes = Vec::deserialize(deserializer)?;
        Self::from_encrypted_bytes(&bytes).map_err(serde::de::Error::custom)
    }
}

#[cfg(feature = "sqlx")]
mod sqlx_impl {
    use super::*;
    use sqlx::{Database, Decode, Encode, Type, encode::IsNull, error::BoxDynError};

    #[cfg(feature = "sqlx_postgres")]
    impl sqlx::postgres::PgHasArrayType for EncryptedSecret {
        fn array_type_info() -> sqlx::postgres::PgTypeInfo {
            // Under the hood, we encode this as a byte array,
            // so our Postgres array type is an array of BYTEA.
            <Vec<u8> as sqlx::postgres::PgHasArrayType>::array_type_info()
        }
    }

    impl<'a, DB: Database> Type<DB> for EncryptedSecret
    where
        &'a [u8]: Type<DB>,
    {
        fn type_info() -> <DB as Database>::TypeInfo {
            <&[u8] as Type<DB>>::type_info()
        }

        fn compatible(ty: &<DB as Database>::TypeInfo) -> bool {
            <&[u8] as Type<DB>>::compatible(ty)
        }
    }

    impl<'a, DB: Database> Encode<'a, DB> for EncryptedSecret
    where
        Vec<u8>: Encode<'a, DB>,
    {
        fn encode_by_ref(
            &self,
            buf: &mut <DB as Database>::ArgumentBuffer<'a>,
        ) -> std::result::Result<IsNull, BoxDynError> {
            <Vec<u8> as Encode<'a, DB>>::encode_by_ref(&self.to_encrypted_bytes(), buf)
        }
        fn produces(&self) -> Option<<DB as Database>::TypeInfo> {
            <Vec<u8> as Encode<'a, DB>>::produces(&self.to_encrypted_bytes())
        }
        fn size_hint(&self) -> usize {
            <Vec<u8> as Encode<'a, DB>>::size_hint(&self.to_encrypted_bytes())
        }
    }

    impl<'a, DB: Database> Decode<'a, DB> for EncryptedSecret
    where
        Vec<u8>: Decode<'a, DB>,
    {
        fn decode(value: <DB as Database>::ValueRef<'a>) -> std::result::Result<Self, BoxDynError> {
            Ok(Self::from_encrypted_bytes(
                &<Vec<u8> as Decode<'a, DB>>::decode(value)?,
            )?)
        }
    }
}

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

    const TEST_KEY: &[u8; 32] = b"thisisa32bytekeyforrustcryptolib";
    const TEST_MESSAGE: &[u8] = b"Super secret message";

    #[test]
    fn test_xchacha20poly1305_roundtrip() {
        let encrypted =
            EncryptedSecret::encrypt(TEST_MESSAGE, TEST_KEY, CipherSuite::XChaCha20Poly1305)
                .expect("Encryption failed");

        assert_eq!(encrypted.algorithm(), CipherSuite::XChaCha20Poly1305);

        let decrypted = encrypted.decrypt(TEST_KEY).expect("Decryption failed");
        assert_eq!(decrypted, TEST_MESSAGE);
    }

    #[test]
    fn test_aes256_gcm_siv_roundtrip() {
        let encrypted = EncryptedSecret::encrypt(TEST_MESSAGE, TEST_KEY, CipherSuite::Aes256GcmSiv)
            .expect("Encryption failed");

        assert_eq!(encrypted.algorithm(), CipherSuite::Aes256GcmSiv);

        let decrypted = encrypted.decrypt(TEST_KEY).expect("Decryption failed");
        assert_eq!(decrypted, TEST_MESSAGE);
    }

    #[test]
    fn test_invalid_key_length() {
        let short_key = b"tooshort";
        let res = EncryptedSecret::encrypt(TEST_MESSAGE, short_key, CipherSuite::XChaCha20Poly1305);

        assert!(matches!(res, Err(Error::InvalidKeyLength)));
    }

    #[test]
    fn test_wrong_key_decryption() {
        let wrong_key: &[u8; 32] = b"thisisaWRONGbytekeyforrustcrypto";
        let encrypted =
            EncryptedSecret::encrypt(TEST_MESSAGE, TEST_KEY, CipherSuite::XChaCha20Poly1305)
                .unwrap();

        let res = encrypted.decrypt(wrong_key);
        assert!(matches!(res, Err(Error::Aead(_))));
    }

    #[test]
    fn test_tampered_ciphertext_fails() {
        let mut encrypted =
            EncryptedSecret::encrypt(TEST_MESSAGE, TEST_KEY, CipherSuite::Aes256GcmSiv).unwrap();

        // Tamper with the last byte (usually the authentication tag)
        let last_idx = encrypted.payload.len() - 1;
        encrypted.payload[last_idx] ^= 1;

        let res = encrypted.decrypt(TEST_KEY);
        assert!(matches!(res, Err(Error::Aead(_))));
    }

    #[test]
    fn test_byte_serialization_roundtrip() {
        let original =
            EncryptedSecret::encrypt(TEST_MESSAGE, TEST_KEY, CipherSuite::XChaCha20Poly1305)
                .unwrap();

        let bytes = original.to_encrypted_bytes();
        let reconstructed = EncryptedSecret::from_encrypted_bytes(&bytes).unwrap();

        assert_eq!(original.algorithm(), reconstructed.algorithm());
        assert_eq!(original.payload, reconstructed.payload);

        let decrypted = reconstructed.decrypt(TEST_KEY).unwrap();
        assert_eq!(decrypted, TEST_MESSAGE);
    }

    #[test]
    fn test_missing_algorithm_byte() {
        let empty_bytes: &[u8] = &[];
        let res = EncryptedSecret::from_encrypted_bytes(empty_bytes);

        assert!(matches!(res, Err(Error::MissingAlgorithm)));
    }

    #[test]
    fn test_unknown_algorithm_byte() {
        // ID 99 doesn't exist
        let invalid_bytes: &[u8] = &[99, 1, 2, 3];
        let res = EncryptedSecret::from_encrypted_bytes(invalid_bytes);

        assert!(matches!(res, Err(Error::UnknownAlgorithm(99))));
    }
}
