ootle-rs 0.10.1

//   Copyright 2026 The Tari Project
//   SPDX-License-Identifier: BSD-3-Clause

use async_trait::async_trait;
use ootle_byte_type::{FromByteType, ToByteType};
use signature::{Keypair, hazmat::PrehashSigner};
use tari_crypto::{
    keys::{PublicKey, SecretKey},
    ristretto::{RistrettoPublicKey, RistrettoSchnorr, RistrettoSecretKey},
};
use tari_ootle_address::{Network, OotleAddress};
use tari_ootle_common_types::{base_layer_hashing::encrypted_data_hasher, engine_types::crypto::OutputBody};
use tari_ootle_wallet_crypto::{DecryptedData, StealthCryptoApi, encrypted_data, kdfs};
use tari_template_lib_types::{EncryptedData, crypto::PedersenCommitmentBytes};

use crate::{
    key_provider,
    key_provider::{DiffieHellmanKdfKeyProvider, LocalKeyProvider, OutputMaskProvider},
    signer,
    signer::StealthKeyPrehashSigner,
    stealth::{BurnClaimKeyProvider, InputDecryptor, StealthProviderError, StealthResult},
};

#[derive(Clone)]
pub struct OotleSecretKey {
    network: Network,
    account_secret: RistrettoSecretKey,
    view_only_secret: RistrettoSecretKey,
}

impl OotleSecretKey {
    /// Create an OotleSecretKey from existing raw Ristretto secret keys.
    pub fn new(network: Network, account_secret: RistrettoSecretKey, view_only_secret: RistrettoSecretKey) -> Self {
        Self {
            network,
            account_secret,
            view_only_secret,
        }
    }

    pub fn random(network: Network) -> Self {
        let mut rng = rand::rng();
        Self::random_with(&mut rng, network)
    }

    pub fn random_with<R: rand::Rng + rand::CryptoRng>(rng: &mut R, network: Network) -> Self {
        let account_secret = RistrettoSecretKey::random(rng);
        let view_only_secret = RistrettoSecretKey::random(rng);
        Self {
            network,
            account_secret,
            view_only_secret,
        }
    }

    pub fn account_secret(&self) -> &RistrettoSecretKey {
        &self.account_secret
    }

    pub fn network(&self) -> Network {
        self.network
    }

    pub fn to_address(&self) -> OotleAddress {
        let account_pk = RistrettoPublicKey::from_secret_key(&self.account_secret);
        let view_only_pk = RistrettoPublicKey::from_secret_key(&self.view_only_secret);
        OotleAddress::new(self.network, view_only_pk.to_byte_type(), account_pk.to_byte_type())
    }

    pub fn view_only_secret(&self) -> &RistrettoSecretKey {
        &self.view_only_secret
    }
}

impl PrehashSigner<(RistrettoSchnorr, RistrettoPublicKey)> for OotleSecretKey {
    fn sign_prehash(&self, prehash: &[u8]) -> signature::Result<(RistrettoSchnorr, RistrettoPublicKey)> {
        let signature = RistrettoSchnorr::sign(&self.account_secret, prehash, &mut rand::rng())
            .expect("sign is infallible (challenge is the correct length)");
        let public_key = self.verifying_key();
        Ok((signature, public_key))
    }
}

impl StealthKeyPrehashSigner<(RistrettoSchnorr, RistrettoPublicKey)> for OotleSecretKey {
    async fn sign_prehash_with_stealth_key(
        &self,
        public_nonce: &RistrettoPublicKey,
        prehash: &[u8],
    ) -> signer::Result<(RistrettoSchnorr, RistrettoPublicKey)> {
        let secret = kdfs::owner_stealth_dh_secret(self.network(), self.account_secret(), public_nonce);
        let signature = RistrettoSchnorr::sign(&secret, prehash, &mut rand::rng())
            .expect("sign is infallible (challenge is the correct length)");
        let public_key = RistrettoPublicKey::from_secret_key(&secret);
        Ok((signature, public_key))
    }
}

#[async_trait]
impl DiffieHellmanKdfKeyProvider for LocalKeyProvider<OotleSecretKey> {
    async fn create_kdf_dh_encrypted_data_key(
        &self,
        public_key: &RistrettoPublicKey,
    ) -> key_provider::Result<RistrettoSecretKey> {
        Ok(kdfs::dh_kdf_aead(
            encrypted_data_hasher(),
            self.credentials().view_only_secret(),
            public_key,
        ))
    }
}

#[async_trait]
impl InputDecryptor for LocalKeyProvider<OotleSecretKey>
where Self: DiffieHellmanKdfKeyProvider
{
    async fn decrypt_input_data(
        &self,
        commitment: &PedersenCommitmentBytes,
        input: &OutputBody,
        skip_memo: bool,
    ) -> StealthResult<DecryptedData> {
        let sender_nonce_pk =
            input
                .public_nonce
                .try_from_byte_type()
                .map_err(|e| StealthProviderError::DecryptionFailed {
                    commitment: *commitment,
                    details: format!(
                        "Malformed public nonce in output ({e}). This should not happen because this is verified by \
                         the validators."
                    ),
                })?;

        let encryption_key = self
            .create_kdf_dh_encrypted_data_key(&sender_nonce_pk)
            .await
            .map_err(|e| StealthProviderError::DecryptionFailed {
                commitment: *commitment,
                details: format!("Failed to derive encryption key: {}", e),
            })?;

        let decrypted = encrypted_data::unblind_output(commitment, &input.encrypted_data, &encryption_key, skip_memo)
            .map_err(|e| StealthProviderError::DecryptionFailed {
            commitment: *commitment,
            details: e.to_string(),
        })?;

        Ok(decrypted)
    }
}

#[async_trait]
impl BurnClaimKeyProvider for LocalKeyProvider<OotleSecretKey> {
    async fn derive_burn_claim_secret(
        &self,
        sender_offset_public_key: &RistrettoPublicKey,
    ) -> StealthResult<RistrettoSecretKey> {
        Ok(StealthCryptoApi::new()
            .derive_burn_claim_stealth_secret(self.credentials().account_secret(), sender_offset_public_key))
    }

    async fn decrypt_burn_claim_output(
        &self,
        encrypted_data: &EncryptedData,
        commitment: &PedersenCommitmentBytes,
        sender_offset_public_key: &RistrettoPublicKey,
    ) -> StealthResult<DecryptedData> {
        // The L1 burn output's value/mask are bound to the account secret (not the view-only key), so we
        // decrypt with the account secret. `skip_memo` because the L1 output carries no memo we need here.
        let decrypted = StealthCryptoApi::new().decrypt_utxo_data(
            encrypted_data,
            commitment,
            self.credentials().account_secret(),
            sender_offset_public_key,
            true,
        )?;
        Ok(decrypted)
    }
}

impl Keypair for OotleSecretKey {
    type VerifyingKey = RistrettoPublicKey;

    fn verifying_key(&self) -> Self::VerifyingKey {
        RistrettoPublicKey::from_secret_key(&self.account_secret)
    }
}

#[async_trait]
impl OutputMaskProvider for OotleSecretKey {
    async fn next_mask(&self) -> key_provider::Result<RistrettoSecretKey> {
        // For simplicity, just generate a random mask here. Another implementation may want to derive it in some
        // deterministic way.
        Ok(RistrettoSecretKey::random(&mut rand::rng()))
    }
}