//! Ring Signatures (LSAG - Linkable Spontaneous Anonymous Group)
//!
//! Ring signatures allow a signer to sign on behalf of a group (ring) without
//! revealing which member actually signed. This provides sender privacy.
//!
//! ## How It Works
//!
//! 1. **Ring Formation**: Collect public keys (including signer's) into a ring
//! 2. **Key Image**: Compute I = x * H(P) where x is the secret key
//! 3. **Signature**: Generate signature that proves knowledge of one secret key
//! 4. **Verification**: Anyone can verify the signature is valid for the ring
//!
//! ## Key Image (Linkability)
//!
//! The key image is a one-way function of the private key. It's:
//! - **Unique**: Each private key produces exactly one key image
//! - **Deterministic**: Same key always produces same image
//! - **Unlinkable**: Cannot determine which public key produced it
//!
//! This prevents double-spending: if the same key image appears twice,
//! the second spend is rejected (even though we don't know who cheated).
//!
//! ## Security Properties
//!
//! - **Unforgeability**: Cannot sign without knowing a private key in the ring
//! - **Anonymity**: Cannot determine which ring member signed
//! - **Linkability**: Can detect if same key signed twice (via key image)

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)
}

/// Minimum ring size for meaningful anonymity
pub const MIN_RING_SIZE: usize = 2;

/// Maximum ring size (for performance)
pub const MAX_RING_SIZE: usize = 32;

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

/// Hash a point to another point (for key images)
fn hash_to_point(point: &RistrettoPoint) -> RistrettoPoint {
    let mut hasher = Sha512::new();
    hasher.update(DOMAIN_SEPARATOR);
    hasher.update(b"HASH_TO_POINT");
    hasher.update(point.compress().as_bytes());
    RistrettoPoint::from_hash(hasher)
}

/// Hash multiple values to a scalar (Fiat-Shamir transform)
fn hash_to_scalar(data: &[&[u8]]) -> Scalar {
    let mut hasher = Sha512::new();
    hasher.update(DOMAIN_SEPARATOR);
    hasher.update(b"RING_HASH");
    for d in data {
        hasher.update(d);
    }
    Scalar::from_hash(hasher)
}

/// A public key in the ring
#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct RingPublicKey {
    point: CompressedRistretto,
}

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

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

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

    /// Create from a point
    pub fn from_point(point: RistrettoPoint) -> Self {
        Self {
            point: point.compress(),
        }
    }
}

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

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

impl BorshDeserialize for RingPublicKey {
    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 private key for ring signing
#[derive(Clone, Zeroize)]
#[zeroize(drop)]
pub struct RingPrivateKey {
    scalar: Scalar,
    #[zeroize(skip)]
    public_key: RingPublicKey,
}

impl RingPrivateKey {
    /// Generate a random keypair
    pub fn generate<R: RngCore + CryptoRng>(rng: &mut R) -> Self {
        let scalar = random_scalar(rng);
        let point = generator_g() * scalar;
        Self {
            scalar,
            public_key: RingPublicKey::from_point(point),
        }
    }

    /// Create from scalar bytes
    pub fn from_bytes(bytes: &[u8; 32]) -> Result<Self> {
        let scalar = Scalar::from_canonical_bytes(*bytes)
            .into_option()
            .ok_or_else(|| Error::InvalidScalar("Invalid scalar bytes".into()))?;
        let point = generator_g() * scalar;
        Ok(Self {
            scalar,
            public_key: RingPublicKey::from_point(point),
        })
    }

    /// Get the public key
    pub fn public_key(&self) -> &RingPublicKey {
        &self.public_key
    }

    /// Compute the key image for this key
    pub fn key_image(&self) -> KeyImage {
        let hp = hash_to_point(&self.public_key.as_point());
        let image = hp * self.scalar;
        KeyImage {
            point: image.compress(),
        }
    }

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

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

/// A key image - used to detect double-spending
#[derive(Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct KeyImage {
    point: CompressedRistretto,
}

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

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

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

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

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

impl BorshDeserialize for KeyImage {
    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 ring signature (LSAG)
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RingSignature {
    /// The key image (for double-spend detection)
    pub key_image: KeyImage,
    /// The challenge value c_0
    c0: [u8; 32],
    /// The response values s_i
    responses: Vec<[u8; 32]>,
}

impl RingSignature {
    /// Sign a message with a ring signature
    ///
    /// # Arguments
    /// * `private_key` - The signer's private key
    /// * `ring` - The ring of public keys (must include signer's key)
    /// * `signer_index` - Index of signer's key in the ring
    /// * `message` - The message to sign
    /// * `rng` - Random number generator
    pub fn sign<R: RngCore + CryptoRng>(
        private_key: &RingPrivateKey,
        ring: &[RingPublicKey],
        signer_index: usize,
        message: &[u8],
        rng: &mut R,
    ) -> Result<Self> {
        let n = ring.len();

        // Validate ring
        if n < MIN_RING_SIZE {
            return Err(Error::InvalidRingSize {
                size: n,
                min: MIN_RING_SIZE,
                max: MAX_RING_SIZE,
            });
        }
        if n > MAX_RING_SIZE {
            return Err(Error::InvalidRingSize {
                size: n,
                min: MIN_RING_SIZE,
                max: MAX_RING_SIZE,
            });
        }
        if signer_index >= n {
            return Err(Error::RingSignatureFailed(
                "Signer index out of bounds".into(),
            ));
        }
        if ring[signer_index] != *private_key.public_key() {
            return Err(Error::RingSignatureFailed(
                "Signer's key not at specified index".into(),
            ));
        }

        // Compute key image
        let key_image = private_key.key_image();
        let image_point = key_image.as_point();

        // Hash the ring and message for domain separation
        let ring_bytes: Vec<u8> = ring.iter().flat_map(|k| k.to_bytes()).collect();

        // Generate random values for non-signer indices
        let mut c = vec![Scalar::ZERO; n];
        let mut s = vec![Scalar::ZERO; n];

        // Step 1: Generate random alpha for the signer
        let alpha = random_scalar(rng);

        // L_pi = alpha * G
        let l_pi = generator_g() * alpha;
        // R_pi = alpha * H(P_pi)
        let hp_pi = hash_to_point(&ring[signer_index].as_point());
        let r_pi = hp_pi * alpha;

        // Step 2: Compute c_{pi+1}
        let c_next = hash_to_scalar(&[
            message,
            &ring_bytes,
            &key_image.to_bytes(),
            &l_pi.compress().to_bytes(),
            &r_pi.compress().to_bytes(),
        ]);
        c[(signer_index + 1) % n] = c_next;

        // Step 3: For each other member, generate random s and compute c
        for offset in 1..n {
            let i = (signer_index + offset) % n;
            let next = (i + 1) % n;

            // Generate random s_i
            s[i] = random_scalar(rng);

            // L_i = s_i * G + c_i * P_i
            let l_i = generator_g() * s[i] + ring[i].as_point() * c[i];
            // R_i = s_i * H(P_i) + c_i * I
            let hp_i = hash_to_point(&ring[i].as_point());
            let r_i = hp_i * s[i] + image_point * c[i];

            // c_{i+1} = H(m, L_i, R_i)
            // For the last iteration, this computes c[signer_index] which we need
            c[next] = hash_to_scalar(&[
                message,
                &ring_bytes,
                &key_image.to_bytes(),
                &l_i.compress().to_bytes(),
                &r_i.compress().to_bytes(),
            ]);
        }

        // Step 4: Close the ring
        // s_pi = alpha - c_pi * x
        s[signer_index] = alpha - c[signer_index] * private_key.scalar;

        Ok(Self {
            key_image,
            c0: c[0].to_bytes(),
            responses: s.iter().map(|s| s.to_bytes()).collect(),
        })
    }

    /// Verify a ring signature
    pub fn verify(&self, ring: &[RingPublicKey], message: &[u8]) -> Result<bool> {
        let n = ring.len();

        if n < MIN_RING_SIZE || n > MAX_RING_SIZE {
            return Err(Error::InvalidRingSize {
                size: n,
                min: MIN_RING_SIZE,
                max: MAX_RING_SIZE,
            });
        }
        if self.responses.len() != n {
            return Err(Error::RingSignatureFailed(
                "Response count doesn't match ring size".into(),
            ));
        }

        let c0 = Scalar::from_canonical_bytes(self.c0)
            .into_option()
            .ok_or_else(|| Error::InvalidScalar("Invalid c0".into()))?;

        let image_point = self.key_image.as_point();

        let ring_bytes: Vec<u8> = ring.iter().flat_map(|k| k.to_bytes()).collect();

        let mut c_current = c0;

        for i in 0..n {
            let s_i = Scalar::from_canonical_bytes(self.responses[i])
                .into_option()
                .ok_or_else(|| Error::InvalidScalar(format!("Invalid s_{}", i)))?;

            // L_i = s_i * G + c_i * P_i
            let l_i = generator_g() * s_i + ring[i].as_point() * c_current;

            // R_i = s_i * H(P_i) + c_i * I
            let hp_i = hash_to_point(&ring[i].as_point());
            let r_i = hp_i * s_i + image_point * c_current;

            // c_{i+1} = H(m, L_i, R_i)
            c_current = hash_to_scalar(&[
                message,
                &ring_bytes,
                &self.key_image.to_bytes(),
                &l_i.compress().to_bytes(),
                &r_i.compress().to_bytes(),
            ]);
        }

        // Check if we closed the ring (c_n should equal c_0)
        Ok(c_current == c0)
    }

    /// Get the number of ring members
    pub fn ring_size(&self) -> usize {
        self.responses.len()
    }
}

impl BorshSerialize for RingSignature {
    fn serialize<W: borsh::io::Write>(&self, writer: &mut W) -> borsh::io::Result<()> {
        BorshSerialize::serialize(&self.key_image, writer)?;
        BorshSerialize::serialize(&self.c0, writer)?;
        BorshSerialize::serialize(&self.responses, writer)?;
        Ok(())
    }
}

impl BorshDeserialize for RingSignature {
    fn deserialize_reader<R: borsh::io::Read>(reader: &mut R) -> borsh::io::Result<Self> {
        let key_image = KeyImage::deserialize_reader(reader)?;
        let c0 = <[u8; 32]>::deserialize_reader(reader)?;
        let responses = Vec::<[u8; 32]>::deserialize_reader(reader)?;
        Ok(Self {
            key_image,
            c0,
            responses,
        })
    }
}

/// Key image tracker for double-spend prevention
pub struct KeyImageTracker {
    used_images: Vec<KeyImage>,
}

impl KeyImageTracker {
    /// Create a new tracker
    pub fn new() -> Self {
        Self {
            used_images: Vec::new(),
        }
    }

    /// Check if a key image has been used
    pub fn is_used(&self, image: &KeyImage) -> bool {
        self.used_images.contains(image)
    }

    /// Mark a key image as used
    /// Returns false if already used (double-spend attempt)
    pub fn mark_used(&mut self, image: KeyImage) -> bool {
        if self.is_used(&image) {
            false
        } else {
            self.used_images.push(image);
            true
        }
    }

    /// Get all used key images
    pub fn used_images(&self) -> &[KeyImage] {
        &self.used_images
    }
}

impl Default for KeyImageTracker {
    fn default() -> Self {
        Self::new()
    }
}

/// Generate a ring of decoy public keys for a transaction
pub fn generate_decoy_ring<R: RngCore + CryptoRng>(
    signer_key: &RingPublicKey,
    decoys: &[RingPublicKey],
    rng: &mut R,
) -> (Vec<RingPublicKey>, usize) {
    let mut ring: Vec<RingPublicKey> = decoys.to_vec();

    // Remove the signer if accidentally included in decoys
    ring.retain(|k| k != signer_key);

    // Insert signer at random position
    let signer_index = if ring.is_empty() {
        0
    } else {
        (rng.next_u32() as usize) % (ring.len() + 1)
    };
    ring.insert(signer_index, *signer_key);

    (ring, signer_index)
}

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

    #[test]
    fn test_keypair_generation() {
        let mut rng = OsRng;
        let key = RingPrivateKey::generate(&mut rng);
        let pubkey = key.public_key();

        // Verify public key is on curve
        assert!(RingPublicKey::from_bytes(&pubkey.to_bytes()).is_ok());
    }

    #[test]
    fn test_key_image_deterministic() {
        let mut rng = OsRng;
        let key = RingPrivateKey::generate(&mut rng);

        let image1 = key.key_image();
        let image2 = key.key_image();

        assert_eq!(image1.to_bytes(), image2.to_bytes());
    }

    #[test]
    fn test_key_image_unique() {
        let mut rng = OsRng;
        let key1 = RingPrivateKey::generate(&mut rng);
        let key2 = RingPrivateKey::generate(&mut rng);

        assert_ne!(key1.key_image().to_bytes(), key2.key_image().to_bytes());
    }

    #[test]
    fn test_ring_signature_basic() {
        let mut rng = OsRng;

        // Generate ring of 4 keys
        let keys: Vec<RingPrivateKey> = (0..4)
            .map(|_| RingPrivateKey::generate(&mut rng))
            .collect();
        let ring: Vec<RingPublicKey> = keys.iter().map(|k| *k.public_key()).collect();

        let signer_index = 2;
        let message = b"Test message";

        let signature = RingSignature::sign(
            &keys[signer_index],
            &ring,
            signer_index,
            message,
            &mut rng,
        )
        .unwrap();

        // Verify
        assert!(signature.verify(&ring, message).unwrap());
    }

    #[test]
    fn test_ring_signature_wrong_message_fails() {
        let mut rng = OsRng;

        let keys: Vec<RingPrivateKey> = (0..4)
            .map(|_| RingPrivateKey::generate(&mut rng))
            .collect();
        let ring: Vec<RingPublicKey> = keys.iter().map(|k| *k.public_key()).collect();

        let signature = RingSignature::sign(
            &keys[0],
            &ring,
            0,
            b"Correct message",
            &mut rng,
        )
        .unwrap();

        // Wrong message should fail
        assert!(!signature.verify(&ring, b"Wrong message").unwrap());
    }

    #[test]
    fn test_ring_signature_wrong_ring_fails() {
        let mut rng = OsRng;

        let keys: Vec<RingPrivateKey> = (0..4)
            .map(|_| RingPrivateKey::generate(&mut rng))
            .collect();
        let ring: Vec<RingPublicKey> = keys.iter().map(|k| *k.public_key()).collect();

        let signature = RingSignature::sign(
            &keys[0],
            &ring,
            0,
            b"Test",
            &mut rng,
        )
        .unwrap();

        // Different ring should fail
        let other_keys: Vec<RingPrivateKey> = (0..4)
            .map(|_| RingPrivateKey::generate(&mut rng))
            .collect();
        let other_ring: Vec<RingPublicKey> = other_keys.iter().map(|k| *k.public_key()).collect();

        assert!(!signature.verify(&other_ring, b"Test").unwrap());
    }

    #[test]
    fn test_key_image_in_signature() {
        let mut rng = OsRng;

        let keys: Vec<RingPrivateKey> = (0..4)
            .map(|_| RingPrivateKey::generate(&mut rng))
            .collect();
        let ring: Vec<RingPublicKey> = keys.iter().map(|k| *k.public_key()).collect();

        let signature = RingSignature::sign(
            &keys[1],
            &ring,
            1,
            b"Test",
            &mut rng,
        )
        .unwrap();

        // Key image should match signer's key image
        assert_eq!(
            signature.key_image.to_bytes(),
            keys[1].key_image().to_bytes()
        );
    }

    #[test]
    fn test_key_image_tracker() {
        let mut rng = OsRng;
        let mut tracker = KeyImageTracker::new();

        let key1 = RingPrivateKey::generate(&mut rng);
        let key2 = RingPrivateKey::generate(&mut rng);

        let image1 = key1.key_image();
        let image2 = key2.key_image();

        // First use succeeds
        assert!(tracker.mark_used(image1));
        assert!(tracker.mark_used(image2));

        // Double use fails
        assert!(!tracker.mark_used(key1.key_image()));
        assert!(!tracker.mark_used(key2.key_image()));

        // Tracker has correct count
        assert_eq!(tracker.used_images().len(), 2);
    }

    #[test]
    fn test_minimum_ring_size() {
        let mut rng = OsRng;

        let key = RingPrivateKey::generate(&mut rng);
        let ring = vec![*key.public_key()]; // Only 1 member

        let result = RingSignature::sign(&key, &ring, 0, b"Test", &mut rng);
        assert!(matches!(result, Err(Error::InvalidRingSize { .. })));
    }

    #[test]
    fn test_generate_decoy_ring() {
        let mut rng = OsRng;

        let signer = RingPrivateKey::generate(&mut rng);
        let decoys: Vec<RingPublicKey> = (0..5)
            .map(|_| *RingPrivateKey::generate(&mut rng).public_key())
            .collect();

        let (ring, signer_index) = generate_decoy_ring(signer.public_key(), &decoys, &mut rng);

        // Ring should contain signer
        assert_eq!(ring[signer_index], *signer.public_key());

        // Ring should have correct size
        assert_eq!(ring.len(), 6);

        // Signer should be able to sign
        let signature = RingSignature::sign(
            &signer,
            &ring,
            signer_index,
            b"Test",
            &mut rng,
        )
        .unwrap();
        assert!(signature.verify(&ring, b"Test").unwrap());
    }

    #[test]
    fn test_serialization() {
        let mut rng = OsRng;

        let keys: Vec<RingPrivateKey> = (0..3)
            .map(|_| RingPrivateKey::generate(&mut rng))
            .collect();
        let ring: Vec<RingPublicKey> = keys.iter().map(|k| *k.public_key()).collect();

        let signature = RingSignature::sign(&keys[0], &ring, 0, b"Test", &mut rng).unwrap();

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

        // Recovered signature should verify
        assert!(recovered.verify(&ring, b"Test").unwrap());
    }

    #[test]
    fn test_signer_at_different_indices() {
        let mut rng = OsRng;

        let keys: Vec<RingPrivateKey> = (0..5)
            .map(|_| RingPrivateKey::generate(&mut rng))
            .collect();
        let ring: Vec<RingPublicKey> = keys.iter().map(|k| *k.public_key()).collect();

        // Test signing from each position
        for i in 0..5 {
            let signature = RingSignature::sign(&keys[i], &ring, i, b"Test", &mut rng).unwrap();
            assert!(
                signature.verify(&ring, b"Test").unwrap(),
                "Failed for signer at index {}",
                i
            );
        }
    }
}