//! Stealth Addresses
//!
//! Stealth addresses provide receiver privacy by allowing senders to generate
//! unique one-time addresses for each payment. The recipient can detect and
//! spend these funds without revealing their identity on-chain.
//!
//! ## How It Works
//!
//! 1. **Bob publishes** his stealth meta-address: (ViewKey, SpendKey)
//! 2. **Alice generates** a one-time stealth address for Bob:
//!    - Creates ephemeral keypair (r, R = r*G)
//!    - Computes shared secret: S = r * ViewKey
//!    - Derives one-time address: P = hash(S) * G + SpendKey
//!    - Includes R in transaction
//! 3. **Bob scans** using his view private key:
//!    - Computes S' = view_secret * R
//!    - Derives P' = hash(S') * G + SpendKey
//!    - If P' matches output address, it's his!
//! 4. **Bob spends** using derived private key:
//!    - spend_key = hash(S) + spend_secret
//!
//! ## Security Properties
//!
//! - **Unlinkability**: Outputs cannot be linked to Bob's public address
//! - **View-only wallets**: Can detect incoming payments without spending
//! - **One-time keys**: Each payment uses a unique address

use alloc::vec::Vec;
use curve25519_dalek::{
    constants::RISTRETTO_BASEPOINT_POINT,
    ristretto::{CompressedRistretto, RistrettoPoint},
    scalar::Scalar,
};
use rand_core::{CryptoRng, RngCore};
use serde::{Deserialize, Serialize};
use borsh::{BorshSerialize, BorshDeserialize};
use sha2::{Sha512, Digest};
use zeroize::Zeroize;

use crate::{Error, Result, DOMAIN_SEPARATOR};

/// Generate a random scalar using the provided RNG
fn random_scalar<R: RngCore + CryptoRng>(rng: &mut R) -> Scalar {
    let mut bytes = [0u8; 64];
    rng.fill_bytes(&mut bytes);
    Scalar::from_bytes_mod_order_wide(&bytes)
}

/// Generator point G
fn generator() -> RistrettoPoint {
    RISTRETTO_BASEPOINT_POINT
}

/// Hash a shared secret to derive a scalar
fn hash_to_scalar(shared_secret: &RistrettoPoint, extra_data: &[u8]) -> Scalar {
    let mut hasher = Sha512::new();
    hasher.update(DOMAIN_SEPARATOR);
    hasher.update(b"STEALTH_DERIVE");
    hasher.update(shared_secret.compress().as_bytes());
    hasher.update(extra_data);

    Scalar::from_hash(hasher)
}

/// A view key (public) - used to detect incoming payments
#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct ViewKey {
    point: CompressedRistretto,
}

impl ViewKey {
    /// Create from a compressed point
    pub fn from_bytes(bytes: &[u8; 32]) -> Result<Self> {
        let point = CompressedRistretto::from_slice(bytes)
            .map_err(|_| Error::InvalidViewKey("Invalid bytes".into()))?;
        // Verify the point is valid
        point
            .decompress()
            .ok_or_else(|| Error::InvalidViewKey("Point not on curve".into()))?;
        Ok(Self { point })
    }

    /// Convert to bytes
    pub fn to_bytes(&self) -> [u8; 32] {
        self.point.to_bytes()
    }

    /// Get the underlying point
    pub fn as_point(&self) -> RistrettoPoint {
        self.point.decompress().expect("ViewKey should be valid")
    }
}

impl core::fmt::Debug for ViewKey {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let bytes = self.point.to_bytes();
        write!(f, "ViewKey({:02x}{:02x}{:02x}{:02x}...)",
            bytes[0], bytes[1], bytes[2], bytes[3])
    }
}

impl BorshSerialize for ViewKey {
    fn serialize<W: borsh::io::Write>(&self, writer: &mut W) -> borsh::io::Result<()> {
        writer.write_all(&self.point.to_bytes())
    }
}

impl BorshDeserialize for ViewKey {
    fn deserialize_reader<R: borsh::io::Read>(reader: &mut R) -> borsh::io::Result<Self> {
        let mut bytes = [0u8; 32];
        reader.read_exact(&mut bytes)?;
        Self::from_bytes(&bytes)
            .map_err(|e| borsh::io::Error::new(borsh::io::ErrorKind::InvalidData, e.to_string()))
    }
}

/// A spend key (public) - represents the ability to spend
#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct SpendKey {
    point: CompressedRistretto,
}

impl SpendKey {
    /// Create from a compressed point
    pub fn from_bytes(bytes: &[u8; 32]) -> Result<Self> {
        let point = CompressedRistretto::from_slice(bytes)
            .map_err(|_| Error::InvalidSpendKey("Invalid bytes".into()))?;
        point
            .decompress()
            .ok_or_else(|| Error::InvalidSpendKey("Point not on curve".into()))?;
        Ok(Self { point })
    }

    /// Convert to bytes
    pub fn to_bytes(&self) -> [u8; 32] {
        self.point.to_bytes()
    }

    /// Get the underlying point
    pub fn as_point(&self) -> RistrettoPoint {
        self.point.decompress().expect("SpendKey should be valid")
    }
}

impl core::fmt::Debug for SpendKey {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        let bytes = self.point.to_bytes();
        write!(f, "SpendKey({:02x}{:02x}{:02x}{:02x}...)",
            bytes[0], bytes[1], bytes[2], bytes[3])
    }
}

impl BorshSerialize for SpendKey {
    fn serialize<W: borsh::io::Write>(&self, writer: &mut W) -> borsh::io::Result<()> {
        writer.write_all(&self.point.to_bytes())
    }
}

impl BorshDeserialize for SpendKey {
    fn deserialize_reader<R: borsh::io::Read>(reader: &mut R) -> borsh::io::Result<Self> {
        let mut bytes = [0u8; 32];
        reader.read_exact(&mut bytes)?;
        Self::from_bytes(&bytes)
            .map_err(|e| borsh::io::Error::new(borsh::io::ErrorKind::InvalidData, e.to_string()))
    }
}

/// A stealth meta-address (public) - published by the recipient
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct StealthMetaAddress {
    /// View key - for detecting payments
    pub view_key: ViewKey,
    /// Spend key - base for spending
    pub spend_key: SpendKey,
}

impl StealthMetaAddress {
    /// Create a new stealth meta-address
    pub fn new(view_key: ViewKey, spend_key: SpendKey) -> Self {
        Self { view_key, spend_key }
    }

    /// Convert to bytes (64 bytes total)
    pub fn to_bytes(&self) -> [u8; 64] {
        let mut bytes = [0u8; 64];
        bytes[..32].copy_from_slice(&self.view_key.to_bytes());
        bytes[32..].copy_from_slice(&self.spend_key.to_bytes());
        bytes
    }

    /// Create from bytes
    pub fn from_bytes(bytes: &[u8; 64]) -> Result<Self> {
        let view_key = ViewKey::from_bytes(bytes[..32].try_into().unwrap())?;
        let spend_key = SpendKey::from_bytes(bytes[32..].try_into().unwrap())?;
        Ok(Self { view_key, spend_key })
    }
}

/// A stealth keypair (private) - held by the recipient
#[derive(Clone, Zeroize)]
#[zeroize(drop)]
pub struct StealthKeyPair {
    /// View secret key
    view_secret: Scalar,
    /// Spend secret key
    spend_secret: Scalar,
    /// Public view key
    #[zeroize(skip)]
    view_key: ViewKey,
    /// Public spend key
    #[zeroize(skip)]
    spend_key: SpendKey,
}

impl StealthKeyPair {
    /// Generate a new random stealth keypair
    pub fn generate<R: RngCore + CryptoRng>(rng: &mut R) -> Self {
        let view_secret = random_scalar(rng);
        let spend_secret = random_scalar(rng);

        let view_point = generator() * view_secret;
        let spend_point = generator() * spend_secret;

        Self {
            view_secret,
            spend_secret,
            view_key: ViewKey {
                point: view_point.compress(),
            },
            spend_key: SpendKey {
                point: spend_point.compress(),
            },
        }
    }

    /// Create from secret keys
    pub fn from_secrets(view_secret: [u8; 32], spend_secret: [u8; 32]) -> Result<Self> {
        let view_secret = Scalar::from_canonical_bytes(view_secret)
            .into_option()
            .ok_or_else(|| Error::InvalidViewKey("Invalid view secret".into()))?;
        let spend_secret = Scalar::from_canonical_bytes(spend_secret)
            .into_option()
            .ok_or_else(|| Error::InvalidSpendKey("Invalid spend secret".into()))?;

        let view_point = generator() * view_secret;
        let spend_point = generator() * spend_secret;

        Ok(Self {
            view_secret,
            spend_secret,
            view_key: ViewKey {
                point: view_point.compress(),
            },
            spend_key: SpendKey {
                point: spend_point.compress(),
            },
        })
    }

    /// Get the public stealth meta-address
    pub fn meta_address(&self) -> StealthMetaAddress {
        StealthMetaAddress {
            view_key: self.view_key,
            spend_key: self.spend_key,
        }
    }

    /// Get the view key
    pub fn view_key(&self) -> &ViewKey {
        &self.view_key
    }

    /// Get the spend key
    pub fn spend_key(&self) -> &SpendKey {
        &self.spend_key
    }

    /// Get the view secret (for creating view-only wallets)
    pub fn view_secret(&self) -> [u8; 32] {
        self.view_secret.to_bytes()
    }

    /// Check if a stealth address belongs to this keypair
    /// Returns the spending key if it matches
    pub fn check_ownership(
        &self,
        stealth_address: &StealthAddress,
    ) -> Option<StealthSpendingKey> {
        // Compute shared secret: S = view_secret * R
        let shared_secret = stealth_address.ephemeral_pubkey.as_point() * self.view_secret;

        // Derive expected stealth address
        let derived_scalar = hash_to_scalar(&shared_secret, &[]);
        let expected_point = generator() * derived_scalar + self.spend_key.as_point();

        // Check if it matches
        if expected_point.compress() == stealth_address.address.point {
            // Compute the spending key
            let spending_scalar = derived_scalar + self.spend_secret;
            Some(StealthSpendingKey {
                scalar: spending_scalar,
            })
        } else {
            None
        }
    }

    /// Scan multiple stealth addresses for ownership (more efficient)
    pub fn scan_addresses<'a>(
        &self,
        addresses: &'a [StealthAddress],
    ) -> Vec<(&'a StealthAddress, StealthSpendingKey)> {
        addresses
            .iter()
            .filter_map(|addr| {
                self.check_ownership(addr).map(|key| (addr, key))
            })
            .collect()
    }
}

/// A view-only wallet - can detect payments but not spend
#[derive(Clone, Zeroize)]
#[zeroize(drop)]
pub struct ViewOnlyWallet {
    view_secret: Scalar,
    #[zeroize(skip)]
    view_key: ViewKey,
    #[zeroize(skip)]
    spend_key: SpendKey,
}

impl ViewOnlyWallet {
    /// Create from a stealth keypair
    pub fn from_keypair(keypair: &StealthKeyPair) -> Self {
        Self {
            view_secret: keypair.view_secret,
            view_key: keypair.view_key,
            spend_key: keypair.spend_key,
        }
    }

    /// Create from view secret and spend key
    pub fn new(view_secret: [u8; 32], spend_key: SpendKey) -> Result<Self> {
        let view_secret = Scalar::from_canonical_bytes(view_secret)
            .into_option()
            .ok_or_else(|| Error::InvalidViewKey("Invalid view secret".into()))?;

        let view_point = generator() * view_secret;

        Ok(Self {
            view_secret,
            view_key: ViewKey {
                point: view_point.compress(),
            },
            spend_key,
        })
    }

    /// Check if a stealth address belongs to this wallet
    /// Returns true if owned, but cannot provide spending key
    pub fn check_ownership(&self, stealth_address: &StealthAddress) -> bool {
        let shared_secret = stealth_address.ephemeral_pubkey.as_point() * self.view_secret;
        let derived_scalar = hash_to_scalar(&shared_secret, &[]);
        let expected_point = generator() * derived_scalar + self.spend_key.as_point();

        expected_point.compress() == stealth_address.address.point
    }

    /// Scan addresses for ownership
    pub fn scan_addresses<'a>(&self, addresses: &'a [StealthAddress]) -> Vec<&'a StealthAddress> {
        addresses
            .iter()
            .filter(|addr| self.check_ownership(addr))
            .collect()
    }
}

/// A one-time stealth address
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct StealthAddress {
    /// The one-time address (P)
    pub address: SpendKey, // Reusing SpendKey type as it's the same format
    /// Ephemeral public key (R) - included in transaction
    pub ephemeral_pubkey: ViewKey, // Reusing ViewKey type
}

impl StealthAddress {
    /// Generate a stealth address for a recipient
    pub fn generate<R: RngCore + CryptoRng>(
        recipient: &StealthMetaAddress,
        rng: &mut R,
    ) -> Self {
        // Generate ephemeral keypair
        let ephemeral_secret = random_scalar(rng);
        let ephemeral_pubkey = generator() * ephemeral_secret;

        // Compute shared secret
        let shared_secret = recipient.view_key.as_point() * ephemeral_secret;

        // Derive one-time address
        let derived_scalar = hash_to_scalar(&shared_secret, &[]);
        let stealth_point = generator() * derived_scalar + recipient.spend_key.as_point();

        Self {
            address: SpendKey {
                point: stealth_point.compress(),
            },
            ephemeral_pubkey: ViewKey {
                point: ephemeral_pubkey.compress(),
            },
        }
    }

    /// Get the address as bytes (32 bytes)
    pub fn address_bytes(&self) -> [u8; 32] {
        self.address.to_bytes()
    }

    /// Convert to full bytes (64 bytes - address + ephemeral)
    pub fn to_bytes(&self) -> [u8; 64] {
        let mut bytes = [0u8; 64];
        bytes[..32].copy_from_slice(&self.address.to_bytes());
        bytes[32..].copy_from_slice(&self.ephemeral_pubkey.to_bytes());
        bytes
    }

    /// Create from bytes
    pub fn from_bytes(bytes: &[u8; 64]) -> Result<Self> {
        let address = SpendKey::from_bytes(bytes[..32].try_into().unwrap())?;
        let ephemeral_pubkey = ViewKey::from_bytes(bytes[32..].try_into().unwrap())?;
        Ok(Self {
            address,
            ephemeral_pubkey,
        })
    }
}

/// A spending key for a stealth address
#[derive(Clone, Zeroize)]
#[zeroize(drop)]
pub struct StealthSpendingKey {
    scalar: Scalar,
}

impl StealthSpendingKey {
    /// Get the corresponding public key
    pub fn public_key(&self) -> SpendKey {
        let point = generator() * self.scalar;
        SpendKey {
            point: point.compress(),
        }
    }

    /// Sign a message (simple Schnorr signature for demonstration)
    pub fn sign<R: RngCore + CryptoRng>(&self, message: &[u8], rng: &mut R) -> StealthSignature {
        // k = random nonce
        let k = random_scalar(rng);
        let r_point = generator() * k;

        // e = H(R || P || m)
        let mut hasher = Sha512::new();
        hasher.update(DOMAIN_SEPARATOR);
        hasher.update(b"STEALTH_SIG");
        hasher.update(r_point.compress().as_bytes());
        hasher.update(self.public_key().to_bytes());
        hasher.update(message);
        let e = Scalar::from_hash(hasher);

        // s = k + e * x
        let s = k + e * self.scalar;

        StealthSignature {
            r: r_point.compress().to_bytes(),
            s: s.to_bytes(),
        }
    }

    /// Get the scalar value (for advanced use)
    pub fn as_scalar(&self) -> &Scalar {
        &self.scalar
    }

    /// Convert to bytes
    pub fn to_bytes(&self) -> [u8; 32] {
        self.scalar.to_bytes()
    }
}

/// A signature from a stealth spending key
#[derive(Clone, Debug, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct StealthSignature {
    r: [u8; 32],
    s: [u8; 32],
}

impl StealthSignature {
    /// Verify the signature against a public key
    pub fn verify(&self, public_key: &SpendKey, message: &[u8]) -> bool {
        let r_point = match CompressedRistretto::from_slice(&self.r)
            .ok()
            .and_then(|c| c.decompress())
        {
            Some(p) => p,
            None => return false,
        };

        let s = match Scalar::from_canonical_bytes(self.s).into_option() {
            Some(s) => s,
            None => return false,
        };

        // e = H(R || P || m)
        let mut hasher = Sha512::new();
        hasher.update(DOMAIN_SEPARATOR);
        hasher.update(b"STEALTH_SIG");
        hasher.update(&self.r);
        hasher.update(public_key.to_bytes());
        hasher.update(message);
        let e = Scalar::from_hash(hasher);

        // Verify: s*G = R + e*P
        let left = generator() * s;
        let right = r_point + public_key.as_point() * e;

        left == right
    }
}

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

    #[test]
    fn test_keypair_generation() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);

        let meta = keypair.meta_address();
        assert_ne!(meta.view_key.to_bytes(), meta.spend_key.to_bytes());
    }

    #[test]
    fn test_stealth_address_generation() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);
        let meta = keypair.meta_address();

        let stealth = StealthAddress::generate(&meta, &mut rng);

        // Different calls produce different addresses
        let stealth2 = StealthAddress::generate(&meta, &mut rng);
        assert_ne!(stealth.address.to_bytes(), stealth2.address.to_bytes());
    }

    #[test]
    fn test_ownership_detection() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);
        let meta = keypair.meta_address();

        let stealth = StealthAddress::generate(&meta, &mut rng);

        // Owner can detect
        let spending_key = keypair.check_ownership(&stealth);
        assert!(spending_key.is_some());

        // Different keypair cannot
        let other = StealthKeyPair::generate(&mut rng);
        assert!(other.check_ownership(&stealth).is_none());
    }

    #[test]
    fn test_spending_key_correctness() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);
        let meta = keypair.meta_address();

        let stealth = StealthAddress::generate(&meta, &mut rng);
        let spending_key = keypair.check_ownership(&stealth).unwrap();

        // Spending key should correspond to stealth address
        assert_eq!(spending_key.public_key().to_bytes(), stealth.address.to_bytes());
    }

    #[test]
    fn test_view_only_wallet() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);
        let meta = keypair.meta_address();

        let view_only = ViewOnlyWallet::from_keypair(&keypair);

        let stealth = StealthAddress::generate(&meta, &mut rng);

        // View-only can detect ownership
        assert!(view_only.check_ownership(&stealth));

        // But can't for other's addresses
        let other_keypair = StealthKeyPair::generate(&mut rng);
        let other_meta = other_keypair.meta_address();
        let other_stealth = StealthAddress::generate(&other_meta, &mut rng);
        assert!(!view_only.check_ownership(&other_stealth));
    }

    #[test]
    fn test_stealth_signature() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);
        let meta = keypair.meta_address();

        let stealth = StealthAddress::generate(&meta, &mut rng);
        let spending_key = keypair.check_ownership(&stealth).unwrap();

        let message = b"Hello, Synor!";
        let signature = spending_key.sign(message, &mut rng);

        // Valid signature
        assert!(signature.verify(&stealth.address, message));

        // Wrong message fails
        assert!(!signature.verify(&stealth.address, b"Wrong message"));

        // Wrong key fails
        let other_key = SpendKey::from_bytes(&[1u8; 32]).unwrap_or(stealth.address);
        if other_key != stealth.address {
            assert!(!signature.verify(&other_key, message));
        }
    }

    #[test]
    fn test_scan_multiple_addresses() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);
        let meta = keypair.meta_address();

        let other = StealthKeyPair::generate(&mut rng);
        let other_meta = other.meta_address();

        // Create mix of addresses
        let mut addresses = Vec::new();
        addresses.push(StealthAddress::generate(&meta, &mut rng)); // Mine
        addresses.push(StealthAddress::generate(&other_meta, &mut rng)); // Not mine
        addresses.push(StealthAddress::generate(&meta, &mut rng)); // Mine
        addresses.push(StealthAddress::generate(&other_meta, &mut rng)); // Not mine

        // Scan should find exactly 2
        let found = keypair.scan_addresses(&addresses);
        assert_eq!(found.len(), 2);
    }

    #[test]
    fn test_serialization() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);
        let meta = keypair.meta_address();
        let stealth = StealthAddress::generate(&meta, &mut rng);

        // Borsh serialization
        let bytes = borsh::to_vec(&stealth).unwrap();
        let recovered: StealthAddress = borsh::from_slice(&bytes).unwrap();

        assert_eq!(stealth.address.to_bytes(), recovered.address.to_bytes());
        assert_eq!(
            stealth.ephemeral_pubkey.to_bytes(),
            recovered.ephemeral_pubkey.to_bytes()
        );
    }

    #[test]
    fn test_meta_address_bytes() {
        let mut rng = OsRng;
        let keypair = StealthKeyPair::generate(&mut rng);
        let meta = keypair.meta_address();

        let bytes = meta.to_bytes();
        let recovered = StealthMetaAddress::from_bytes(&bytes).unwrap();

        assert_eq!(meta.view_key.to_bytes(), recovered.view_key.to_bytes());
        assert_eq!(meta.spend_key.to_bytes(), recovered.spend_key.to_bytes());
    }
}
