use ark_crypto_primitives::sponge::{
    constraints::{AbsorbGadget, CryptographicSpongeVar},
    poseidon::constraints::PoseidonSpongeVar,
    Absorb,
};
use ark_ec::CurveGroup;
use ark_ff::PrimeField;
use ark_r1cs_std::{alloc::AllocVar, fields::fp::FpVar, ToConstraintFieldGadget};
use ark_relations::r1cs::SynthesisError;

use crate::{transcript::Transcript, Error};

use super::circuits::CF1;

pub trait CommittedInstanceOps<C: CurveGroup> {
    /// The in-circuit representation of the committed instance.
    type Var: AllocVar<Self, CF1<C>> + CommittedInstanceVarOps<C>;
    /// `hash` implements the committed instance hash compatible with the
    /// in-circuit implementation from `CommittedInstanceVarOps::hash`.
    ///
    /// Returns `H(i, z_0, z_i, U_i)`, where `i` can be `i` but also `i+1`, and
    /// `U_i` is the committed instance `self`.
    fn hash<T: Transcript<CF1<C>>>(
        &self,
        sponge: &T,
        pp_hash: CF1<C>, // public params hash
        i: CF1<C>,
        z_0: &[CF1<C>],
        z_i: &[CF1<C>],
    ) -> CF1<C>
    where
        CF1<C>: Absorb,
        Self: Sized + Absorb,
    {
        let mut sponge = sponge.clone();
        sponge.absorb(&pp_hash);
        sponge.absorb(&i);
        sponge.absorb(&z_0);
        sponge.absorb(&z_i);
        sponge.absorb(&self);
        sponge.squeeze_field_elements(1)[0]
    }

    /// Returns the commitments contained in the committed instance.
    fn get_commitments(&self) -> Vec<C>;

    /// Returns `true` if the committed instance is an incoming instance, and
    /// `false` if it is a running instance.
    fn is_incoming(&self) -> bool;

    /// Checks if the committed instance is an incoming instance.
    fn check_incoming(&self) -> Result<(), Error> {
        self.is_incoming()
            .then_some(())
            .ok_or(Error::NotIncomingCommittedInstance)
    }
}

pub trait CommittedInstanceVarOps<C: CurveGroup> {
    type PointVar: ToConstraintFieldGadget<CF1<C>>;
    /// `hash` implements the in-circuit committed instance hash compatible with
    /// the native implementation from `CommittedInstanceOps::hash`.
    /// Returns `H(i, z_0, z_i, U_i)`, where `i` can be `i` but also `i+1`, and
    /// `U_i` is the committed instance `self`.
    ///
    /// Additionally it returns the in-circuit representation of the committed
    /// instance `self` as a vector of field elements, so they can be reused in
    /// other gadgets avoiding recalculating (reconstraining) them.
    #[allow(clippy::type_complexity)]
    fn hash(
        &self,
        sponge: &PoseidonSpongeVar<CF1<C>>,
        pp_hash: &FpVar<CF1<C>>,
        i: &FpVar<CF1<C>>,
        z_0: &[FpVar<CF1<C>>],
        z_i: &[FpVar<CF1<C>>],
    ) -> Result<(FpVar<CF1<C>>, Vec<FpVar<CF1<C>>>), SynthesisError>
    where
        Self: AbsorbGadget<CF1<C>>,
    {
        let mut sponge = sponge.clone();
        let U_vec = self.to_sponge_field_elements()?;
        sponge.absorb(&pp_hash)?;
        sponge.absorb(&i)?;
        sponge.absorb(&z_0)?;
        sponge.absorb(&z_i)?;
        sponge.absorb(&U_vec)?;
        Ok((sponge.squeeze_field_elements(1)?.pop().unwrap(), U_vec))
    }

    /// Returns the commitments contained in the committed instance.
    fn get_commitments(&self) -> Vec<Self::PointVar>;

    /// Returns the public inputs contained in the committed instance.
    fn get_public_inputs(&self) -> &[FpVar<CF1<C>>];

    /// Generates constraints to enforce that the committed instance is an
    /// incoming instance.
    fn enforce_incoming(&self) -> Result<(), SynthesisError>;

    /// Generates constraints to enforce that the committed instance `self` is
    /// partially equal to another committed instance `other`.
    /// Here, only field elements are compared, while commitments (points) are
    /// not.
    fn enforce_partial_equal(&self, other: &Self) -> Result<(), SynthesisError>;
}

pub trait WitnessOps<F: PrimeField> {
    /// The in-circuit representation of the witness.
    type Var: AllocVar<Self, F> + WitnessVarOps<F>;

    /// Returns the openings (i.e., the values being committed to and the
    /// randomness) contained in the witness.
    fn get_openings(&self) -> Vec<(&[F], F)>;
}

pub trait WitnessVarOps<F: PrimeField> {
    /// Returns the openings (i.e., the values being committed to and the
    /// randomness) contained in the witness.
    fn get_openings(&self) -> Vec<(&[FpVar<F>], FpVar<F>)>;
}