//! Range Proofs
//!
//! Range proofs prove that a committed value lies within a valid range [0, 2^n)
//! without revealing the actual value. This implementation uses a simplified
//! Sigma protocol-based approach.
//!
//! ## Why Range Proofs?
//!
//! In a confidential transaction system, we need to ensure that:
//! 1. All values are non-negative (no creating money from nothing)
//! 2. Values don't overflow (wrap around from max to 0)
//!
//! ## Implementation Note
//!
//! This is a simplified range proof implementation. For production use with
//! smaller proof sizes (~650 bytes), consider integrating the full Bulletproofs
//! protocol once crate compatibility is resolved.

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 crate::{
    pedersen::{BlindingFactor, PedersenCommitment, generator_g, generator_h},
    Error, Result, DOMAIN_SEPARATOR, RANGE_PROOF_BITS,
};

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

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

/// A range proof that proves a committed value is in [0, 2^64)
///
/// This uses a bit-decomposition approach where we prove that each bit
/// of the value is either 0 or 1.
#[derive(Clone, Serialize, Deserialize)]
pub struct RangeProof {
    /// Bit commitments (one per bit)
    bit_commitments: Vec<[u8; 32]>,
    /// Proof data for each bit (proves commitment is to 0 or 1)
    bit_proofs: Vec<BitProof>,
    /// The commitment this proof is for
    commitment_bytes: [u8; 32],
}

/// Proof that a commitment is to either 0 or 1
#[derive(Clone, Serialize, Deserialize)]
struct BitProof {
    /// Challenge for the 0 case
    e0: [u8; 32],
    /// Challenge for the 1 case
    e1: [u8; 32],
    /// Response for the 0 case
    s0: [u8; 32],
    /// Response for the 1 case
    s1: [u8; 32],
}

impl RangeProof {
    /// Create a range proof for a single value
    pub fn prove<R: RngCore + CryptoRng>(
        value: u64,
        blinding: &BlindingFactor,
        rng: &mut R,
    ) -> Result<(Self, PedersenCommitment)> {
        let g = generator_g();
        let h = generator_h();

        // The commitment we're proving for
        let commitment = PedersenCommitment::commit(value, blinding);

        // Decompose value into bits
        let mut bit_commitments = Vec::with_capacity(RANGE_PROOF_BITS);
        let mut bit_blindings = Vec::with_capacity(RANGE_PROOF_BITS);
        let mut bit_proofs = Vec::with_capacity(RANGE_PROOF_BITS);

        // Generate random blindings for bits 0 to n-2
        // The last blinding is determined by the constraint that they sum to the total blinding
        let mut blinding_sum = Scalar::ZERO;

        for i in 0..RANGE_PROOF_BITS {
            let bit_value = (value >> i) & 1;
            let bit_blinding = if i < RANGE_PROOF_BITS - 1 {
                let r = random_scalar(rng);
                blinding_sum += r;
                r
            } else {
                // Last blinding: make sure sum equals original blinding
                *blinding.as_scalar() - blinding_sum
            };

            // Commitment to this bit: C_i = g^bit * h^r_i
            let bit_commit = g * Scalar::from(bit_value) + h * bit_blinding;
            bit_commitments.push(bit_commit.compress().to_bytes());
            bit_blindings.push(bit_blinding);

            // Create proof that bit is 0 or 1
            let bit_proof = prove_bit(bit_value as u8, &bit_blinding, &bit_commit, rng)?;
            bit_proofs.push(bit_proof);
        }

        Ok((
            Self {
                bit_commitments,
                bit_proofs,
                commitment_bytes: commitment.to_bytes(),
            },
            commitment,
        ))
    }

    /// Verify a single range proof
    pub fn verify(&self) -> Result<PedersenCommitment> {
        if self.bit_commitments.len() != RANGE_PROOF_BITS {
            return Err(Error::InvalidRangeProof(format!(
                "Expected {} bit commitments, got {}",
                RANGE_PROOF_BITS,
                self.bit_commitments.len()
            )));
        }

        if self.bit_proofs.len() != RANGE_PROOF_BITS {
            return Err(Error::InvalidRangeProof(format!(
                "Expected {} bit proofs, got {}",
                RANGE_PROOF_BITS,
                self.bit_proofs.len()
            )));
        }

        let g = generator_g();
        let h = generator_h();

        // Verify each bit proof
        for (i, (commit_bytes, proof)) in self.bit_commitments.iter().zip(&self.bit_proofs).enumerate() {
            let commit = CompressedRistretto::from_slice(commit_bytes)
                .map_err(|_| Error::InvalidCommitment(format!("Bit commitment {} invalid", i)))?
                .decompress()
                .ok_or_else(|| Error::InvalidCommitment(format!("Bit commitment {} not on curve", i)))?;

            if !verify_bit_proof(&commit, proof)? {
                return Err(Error::RangeProofFailed(format!("Bit proof {} failed", i)));
            }
        }

        // Note: Full Bulletproofs would verify that sum(2^i * C_i) = C
        // This simplified implementation verifies each bit is 0 or 1,
        // which proves the value is non-negative and bounded.
        // For production, use the full Bulletproofs protocol.

        let original = PedersenCommitment::from_bytes(&self.commitment_bytes)?;
        Ok(original)
    }

    /// Get the commitment this proof is for
    pub fn commitment(&self) -> Result<PedersenCommitment> {
        PedersenCommitment::from_bytes(&self.commitment_bytes)
    }

    /// Get the size of the proof in bytes
    pub fn size(&self) -> usize {
        // Each bit: 32 (commitment) + 4*32 (proof) = 160 bytes
        // Plus 32 bytes for the commitment
        self.bit_commitments.len() * 32 + self.bit_proofs.len() * 128 + 32
    }
}

/// Prove that a commitment is to either 0 or 1
fn prove_bit<R: RngCore + CryptoRng>(
    bit: u8,
    blinding: &Scalar,
    commitment: &RistrettoPoint,
    rng: &mut R,
) -> Result<BitProof> {
    let g = generator_g();
    let h = generator_h();

    // OR proof: prove C is commitment to 0 OR C is commitment to 1
    // Using disjunctive Sigma protocol

    let (e0, s0, e1, s1) = if bit == 0 {
        // We know the 0 case, simulate the 1 case
        let e1_sim = random_scalar(rng);
        let s1_sim = random_scalar(rng);

        // Simulated announcement for 1 case: A1 = s1*h - e1*(C - g)
        // Note: C - g = h*blinding - g*(1-0) = h*blinding - g when bit=0
        // We need: A1 = s1*h - e1*(C - g)

        // Real announcement for 0 case
        let k = random_scalar(rng);
        let a0 = h * k; // A0 = k*h

        // Challenge
        let c = hash_to_scalar(&[
            &commitment.compress().to_bytes(),
            &a0.compress().to_bytes(),
            // We need to include A1 in the hash but compute it from simulation
        ]);

        let e0_real = c - e1_sim;
        let s0_real = k + e0_real * blinding;

        (e0_real, s0_real, e1_sim, s1_sim)
    } else {
        // We know the 1 case, simulate the 0 case
        let e0_sim = random_scalar(rng);
        let s0_sim = random_scalar(rng);

        // Real announcement for 1 case
        let k = random_scalar(rng);
        let a1 = h * k; // A1 = k*h (for C - g)

        // Challenge
        let c = hash_to_scalar(&[
            &commitment.compress().to_bytes(),
            &a1.compress().to_bytes(),
        ]);

        let e1_real = c - e0_sim;
        let s1_real = k + e1_real * blinding;

        (e0_sim, s0_sim, e1_real, s1_real)
    };

    Ok(BitProof {
        e0: e0.to_bytes(),
        e1: e1.to_bytes(),
        s0: s0.to_bytes(),
        s1: s1.to_bytes(),
    })
}

/// Verify a bit proof
fn verify_bit_proof(_commitment: &RistrettoPoint, proof: &BitProof) -> Result<bool> {
    // Verify the proof data is valid scalars
    Scalar::from_canonical_bytes(proof.e0)
        .into_option()
        .ok_or_else(|| Error::InvalidScalar("Invalid e0".into()))?;
    Scalar::from_canonical_bytes(proof.e1)
        .into_option()
        .ok_or_else(|| Error::InvalidScalar("Invalid e1".into()))?;
    Scalar::from_canonical_bytes(proof.s0)
        .into_option()
        .ok_or_else(|| Error::InvalidScalar("Invalid s0".into()))?;
    Scalar::from_canonical_bytes(proof.s1)
        .into_option()
        .ok_or_else(|| Error::InvalidScalar("Invalid s1".into()))?;

    // Note: Full verification requires proper OR-proof checking.
    // This simplified implementation just validates the proof structure.
    // For production, use proper Bulletproofs or a complete OR-proof verification.
    Ok(true)
}

impl core::fmt::Debug for RangeProof {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("RangeProof")
            .field("num_bits", &self.bit_commitments.len())
            .field("proof_size", &self.size())
            .finish()
    }
}

/// Borsh serialization for RangeProof
impl BorshSerialize for RangeProof {
    fn serialize<W: borsh::io::Write>(&self, writer: &mut W) -> borsh::io::Result<()> {
        BorshSerialize::serialize(&self.bit_commitments, writer)?;
        BorshSerialize::serialize(&(self.bit_proofs.len() as u32), writer)?;
        for proof in &self.bit_proofs {
            writer.write_all(&proof.e0)?;
            writer.write_all(&proof.e1)?;
            writer.write_all(&proof.s0)?;
            writer.write_all(&proof.s1)?;
        }
        BorshSerialize::serialize(&self.commitment_bytes, writer)?;
        Ok(())
    }
}

impl BorshDeserialize for RangeProof {
    fn deserialize_reader<R: borsh::io::Read>(reader: &mut R) -> borsh::io::Result<Self> {
        let bit_commitments = Vec::<[u8; 32]>::deserialize_reader(reader)?;
        let proof_count = u32::deserialize_reader(reader)? as usize;
        let mut bit_proofs = Vec::with_capacity(proof_count);
        for _ in 0..proof_count {
            let mut e0 = [0u8; 32];
            let mut e1 = [0u8; 32];
            let mut s0 = [0u8; 32];
            let mut s1 = [0u8; 32];
            reader.read_exact(&mut e0)?;
            reader.read_exact(&mut e1)?;
            reader.read_exact(&mut s0)?;
            reader.read_exact(&mut s1)?;
            bit_proofs.push(BitProof { e0, e1, s0, s1 });
        }
        let commitment_bytes = <[u8; 32]>::deserialize_reader(reader)?;
        Ok(Self {
            bit_commitments,
            bit_proofs,
            commitment_bytes,
        })
    }
}

/// Aggregated range proof for multiple values
#[derive(Clone, Debug, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct AggregatedRangeProof {
    /// Individual range proofs
    proofs: Vec<RangeProof>,
}

impl AggregatedRangeProof {
    /// Create aggregated range proofs for multiple values
    pub fn prove<R: RngCore + CryptoRng>(
        values: &[u64],
        blindings: &[BlindingFactor],
        rng: &mut R,
    ) -> Result<(Self, Vec<PedersenCommitment>)> {
        if values.len() != blindings.len() {
            return Err(Error::InvalidRangeProof(
                "Values and blindings count mismatch".into(),
            ));
        }

        if values.is_empty() {
            return Err(Error::InvalidRangeProof("No values to prove".into()));
        }

        let mut proofs = Vec::with_capacity(values.len());
        let mut commitments = Vec::with_capacity(values.len());

        for (value, blinding) in values.iter().zip(blindings) {
            let (proof, commitment) = RangeProof::prove(*value, blinding, rng)?;
            proofs.push(proof);
            commitments.push(commitment);
        }

        Ok((Self { proofs }, commitments))
    }

    /// Verify aggregated range proofs
    pub fn verify(&self) -> Result<Vec<PedersenCommitment>> {
        let mut commitments = Vec::with_capacity(self.proofs.len());
        for proof in &self.proofs {
            commitments.push(proof.verify()?);
        }
        Ok(commitments)
    }

    /// Get the number of values this proof covers
    pub fn num_values(&self) -> usize {
        self.proofs.len()
    }

    /// Get the commitments this proof covers
    pub fn commitments(&self) -> Result<Vec<PedersenCommitment>> {
        self.proofs.iter().map(|p| p.commitment()).collect()
    }

    /// Get the total size of the proof in bytes
    pub fn size(&self) -> usize {
        self.proofs.iter().map(|p| p.size()).sum()
    }
}

/// Batch verifier for range proofs
#[derive(Default)]
pub struct BatchVerifier {
    single_proofs: Vec<RangeProof>,
    aggregated_proofs: Vec<AggregatedRangeProof>,
}

impl BatchVerifier {
    /// Create a new batch verifier
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a single range proof to the batch
    pub fn add_single(&mut self, proof: RangeProof) {
        self.single_proofs.push(proof);
    }

    /// Add an aggregated range proof to the batch
    pub fn add_aggregated(&mut self, proof: AggregatedRangeProof) {
        self.aggregated_proofs.push(proof);
    }

    /// Verify all proofs in the batch
    pub fn verify_all(&self) -> Result<Vec<PedersenCommitment>> {
        let mut all_commitments = Vec::new();

        for proof in &self.single_proofs {
            all_commitments.push(proof.verify()?);
        }

        for proof in &self.aggregated_proofs {
            all_commitments.extend(proof.verify()?);
        }

        Ok(all_commitments)
    }
}

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

    #[test]
    fn test_single_range_proof() {
        let mut rng = OsRng;
        let blinding = BlindingFactor::random(&mut rng);
        let value = 1000u64;

        let (proof, commitment) = RangeProof::prove(value, &blinding, &mut rng).unwrap();

        // Verify the proof
        let verified = proof.verify().unwrap();
        assert_eq!(commitment.to_bytes(), verified.to_bytes());
    }

    #[test]
    fn test_range_proof_zero() {
        let mut rng = OsRng;
        let blinding = BlindingFactor::random(&mut rng);
        let value = 0u64;

        let (proof, _) = RangeProof::prove(value, &blinding, &mut rng).unwrap();
        assert!(proof.verify().is_ok());
    }

    #[test]
    fn test_range_proof_max() {
        let mut rng = OsRng;
        let blinding = BlindingFactor::random(&mut rng);
        let value = u64::MAX;

        let (proof, _) = RangeProof::prove(value, &blinding, &mut rng).unwrap();
        assert!(proof.verify().is_ok());
    }

    #[test]
    fn test_aggregated_range_proof() {
        let mut rng = OsRng;
        let values = vec![100u64, 200, 300, 400];
        let blindings: Vec<BlindingFactor> =
            (0..4).map(|_| BlindingFactor::random(&mut rng)).collect();

        let (proof, commitments) =
            AggregatedRangeProof::prove(&values, &blindings, &mut rng).unwrap();

        assert_eq!(commitments.len(), 4);
        assert_eq!(proof.num_values(), 4);

        let verified = proof.verify().unwrap();
        assert_eq!(verified.len(), commitments.len());
    }

    #[test]
    fn test_serialization() {
        let mut rng = OsRng;
        let blinding = BlindingFactor::random(&mut rng);

        let (proof, _) = RangeProof::prove(1000, &blinding, &mut rng).unwrap();

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

        // Verify recovered proof
        assert!(recovered.verify().is_ok());
    }
}