//! Fiat-Shamir transcript for non-interactive proofs.
//!
//! This module wraps the production-grade [Merlin](https://merlin.cool) transcript
//! library, which provides STROBE-based transcript construction with automatic
//! domain separation and message framing.

use crate::traits::PairingEngine;

/// Protocol version for transcript domain separation.
pub const TRANSCRIPT_VERSION: &str = "SAMAHARAM_V2";

/// Transcript for Fiat-Shamir transform.
///
/// Wraps `merlin::Transcript` to provide type-safe methods for
/// appending elliptic curve points and field elements.
///
/// # Security Features (via Merlin)
/// - **STROBE-based**: Uses Keccak-f permutation, not plain SHA-256
/// - **Automatic framing**: Length-prefixed messages prevent concatenation attacks
/// - **Domain separation**: Protocol labels are bound to the transcript
#[derive(Clone)]
pub struct Transcript {
    inner: merlin::Transcript,
}

impl std::fmt::Debug for Transcript {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Transcript").field("version", &TRANSCRIPT_VERSION).finish()
    }
}

impl Transcript {
    /// Create a new transcript with a domain separator.
    ///
    /// The domain separator is combined with the protocol version to ensure
    /// transcripts from different protocols or versions don't collide.
    pub fn new(domain: &str) -> Self {
        // Merlin requires static labels, so we use a fixed protocol label
        // and append the domain as a message
        let mut inner = merlin::Transcript::new(b"SAMAHARAM_V2");
        inner.append_message(b"domain", domain.as_bytes());
        Self { inner }
    }

    /// Append a message to the transcript.
    ///
    /// Messages are automatically length-prefixed by Merlin to prevent
    /// ambiguous encodings.
    pub fn append_message(&mut self, label: &str, message: &[u8]) {
        // Convert label to bytes and use a fixed prefix
        self.inner.append_message(b"msg", label.as_bytes());
        self.inner.append_message(b"data", message);
    }

    /// Append a G1 point to the transcript using canonical serialization.
    ///
    /// Uses `GroupEncoding::to_bytes()` for deterministic, platform-independent
    /// serialization. This is critical for Fiat-Shamir security.
    pub fn append_g1<E: PairingEngine>(&mut self, label: &str, point: &E::G1Affine) {
        use group::GroupEncoding;
        let bytes = point.to_bytes();
        self.inner.append_message(b"g1_label", label.as_bytes());
        self.inner.append_message(b"g1_point", bytes.as_ref());
    }

    /// Append a G2 point to the transcript using canonical serialization.
    ///
    /// Uses `GroupEncoding::to_bytes()` for deterministic serialization.
    pub fn append_g2<E: PairingEngine>(&mut self, label: &str, point: &E::G2Affine) {
        use group::GroupEncoding;
        let bytes = point.to_bytes();
        self.inner.append_message(b"g2_label", label.as_bytes());
        self.inner.append_message(b"g2_point", bytes.as_ref());
    }

    /// Append a scalar to the transcript using canonical representation.
    pub fn append_scalar<E: PairingEngine>(&mut self, label: &str, scalar: &E::Fr) {
        use ff::PrimeField;
        let repr = scalar.to_repr();
        self.inner.append_message(b"scalar_label", label.as_bytes());
        self.inner.append_message(b"scalar", repr.as_ref());
    }

    /// Get a challenge scalar from the transcript.
    ///
    /// Uses Merlin's `challenge_bytes` to extract pseudorandom bytes,
    /// then applies rejection sampling to get a uniform field element.
    pub fn challenge_scalar<E: PairingEngine>(&mut self, label: &str) -> E::Fr {
        use ff::PrimeField;

        // Append the challenge label to the transcript
        self.inner.append_message(b"challenge_label", label.as_bytes());

        // Get enough bytes for rejection sampling (64 bytes = 512 bits)
        let mut challenge_bytes = [0u8; 64];
        self.inner.challenge_bytes(b"challenge", &mut challenge_bytes);

        // Try to construct a valid scalar using rejection sampling
        let mut repr = <E::Fr as PrimeField>::Repr::default();
        let len = repr.as_ref().len().min(32);

        // Use different portions of the challenge bytes with a nonce for rejection sampling
        let mut nonce = 0u64;
        loop {
            // Mix the nonce into the challenge bytes
            let mut mixed = [0u8; 40];
            mixed[..32].copy_from_slice(&challenge_bytes[..32]);
            mixed[32..40].copy_from_slice(&nonce.to_le_bytes());

            // Hash to get bytes for the scalar
            use sha2::{Digest, Sha256};
            let hash = Sha256::digest(mixed);

            repr.as_mut()[..len].copy_from_slice(&hash[..len]);
            if let Some(scalar) = E::Fr::from_repr(repr).into_option() {
                return scalar;
            }
            nonce += 1;

            // Safety valve: should never happen with proper field size
            if nonce > 1000 {
                panic!("Failed to sample scalar after 1000 attempts");
            }
        }
    }

    /// Fork the transcript for parallel operations.
    ///
    /// Creates an independent transcript that shares the same history
    /// but diverges after the fork point.
    pub fn fork(&self, label: &str) -> Self {
        let mut forked = self.clone();
        forked.inner.append_message(b"fork", label.as_bytes());
        forked
    }

    /// Get access to the underlying Merlin transcript.
    ///
    /// This is useful for advanced use cases that need direct Merlin access.
    pub fn inner(&self) -> &merlin::Transcript {
        &self.inner
    }

    /// Get mutable access to the underlying Merlin transcript.
    pub fn inner_mut(&mut self) -> &mut merlin::Transcript {
        &mut self.inner
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::backend::bn254::Bn254;
    use halo2curves::bn256::Fr;

    #[test]
    fn transcript_produces_deterministic_challenges() {
        let mut t1 = Transcript::new("test");
        t1.append_message("input", b"hello");
        let c1: Fr = t1.challenge_scalar::<Bn254>("challenge");

        let mut t2 = Transcript::new("test");
        t2.append_message("input", b"hello");
        let c2: Fr = t2.challenge_scalar::<Bn254>("challenge");

        assert_eq!(c1, c2);
    }

    #[test]
    fn transcript_different_inputs_different_challenges() {
        let mut t1 = Transcript::new("test");
        t1.append_message("input", b"hello");
        let c1: Fr = t1.challenge_scalar::<Bn254>("challenge");

        let mut t2 = Transcript::new("test");
        t2.append_message("input", b"world");
        let c2: Fr = t2.challenge_scalar::<Bn254>("challenge");

        assert_ne!(c1, c2);
    }

    #[test]
    fn transcript_domain_separation() {
        let mut t1 = Transcript::new("domain1");
        t1.append_message("input", b"same");
        let c1: Fr = t1.challenge_scalar::<Bn254>("challenge");

        let mut t2 = Transcript::new("domain2");
        t2.append_message("input", b"same");
        let c2: Fr = t2.challenge_scalar::<Bn254>("challenge");

        assert_ne!(c1, c2);
    }

    #[test]
    fn transcript_fork_produces_different_challenges() {
        let mut t1 = Transcript::new("test");
        t1.append_message("shared", b"data");

        let mut forked = t1.fork("branch1");
        let c_forked: Fr = forked.challenge_scalar::<Bn254>("challenge");

        let c_original: Fr = t1.challenge_scalar::<Bn254>("challenge");

        assert_ne!(c_forked, c_original);
    }

    #[test]
    fn transcript_uses_merlin_internally() {
        let t = Transcript::new("test");
        // Verify we can access the inner Merlin transcript
        let _inner = t.inner();
    }
}
