use bellman::gadgets::boolean::Boolean;
use bellman::gadgets::sha256::sha256_block_no_padding;
use bellman::{ConstraintSystem, SynthesisError};
use pairing::Engine;

fn prf<E, CS>(
    cs: CS,
    a: bool,
    b: bool,
    c: bool,
    d: bool,
    x: &[Boolean],
    y: &[Boolean],
) -> Result<Vec<Boolean>, SynthesisError>
where
    E: Engine,
    CS: ConstraintSystem<E>,
{
    assert_eq!(x.len(), 252);
    assert_eq!(y.len(), 256);

    let mut image = vec![];
    image.push(Boolean::constant(a));
    image.push(Boolean::constant(b));
    image.push(Boolean::constant(c));
    image.push(Boolean::constant(d));
    image.extend(x.iter().cloned());
    image.extend(y.iter().cloned());

    assert_eq!(image.len(), 512);

    sha256_block_no_padding(cs, &image)
}

pub fn prf_a_pk<E, CS>(cs: CS, a_sk: &[Boolean]) -> Result<Vec<Boolean>, SynthesisError>
where
    E: Engine,
    CS: ConstraintSystem<E>,
{
    prf(
        cs,
        true,
        true,
        false,
        false,
        a_sk,
        &(0..256)
            .map(|_| Boolean::constant(false))
            .collect::<Vec<_>>(),
    )
}

pub fn prf_nf<E, CS>(
    cs: CS,
    a_sk: &[Boolean],
    rho: &[Boolean],
) -> Result<Vec<Boolean>, SynthesisError>
where
    E: Engine,
    CS: ConstraintSystem<E>,
{
    prf(cs, true, true, true, false, a_sk, rho)
}

pub fn prf_pk<E, CS>(
    cs: CS,
    a_sk: &[Boolean],
    h_sig: &[Boolean],
    nonce: bool,
) -> Result<Vec<Boolean>, SynthesisError>
where
    E: Engine,
    CS: ConstraintSystem<E>,
{
    prf(cs, false, nonce, false, false, a_sk, h_sig)
}

pub fn prf_rho<E, CS>(
    cs: CS,
    phi: &[Boolean],
    h_sig: &[Boolean],
    nonce: bool,
) -> Result<Vec<Boolean>, SynthesisError>
where
    E: Engine,
    CS: ConstraintSystem<E>,
{
    prf(cs, false, nonce, true, false, phi, h_sig)
}