use crate::{common::*, random_oracle::*};
use itertools::izip;

/// The common data known to the prover and verifier, i.e., public values and
/// coefficients.
pub trait SigmaProtocol: Sized {
    /// The first message of the prover.
    type CommitMessage: Serial;
    /// The second message in the sigma protocol sent by the verifier to the
    /// prover.
    type ProtocolChallenge;
    /// The third message in the sigma protocol sent by the prover in response
    /// to the challenge.
    type Response: Serialize;
    /// Prover's state after the first message. Essentially any randomness
    /// they've chosen to generate the commit message.
    type ProverState;
    /// The prover's secret data. This corresponds to what is usually called
    /// "witness".
    type SecretData;

    /// Given a random oracle, feed it all of the public input of this instance.
    fn public(&self, ro: &mut RandomOracle);

    /// Compute the first message of the prover. We allow this function
    /// to return 'None' if the inputs are inconsistent.
    /// The arguments are
    /// - self -- the prover's public and secret data
    /// - secret -- secret data known to the prover only
    /// - csprng -- a cryptographically secure random number generator
    fn compute_commit_message<R: rand::Rng>(
        &self,
        csprng: &mut R,
    ) -> Option<(Self::CommitMessage, Self::ProverState)>;

    /// Convert a random_oracle Challenge to a ProtocolChallenge
    fn get_challenge(&self, challenge: &Challenge) -> Self::ProtocolChallenge;

    /// Third message generated by the prover. We allow this function to return
    /// 'None' if the inputs are inconsistent.
    /// - self -- the prover's public and secret data
    /// - secret -- secret data known to the prover only
    /// - state -- the state returned from the 'compute_commit_message' call
    /// - challenge -- computed challenge
    /// This function is pure and deterministic.
    fn compute_response(
        &self,
        secret: Self::SecretData,
        state: Self::ProverState,
        challenge: &Self::ProtocolChallenge,
    ) -> Option<Self::Response>;

    /// Combine the public outputs from the third message with the challenge
    /// and produce a CommitMessage for verification.
    /// - challenge -- computed challenge
    /// - self -- the verifier's public data
    /// This function is pure and deterministic.
    /// It is allowed to return 'None' if some of the input data is malformed,
    /// e.g., vectors of inconsistent lengths.
    fn extract_commit_message(
        &self,
        challenge: &Self::ProtocolChallenge,
        response: &Self::Response,
    ) -> Option<Self::CommitMessage>;

    #[cfg(test)]
    /// Function used for testing. Generates valid input for this sigma proof.
    /// The 'data_size' parameter can be used to adjust the size of the
    /// generated data.
    fn with_valid_data<R: rand::Rng>(
        data_size: usize,
        csprng: &mut R,
        f: impl FnOnce(Self, Self::SecretData, &mut R),
    );
}

#[derive(Debug, Serialize, Eq, PartialEq, SerdeBase16Serialize, Clone)]
/// Generic structure to contain a single sigma proof.
pub struct SigmaProof<R: Serialize> {
    pub challenge: Challenge,
    pub response:  R,
}

#[derive(Serialize)]
pub struct AndResponse<R1: Serialize, R2: Serialize> {
    pub r1: R1,
    pub r2: R2,
}

/// An adapter to combine multiple provers or multiple verifiers.
/// The marker type C is for convenience in use with the
/// SigmaProtocolProver/Verifier traits below.
pub struct AndAdapter<P1, P2> {
    pub first:  P1,
    pub second: P2,
}

impl<P1: SigmaProtocol, P2: SigmaProtocol> SigmaProtocol for AndAdapter<P1, P2> {
    type CommitMessage = (P1::CommitMessage, P2::CommitMessage);
    // FIXME: The challenge pair is inefficient, since we will typically
    // have a pair of the same element, however in the grand scheme of things
    // this is a minor inefficiency.
    type ProtocolChallenge = (P1::ProtocolChallenge, P2::ProtocolChallenge);
    type ProverState = (P1::ProverState, P2::ProverState);
    type Response = AndResponse<P1::Response, P2::Response>;
    type SecretData = (P1::SecretData, P2::SecretData);

    fn public(&self, ro: &mut RandomOracle) {
        self.first.public(ro);
        self.second.public(ro)
    }

    fn compute_commit_message<R: rand::Rng>(
        &self,
        csprng: &mut R,
    ) -> Option<(Self::CommitMessage, Self::ProverState)> {
        let (m1, s1) = self.first.compute_commit_message(csprng)?;
        let (m2, s2) = self.second.compute_commit_message(csprng)?;
        Some(((m1, m2), (s1, s2)))
    }

    fn get_challenge(&self, challenge: &Challenge) -> Self::ProtocolChallenge {
        (
            self.first.get_challenge(challenge),
            self.second.get_challenge(challenge),
        )
    }

    fn compute_response(
        &self,
        secret: Self::SecretData,
        state: Self::ProverState,
        challenge: &Self::ProtocolChallenge,
    ) -> Option<Self::Response> {
        let r1 = self
            .first
            .compute_response(secret.0, state.0, &challenge.0)?;
        let r2 = self
            .second
            .compute_response(secret.1, state.1, &challenge.1)?;
        Some(AndResponse { r1, r2 })
    }

    fn extract_commit_message(
        &self,
        challenge: &Self::ProtocolChallenge,
        response: &Self::Response,
    ) -> Option<Self::CommitMessage> {
        let p1 = self
            .first
            .extract_commit_message(&challenge.0, &response.r1)?;
        let p2 = self
            .second
            .extract_commit_message(&challenge.1, &response.r2)?;
        Some((p1, p2))
    }

    #[cfg(test)]
    fn with_valid_data<R: rand::Rng>(
        data_size: usize,
        csprng: &mut R,
        f: impl FnOnce(Self, Self::SecretData, &mut R),
    ) {
        P1::with_valid_data(data_size, csprng, |first, s1, csprng| {
            P2::with_valid_data(data_size, csprng, |second, s2, csprng| {
                f(AndAdapter { first, second }, (s1, s2), csprng)
            })
        })
    }
}

impl<P1: SigmaProtocol, P2: SigmaProtocol> AndAdapter<P1, P2> {
    /// Extend the current adapter with a new prover.
    pub fn add_prover<P3: SigmaProtocol>(
        self,
        additional_prover: P3,
    ) -> AndAdapter<AndAdapter<P1, P2>, P3> {
        AndAdapter {
            first:  self,
            second: additional_prover,
        }
    }
}

/// ## This section provides an and-like adapter, but where we combine
/// multiple proofs of the same kind, only with different parameters.
#[derive(Serialize)]
pub struct ReplicateResponse<R: Serialize> {
    #[size_length = 4]
    pub responses: Vec<R>,
}

/// An adapter to combine multiple provers or multiple verifiers.
/// The marker type C is for convenience in use with the
/// SigmaProtocolProver/Verifier traits below.
pub struct ReplicateAdapter<P> {
    // Protocols to combine. This vector is assumed to be non-empty.
    pub protocols: Vec<P>,
}

#[derive(Serial)]
pub struct ReplicatePoints<P: Serial> {
    #[size_length = 4]
    pub points: Vec<P>,
}

impl<P: SigmaProtocol> SigmaProtocol for ReplicateAdapter<P> {
    type CommitMessage = ReplicatePoints<P::CommitMessage>;
    // We assume that all protocols use the same challenge map.
    type ProtocolChallenge = P::ProtocolChallenge;
    type ProverState = Vec<P::ProverState>;
    type Response = ReplicateResponse<P::Response>;
    type SecretData = Vec<P::SecretData>;

    fn public(&self, ro: &mut RandomOracle) {
        // add all public data in sequence from left to right
        self.protocols.iter().for_each(|p| p.public(ro))
    }

    fn compute_commit_message<R: rand::Rng>(
        &self,
        csprng: &mut R,
    ) -> Option<(Self::CommitMessage, Self::ProverState)> {
        let n = self.protocols.len();
        let mut ms = Vec::with_capacity(n);
        let mut ss = Vec::with_capacity(n);
        for p in self.protocols.iter() {
            let (m, s) = p.compute_commit_message(csprng)?;
            ms.push(m);
            ss.push(s);
        }
        Some((ReplicatePoints { points: ms }, ss))
    }

    fn get_challenge(&self, challenge: &Challenge) -> Self::ProtocolChallenge {
        // FIXME: This relies on the replicate adapter precondition.
        self.protocols
            .first()
            .expect("Protocols is non-empty.")
            .get_challenge(challenge)
    }

    fn compute_response(
        &self,
        secret: Self::SecretData,
        state: Self::ProverState,
        challenge: &Self::ProtocolChallenge,
    ) -> Option<Self::Response> {
        let n = self.protocols.len();
        if secret.len() != n {
            return None;
        }
        if state.len() != n {
            return None;
        }
        let mut rs: Vec<<P as SigmaProtocol>::Response> = Vec::with_capacity(n);
        for (p, sec, state) in izip!(self.protocols.iter(), secret, state) {
            rs.push(p.compute_response(sec, state, challenge)?);
        }
        Some(ReplicateResponse { responses: rs })
    }

    fn extract_commit_message(
        &self,
        challenge: &Self::ProtocolChallenge,
        response: &Self::Response,
    ) -> Option<Self::CommitMessage> {
        let n = self.protocols.len();
        if response.responses.len() != n {
            return None;
        }
        let mut points = Vec::with_capacity(n);
        for (p, res) in izip!(self.protocols.iter(), response.responses.iter()) {
            points.push(p.extract_commit_message(challenge, res)?);
        }
        Some(ReplicatePoints { points })
    }

    #[cfg(test)]
    fn with_valid_data<R: rand::Rng>(
        _data_size: usize,
        _csprng: &mut R,
        _f: impl FnOnce(Self, Self::SecretData, &mut R),
    ) {
        todo!()
    }
}

impl<P: SigmaProtocol> ReplicateAdapter<P> {
    /// Extend the current adapter with a new prover.
    pub fn add_prover(&mut self, additional_protocol: P) {
        self.protocols.push(additional_protocol)
    }
}

/// Given a sigma protocol prover and a context (in the form of the random
/// oracle), produce a sigma proof and update the context. This function can
/// return 'None' if the input data is inconsistent.
pub fn prove<R: rand::Rng, D: SigmaProtocol>(
    ro: &mut RandomOracle,
    prover: &D,
    secret: D::SecretData,
    csprng: &mut R,
) -> Option<SigmaProof<D::Response>> {
    let (commit_message, state) = prover.compute_commit_message(csprng)?;
    prover.public(ro);
    // For legacy reasons the label `point` is used when adding the commit message
    // to the RO
    ro.append_message("point", &commit_message);
    let challenge_bytes = ro.split().get_challenge();
    let challenge = prover.get_challenge(&challenge_bytes);
    let response = prover.compute_response(secret, state, &challenge)?;
    Some(SigmaProof {
        challenge: challenge_bytes,
        response,
    })
}

/// Given a single sigma proof and a context in the form of an instantiated
/// random oracle, verify the sigma proof and update the state of the context.
pub fn verify<D: SigmaProtocol>(
    ro: &mut RandomOracle,
    verifier: &D,
    proof: &SigmaProof<D::Response>,
) -> bool {
    let challenge = verifier.get_challenge(&proof.challenge);
    match verifier.extract_commit_message(&challenge, &proof.response) {
        None => false,
        Some(ref point) => {
            verifier.public(ro);
            ro.append_message("point", &point);
            let computed_challenge = ro.split().get_challenge();
            computed_challenge == proof.challenge
        }
    }
}

#[cfg(test)]
/// Generate a random challenge prefix of random length used for testing.
pub fn generate_challenge_prefix<R: rand::Rng>(csprng: &mut R) -> Vec<u8> {
    // length of the challenge
    let l = csprng.gen_range(0, 1000);
    let mut challenge_prefix = vec![0; l];
    for v in challenge_prefix.iter_mut() {
        *v = csprng.gen();
    }
    challenge_prefix
}