/*!
 * BBS ProofGen and ProofVerify — zero-knowledge selective disclosure.
 *
 * Per IETF draft §3.4.3 (ProofGen) and §3.4.4 (ProofVerify):
 *
 * ProofGen creates a zero-knowledge proof that the prover knows a valid
 * BBS signature over a set of messages, selectively disclosing a subset.
 *
 * The proof is unlinkable: different proofs from the same signature
 * cannot be correlated by verifiers.
 */

use std::collections::HashSet;

use bls12_381_plus::{
    G1Affine, G1Projective, G2Affine, G2Prepared, G2Projective, Scalar, multi_miller_loop,
};
use ff::Field;
use group::Group;
use rand::Rng;

use crate::ciphersuite::Ciphersuite;
use crate::error::{BbsError, Result};
use crate::generators::{calculate_domain, create_generators, p1_generator, point_to_bytes};
use crate::hash::{hash_to_scalar, messages_to_scalars, scalar_to_bytes};
use crate::signature::compute_b;
use crate::types::{Proof, PublicKey, Signature};

/// Generate a zero-knowledge proof of selective disclosure.
///
/// # Arguments
///
/// * `pk` - The signer's public key
/// * `signature` - The BBS signature to prove knowledge of
/// * `header` - The same header used during signing
/// * `presentation_header` - Session-specific header (binds proof to session)
/// * `messages` - All messages that were signed
/// * `disclosed_indexes` - Which message indexes to reveal (0-based)
/// * `cs` - The ciphersuite
///
/// # Returns
///
/// A zero-knowledge proof. The verifier receives only the disclosed messages
/// and can verify the proof without learning anything about undisclosed messages.
pub fn core_proof_gen(
    pk: &PublicKey,
    signature: &Signature,
    header: &[u8],
    presentation_header: &[u8],
    messages: &[&[u8]],
    disclosed_indexes: &[usize],
    cs: Ciphersuite,
) -> Result<Proof> {
    let l = messages.len();

    // Validate disclosed indexes: bounds check and no duplicates
    {
        let mut seen = HashSet::with_capacity(disclosed_indexes.len());
        for &idx in disclosed_indexes {
            if idx >= l {
                return Err(BbsError::InvalidIndex(format!(
                    "disclosed index {idx} >= message count {l}"
                )));
            }
            if !seen.insert(idx) {
                return Err(BbsError::InvalidIndex(format!(
                    "duplicate disclosed index: {idx}"
                )));
            }
        }
    }

    // 1. Convert messages to scalars
    let msg_scalars = messages_to_scalars(messages, cs)?;

    // 2. Create generators
    let generators = create_generators(l + 1, cs)?;
    let q1 = &generators[0];
    let h_generators = &generators[1..];

    // 3. Calculate domain
    let domain = calculate_domain(pk, q1, h_generators, header, cs)?;

    // Partition into disclosed/undisclosed
    let mut undisclosed_indexes: Vec<usize> =
        (0..l).filter(|i| !disclosed_indexes.contains(i)).collect();
    undisclosed_indexes.sort();
    let u = undisclosed_indexes.len();

    // 4. Generate random scalars
    let mut rng = rand::rng();
    let r1 = random_nonzero_scalar(&mut rng);
    let r2 = random_nonzero_scalar(&mut rng);
    let e_tilde = random_nonzero_scalar(&mut rng);
    let r1_tilde = random_nonzero_scalar(&mut rng);
    let r3_tilde = random_nonzero_scalar(&mut rng);

    let mut m_tildes: Vec<Scalar> = Vec::with_capacity(u);
    for _ in 0..u {
        m_tildes.push(random_nonzero_scalar(&mut rng));
    }

    // 5. Compute B = P1 + Q1*domain + H1*msg1 + ... + HL*msgL
    let b = compute_b(q1, &domain, h_generators, &msg_scalars);

    // 6. Blinding: D = B * r2, Abar = A * (r1 * r2), Bbar = D * r1 - Abar * e
    let d = b * r2;
    let r1_r2 = r1 * r2;
    let abar = signature.a * r1_r2;
    let bbar = d * r1 - abar * signature.e;

    // r3 = r2^(-1)
    let r3 = r2.invert().unwrap();

    // 7. Compute T1 = Abar * e~ + D * r1~
    let t1 = abar * e_tilde + d * r1_tilde;

    // 8. Compute T2 = D * r3~ + sum(H_j * m~_j for undisclosed j)
    let mut t2 = d * r3_tilde;
    for (k, &j) in undisclosed_indexes.iter().enumerate() {
        t2 += h_generators[j] * m_tildes[k];
    }

    // 9. Fiat-Shamir challenge (uses only disclosed message scalars)
    let disclosed_msg_scalars: Vec<Scalar> =
        disclosed_indexes.iter().map(|&i| msg_scalars[i]).collect();

    let challenge = compute_challenge(
        &abar,
        &bbar,
        &d,
        &t1,
        &t2,
        disclosed_indexes,
        &disclosed_msg_scalars,
        &domain,
        presentation_header,
        cs,
    )?;

    // 10. Compute response scalars
    let e_hat = e_tilde + challenge * signature.e;
    let r1_hat = r1_tilde - challenge * r1;
    let r3_hat = r3_tilde - challenge * r3;

    let mut m_hats: Vec<Scalar> = Vec::with_capacity(u);
    for (k, &j) in undisclosed_indexes.iter().enumerate() {
        // Per IETF draft: m^_j = m~_j + msg_j * c (addition, not subtraction)
        m_hats.push(m_tildes[k] + challenge * msg_scalars[j]);
    }

    // 11. Serialize proof
    let mut proof_bytes = Vec::new();
    proof_bytes.extend_from_slice(&point_to_bytes(&abar));
    proof_bytes.extend_from_slice(&point_to_bytes(&bbar));
    proof_bytes.extend_from_slice(&point_to_bytes(&d));
    proof_bytes.extend_from_slice(&scalar_to_bytes(&e_hat));
    proof_bytes.extend_from_slice(&scalar_to_bytes(&r1_hat));
    proof_bytes.extend_from_slice(&scalar_to_bytes(&r3_hat));
    for m_hat in &m_hats {
        proof_bytes.extend_from_slice(&scalar_to_bytes(m_hat));
    }
    proof_bytes.extend_from_slice(&scalar_to_bytes(&challenge));

    Ok(Proof::from_bytes(&proof_bytes))
}

/// Verify a BBS zero-knowledge proof.
///
/// # Arguments
///
/// * `pk` - The signer's public key
/// * `proof` - The proof to verify
/// * `header` - The same header used during signing
/// * `presentation_header` - The same presentation header used in proof generation
/// * `disclosed_messages` - The messages that were disclosed (in order of disclosed_indexes)
/// * `disclosed_indexes` - The indexes of the disclosed messages
/// * `cs` - The ciphersuite
pub fn core_proof_verify(
    pk: &PublicKey,
    proof: &Proof,
    header: &[u8],
    presentation_header: &[u8],
    disclosed_messages: &[&[u8]],
    disclosed_indexes: &[usize],
    cs: Ciphersuite,
) -> Result<bool> {
    // Validate public key
    pk.validate()?;

    let proof_bytes = proof.to_bytes();

    // Determine U (undisclosed count) from proof length
    let r = disclosed_indexes.len();
    let scalar_len = cs.octet_scalar_length();
    let point_len = cs.octet_point_length();
    let min_len = 3 * point_len + 4 * scalar_len;

    if proof_bytes.len() < min_len {
        return Err(BbsError::InvalidProof("proof too short".into()));
    }

    let u = (proof_bytes.len() - min_len) / scalar_len;
    let l = r + u; // Total message count

    if disclosed_messages.len() != r {
        return Err(BbsError::InvalidProof(
            "disclosed message count mismatch".into(),
        ));
    }

    // 1. Deserialize proof components
    let mut offset = 0;

    let mut abar_bytes = [0u8; 48];
    abar_bytes.copy_from_slice(&proof_bytes[offset..offset + 48]);
    let abar = crate::generators::point_from_bytes(&abar_bytes)
        .ok_or_else(|| BbsError::InvalidProof("invalid Abar point".into()))?;
    offset += 48;

    let mut bbar_bytes = [0u8; 48];
    bbar_bytes.copy_from_slice(&proof_bytes[offset..offset + 48]);
    let bbar = crate::generators::point_from_bytes(&bbar_bytes)
        .ok_or_else(|| BbsError::InvalidProof("invalid Bbar point".into()))?;
    offset += 48;

    let mut d_bytes = [0u8; 48];
    d_bytes.copy_from_slice(&proof_bytes[offset..offset + 48]);
    let d = crate::generators::point_from_bytes(&d_bytes)
        .ok_or_else(|| BbsError::InvalidProof("invalid D point".into()))?;
    offset += 48;

    let e_hat = read_scalar(proof_bytes, &mut offset)?;
    let r1_hat = read_scalar(proof_bytes, &mut offset)?;
    let r3_hat = read_scalar(proof_bytes, &mut offset)?;

    let mut m_hats = Vec::with_capacity(u);
    for _ in 0..u {
        m_hats.push(read_scalar(proof_bytes, &mut offset)?);
    }

    let challenge = read_scalar(proof_bytes, &mut offset)?;

    // 2. Convert disclosed messages to scalars
    let disclosed_scalars = messages_to_scalars(disclosed_messages, cs)?;

    // 3. Create generators for L+1 messages
    let generators = create_generators(l + 1, cs)?;
    let q1 = &generators[0];
    let h_generators = &generators[1..];

    // 4. Calculate domain
    let domain = calculate_domain(pk, q1, h_generators, header, cs)?;

    // 5. Determine undisclosed indexes
    let undisclosed_indexes: Vec<usize> =
        (0..l).filter(|i| !disclosed_indexes.contains(i)).collect();

    // 6. Recompute T1 = Abar * e^ + D * r1^ + Bbar * challenge
    let t1 = abar * e_hat + d * r1_hat + bbar * challenge;

    // Per the spec, ProofVerifyInit computes T2 as:
    // Bv = P1 + Q1*domain + sum(H_i*msg_i for disclosed)
    // T2 = Bv*c + D*r3^ + sum(H_j*m^_j for undisclosed) - Bbar*c
    //
    // This is because during ProofGen:
    // T2 = D*r3~ + sum(H_j*m~_j)
    // And the relationship Bbar = D*r1 - Abar*e means:
    // T2_verify = Bv*c + D*r3^ + sum(H_j*m^_j) - Bbar*c
    let mut bv = p1_generator() + *q1 * domain;
    for (k, &i) in disclosed_indexes.iter().enumerate() {
        bv += h_generators[i] * disclosed_scalars[k];
    }

    let mut t2 = bv * challenge + d * r3_hat;
    for (k, &j) in undisclosed_indexes.iter().enumerate() {
        t2 += h_generators[j] * m_hats[k];
    }

    // 8. Recompute challenge
    let recomputed_challenge = compute_challenge(
        &abar,
        &bbar,
        &d,
        &t1,
        &t2,
        disclosed_indexes,
        &disclosed_scalars,
        &domain,
        presentation_header,
        cs,
    )?;

    // 9. Check challenge matches
    if challenge != recomputed_challenge {
        return Ok(false);
    }

    // 10. Pairing check: e(Abar, W) * e(Bbar, -BP2) == 1_GT
    let w = pk.0;
    let bp2 = G2Projective::generator();

    let abar_affine = G1Affine::from(&abar);
    let neg_bbar_affine = G1Affine::from(&(-bbar));
    let w_affine = G2Affine::from(&w);
    let bp2_affine = G2Affine::from(&bp2);

    let result = multi_miller_loop(&[
        (&abar_affine, &G2Prepared::from(w_affine)),
        (&neg_bbar_affine, &G2Prepared::from(bp2_affine)),
    ])
    .final_exponentiation();

    Ok(bool::from(result.is_identity()))
}

/// Compute the Fiat-Shamir challenge.
#[allow(clippy::too_many_arguments)]
fn compute_challenge(
    abar: &G1Projective,
    bbar: &G1Projective,
    d: &G1Projective,
    t1: &G1Projective,
    t2: &G1Projective,
    disclosed_indexes: &[usize],
    disclosed_scalars: &[Scalar],
    domain: &Scalar,
    presentation_header: &[u8],
    cs: Ciphersuite,
) -> Result<Scalar> {
    let challenge_dst = [cs.api_id().as_slice(), b"H2S_"].concat();

    let mut data = Vec::new();
    data.extend_from_slice(&point_to_bytes(abar));
    data.extend_from_slice(&point_to_bytes(bbar));
    data.extend_from_slice(&point_to_bytes(d));
    data.extend_from_slice(&point_to_bytes(t1));
    data.extend_from_slice(&point_to_bytes(t2));

    // Disclosed indexes and messages
    data.extend_from_slice(&(disclosed_indexes.len() as u64).to_be_bytes());
    for &idx in disclosed_indexes {
        data.extend_from_slice(&(idx as u64).to_be_bytes());
    }
    for scalar in disclosed_scalars {
        data.extend_from_slice(&scalar_to_bytes(scalar));
    }

    data.extend_from_slice(&scalar_to_bytes(domain));
    data.extend_from_slice(&(presentation_header.len() as u64).to_be_bytes());
    data.extend_from_slice(presentation_header);

    hash_to_scalar(&data, &challenge_dst, cs)
}

/// Generate a random nonzero scalar.
fn random_nonzero_scalar(rng: &mut impl Rng) -> Scalar {
    loop {
        let bytes: [u8; 48] = rng.random();
        let s = Scalar::from_okm(&bytes);
        if !bool::from(s.is_zero()) {
            return s;
        }
    }
}

/// Read a 32-byte scalar from a buffer at the given offset.
fn read_scalar(buf: &[u8], offset: &mut usize) -> Result<Scalar> {
    if *offset + 32 > buf.len() {
        return Err(BbsError::InvalidProof("proof too short for scalar".into()));
    }
    let mut bytes = [0u8; 32];
    bytes.copy_from_slice(&buf[*offset..*offset + 32]);
    *offset += 32;

    let s = Scalar::from_be_bytes(&bytes);
    if s.is_some().into() {
        Ok(s.unwrap())
    } else {
        Err(BbsError::Deserialization("invalid scalar in proof".into()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::SecretKey;
    use bls12_381_plus::Scalar;

    fn test_keypair() -> (SecretKey, PublicKey) {
        let sk_scalar = Scalar::from(42u64);
        let sk = SecretKey(sk_scalar);
        let pk = PublicKey(G2Projective::generator() * sk_scalar);
        (sk, pk)
    }

    #[test]
    fn proof_single_message_single_disclosed() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![b"hello"];
        let sig =
            crate::signature::core_sign(&sk, &pk, b"", &messages, Ciphersuite::Bls12381Sha256)
                .unwrap();

        let proof = core_proof_gen(
            &pk,
            &sig,
            b"",
            b"ph",
            &messages,
            &[0],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        let valid = core_proof_verify(
            &pk,
            &proof,
            b"",
            b"ph",
            &[b"hello"],
            &[0],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        assert!(valid);
    }

    #[test]
    fn proof_multiple_messages_partial_disclosure() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![b"msg0", b"msg1", b"msg2", b"msg3"];
        let sig = crate::signature::core_sign(
            &sk,
            &pk,
            b"header",
            &messages,
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        // Disclose only msg0 and msg2
        let proof = core_proof_gen(
            &pk,
            &sig,
            b"header",
            b"session-ph",
            &messages,
            &[0, 2],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        let valid = core_proof_verify(
            &pk,
            &proof,
            b"header",
            b"session-ph",
            &[b"msg0", b"msg2"],
            &[0, 2],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        assert!(valid);
    }

    #[test]
    fn proof_all_disclosed() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![b"a", b"b", b"c"];
        let sig =
            crate::signature::core_sign(&sk, &pk, b"", &messages, Ciphersuite::Bls12381Sha256)
                .unwrap();

        let proof = core_proof_gen(
            &pk,
            &sig,
            b"",
            b"ph",
            &messages,
            &[0, 1, 2],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        let valid = core_proof_verify(
            &pk,
            &proof,
            b"",
            b"ph",
            &[b"a", b"b", b"c"],
            &[0, 1, 2],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        assert!(valid);
    }

    #[test]
    fn proof_none_disclosed() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![b"secret1", b"secret2"];
        let sig =
            crate::signature::core_sign(&sk, &pk, b"", &messages, Ciphersuite::Bls12381Sha256)
                .unwrap();

        // Disclose nothing — prove existence without revealing any messages
        let proof = core_proof_gen(
            &pk,
            &sig,
            b"",
            b"ph",
            &messages,
            &[],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        let valid = core_proof_verify(
            &pk,
            &proof,
            b"",
            b"ph",
            &[],
            &[],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        assert!(valid);
    }

    #[test]
    fn proof_wrong_disclosed_message_fails() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![b"real"];
        let sig =
            crate::signature::core_sign(&sk, &pk, b"", &messages, Ciphersuite::Bls12381Sha256)
                .unwrap();

        let proof = core_proof_gen(
            &pk,
            &sig,
            b"",
            b"ph",
            &messages,
            &[0],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        let valid = core_proof_verify(
            &pk,
            &proof,
            b"",
            b"ph",
            &[b"fake"],
            &[0],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        assert!(!valid);
    }

    #[test]
    fn proof_wrong_presentation_header_fails() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![b"msg"];
        let sig =
            crate::signature::core_sign(&sk, &pk, b"", &messages, Ciphersuite::Bls12381Sha256)
                .unwrap();

        let proof = core_proof_gen(
            &pk,
            &sig,
            b"",
            b"ph1",
            &messages,
            &[0],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        let valid = core_proof_verify(
            &pk,
            &proof,
            b"",
            b"ph2",
            &[b"msg"],
            &[0],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        assert!(!valid);
    }

    #[test]
    fn proofs_are_unlinkable() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![b"msg"];
        let sig =
            crate::signature::core_sign(&sk, &pk, b"", &messages, Ciphersuite::Bls12381Sha256)
                .unwrap();

        let proof1 = core_proof_gen(
            &pk,
            &sig,
            b"",
            b"session1",
            &messages,
            &[0],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        let proof2 = core_proof_gen(
            &pk,
            &sig,
            b"",
            b"session2",
            &messages,
            &[0],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        // Proofs should be different (randomized)
        assert_ne!(proof1.to_bytes(), proof2.to_bytes());

        // Both should verify
        assert!(
            core_proof_verify(
                &pk,
                &proof1,
                b"",
                b"session1",
                &[b"msg"],
                &[0],
                Ciphersuite::Bls12381Sha256,
            )
            .unwrap()
        );
        assert!(
            core_proof_verify(
                &pk,
                &proof2,
                b"",
                b"session2",
                &[b"msg"],
                &[0],
                Ciphersuite::Bls12381Sha256,
            )
            .unwrap()
        );
    }

    #[test]
    fn proof_invalid_index_rejected() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![b"msg"];
        let sig =
            crate::signature::core_sign(&sk, &pk, b"", &messages, Ciphersuite::Bls12381Sha256)
                .unwrap();

        let result = core_proof_gen(
            &pk,
            &sig,
            b"",
            b"ph",
            &messages,
            &[5], // Index 5 out of bounds
            Ciphersuite::Bls12381Sha256,
        );

        assert!(result.is_err());
    }

    #[test]
    fn proof_ten_messages_selective() {
        let (sk, pk) = test_keypair();
        let messages: Vec<&[u8]> = vec![
            b"m0", b"m1", b"m2", b"m3", b"m4", b"m5", b"m6", b"m7", b"m8", b"m9",
        ];
        let sig =
            crate::signature::core_sign(&sk, &pk, b"hdr", &messages, Ciphersuite::Bls12381Sha256)
                .unwrap();

        // Disclose messages 2, 5, 7
        let proof = core_proof_gen(
            &pk,
            &sig,
            b"hdr",
            b"ph",
            &messages,
            &[2, 5, 7],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        let valid = core_proof_verify(
            &pk,
            &proof,
            b"hdr",
            b"ph",
            &[b"m2", b"m5", b"m7"],
            &[2, 5, 7],
            Ciphersuite::Bls12381Sha256,
        )
        .unwrap();

        assert!(valid);
    }
}
