use blake2::Blake2b512;
use curve25519_dalek::{
    constants::RISTRETTO_BASEPOINT_TABLE,
    ristretto::{CompressedRistretto, RistrettoPoint},
    scalar::Scalar,
};
use digest::Digest;
use rand::{CryptoRng, RngCore};
use zeroize::{Zeroize, ZeroizeOnDrop};

use crate::{Error, Result};

const SCHNORR_DOMAIN: &[u8] = b"zkac-schnorr-v1";

pub const SIGNATURE_LEN: usize = 64;

#[derive(Clone, Zeroize, ZeroizeOnDrop)]
pub struct SecretKey {
    scalar: Scalar,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PublicKey {
    pub(crate) point: RistrettoPoint,
    pub(crate) compressed: CompressedRistretto,
}

pub struct Signature {
    r: CompressedRistretto,
    s: Scalar,
}

#[derive(Zeroize, ZeroizeOnDrop)]
pub struct Keypair {
    secret: SecretKey,
    #[zeroize(skip)]
    public: PublicKey,
}

/// Hash (R || pk || msg) into a Scalar challenge.
fn challenge(r: &CompressedRistretto, pk: &CompressedRistretto, msg: &[u8]) -> Scalar {
    let mut h = Blake2b512::new();
    h.update(SCHNORR_DOMAIN);
    h.update(r.as_bytes());
    h.update(pk.as_bytes());
    h.update(msg);
    Scalar::from_hash(h)
}

impl Keypair {
    pub fn generate<R: CryptoRng + RngCore>(rng: &mut R) -> Self {
        let scalar = Scalar::random(rng);
        Self::from_scalar(scalar)
    }

    /// 32-byte canonical encoding of the secret scalar (for persistence).
    pub fn secret_key_bytes(&self) -> [u8; 32] {
        self.secret.scalar.to_bytes()
    }

    /// Restore from [`secret_key_bytes`](Self::secret_key_bytes).
    pub fn from_secret_key_bytes(bytes: &[u8; 32]) -> Result<Self> {
        let scalar = Option::from(Scalar::from_canonical_bytes(*bytes))
            .ok_or_else(|| Error::DeserializationError("invalid secret key scalar"))?;
        Ok(Self::from_scalar(scalar))
    }

    fn from_scalar(scalar: Scalar) -> Self {
        let point = &scalar * RISTRETTO_BASEPOINT_TABLE;
        Keypair {
            secret: SecretKey { scalar },
            public: PublicKey {
                point,
                compressed: point.compress(),
            },
        }
    }

    pub fn public(&self) -> &PublicKey {
        &self.public
    }

    pub fn sign(&self, msg: &[u8]) -> Signature {
        // Deterministic nonce: H("zkac-schnorr-nonce" || sk || msg) → scalar k.
        // Avoids dependency on RNG quality at signing time.
        let mut nh = Blake2b512::new();
        nh.update(b"zkac-schnorr-nonce");
        nh.update(self.secret.scalar.as_bytes());
        nh.update(msg);
        let k = Scalar::from_hash(nh);

        let r_point = &k * RISTRETTO_BASEPOINT_TABLE;
        let r_compressed = r_point.compress();
        let e = challenge(&r_compressed, &self.public.compressed, msg);
        let s = k + e * self.secret.scalar;

        Signature { r: r_compressed, s }
    }
}

impl PublicKey {
    pub fn to_bytes(&self) -> [u8; 32] {
        self.compressed.to_bytes()
    }

    pub fn from_bytes(bytes: [u8; 32]) -> Result<Self> {
        let compressed = CompressedRistretto::from_slice(&bytes)
            .map_err(|_| Error::DeserializationError("invalid public key length"))?;
        let point = compressed
            .decompress()
            .ok_or(Error::DeserializationError("invalid ristretto point"))?;
        Ok(PublicKey { point, compressed })
    }

    pub fn as_compressed(&self) -> &CompressedRistretto {
        &self.compressed
    }

    pub fn verify(&self, msg: &[u8], sig: &Signature) -> Result<()> {
        let r_point = sig.r
            .decompress()
            .ok_or(Error::DeserializationError("invalid signature R point"))?;
        let e = challenge(&sig.r, &self.compressed, msg);
        // Check: s * G == R + e * pk
        let lhs = &sig.s * RISTRETTO_BASEPOINT_TABLE;
        let rhs = r_point + e * self.point;
        if lhs == rhs {
            Ok(())
        } else {
            Err(Error::IdentityVerificationFailed("schnorr signature invalid"))
        }
    }
}

impl Signature {
    pub fn to_bytes(&self) -> [u8; SIGNATURE_LEN] {
        let mut buf = [0u8; SIGNATURE_LEN];
        buf[..32].copy_from_slice(self.r.as_bytes());
        buf[32..].copy_from_slice(self.s.as_bytes());
        buf
    }

    pub fn from_bytes(bytes: &[u8; SIGNATURE_LEN]) -> Result<Self> {
        let r = CompressedRistretto::from_slice(&bytes[..32])
            .map_err(|_| Error::DeserializationError("invalid signature R"))?;
        let s_bytes: [u8; 32] = bytes[32..].try_into().unwrap();
        let s = Option::from(Scalar::from_canonical_bytes(s_bytes))
            .ok_or(Error::DeserializationError("invalid signature scalar"))?;
        Ok(Signature { r, s })
    }
}

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

    #[test]
    fn pubkey_serialization() {
        let kp = Keypair::generate(&mut OsRng);
        let bytes = kp.public().to_bytes();
        let pk2 = PublicKey::from_bytes(bytes).unwrap();
        assert_eq!(*kp.public(), pk2);
    }

    #[test]
    fn sign_verify_roundtrip() {
        let kp = Keypair::generate(&mut OsRng);
        let msg = b"hello world";
        let sig = kp.sign(msg);
        kp.public().verify(msg, &sig).unwrap();
    }

    #[test]
    fn signature_serialization() {
        let kp = Keypair::generate(&mut OsRng);
        let sig = kp.sign(b"test");
        let bytes = sig.to_bytes();
        let sig2 = Signature::from_bytes(&bytes).unwrap();
        kp.public().verify(b"test", &sig2).unwrap();
    }

    #[test]
    fn wrong_message_rejected() {
        let kp = Keypair::generate(&mut OsRng);
        let sig = kp.sign(b"correct");
        assert!(kp.public().verify(b"wrong", &sig).is_err());
    }

    #[test]
    fn wrong_key_rejected() {
        let kp1 = Keypair::generate(&mut OsRng);
        let kp2 = Keypair::generate(&mut OsRng);
        let sig = kp1.sign(b"msg");
        assert!(kp2.public().verify(b"msg", &sig).is_err());
    }

    #[test]
    fn deterministic_signature() {
        let kp = Keypair::generate(&mut OsRng);
        let s1 = kp.sign(b"same msg");
        let s2 = kp.sign(b"same msg");
        assert_eq!(s1.to_bytes(), s2.to_bytes());
    }

    #[test]
    fn keypair_secret_roundtrip() {
        let kp = Keypair::generate(&mut OsRng);
        let bytes = kp.secret_key_bytes();
        let kp2 = Keypair::from_secret_key_bytes(&bytes).unwrap();
        assert_eq!(kp.public().to_bytes(), kp2.public().to_bytes());
        assert_eq!(kp.sign(b"m").to_bytes(), kp2.sign(b"m").to_bytes());
    }
}