// TODO: After finalizing vkeys, decide on a vkey verification solution. Then, clean up the
// constants and error enum in this file and update sp1_prover::tests::sp1_verifier_valid.

use alloc::boxed::Box;

use sp1_hypercube::{MachineVerifierConfigError, SP1InnerPcs, VcsError};
use sp1_primitives::{SP1Field, SP1GlobalContext};
use strum::IntoDiscriminant;
use thiserror::Error;

// NOTE: that all these constants and types are checked by sp1_prover::tests::sp1_verifier_valid.
// If you add a new proof, you MUST add to the test in that crate.

mod config;
mod internal;

use crate::{SP1Proof, SP1ProofMode};

pub use self::internal::SP1CompressedVerifier;
pub use config::{RECURSION_LOG_STACKING_HEIGHT, RECURSION_MAX_LOG_ROW_COUNT};

/// A reason why the verifier rejects a given proof.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum CompressedError {
    #[error("invalid proof type")]
    InvalidProofType(SP1ProofMode),
    #[error("failed to deserialize proof: {0}")]
    DeserializeProof(Box<bincode::ErrorKind>),
    #[error("failed to deserialize vkey hash: {0}")]
    DeserializeVkeyHash(Box<bincode::ErrorKind>),
    #[error("failed to verify proof: {0}")]
    ProofRejected(#[from] MachineVerifierConfigError<SP1GlobalContext, SP1InnerPcs>),
    #[error("given public values do not match the commitment in the proof")]
    PublicValuesMismatch,
    #[error("incorrect vkey")]
    InvalidVkey(VcsError),
}

/// A verifier for SP1 "compressed" proofs given as raw bytes.
///
/// This struct priovides a convenient interface for verifying compressed proofs given as raw bytes.
/// If possible, it's usually better to use the [SP1CompressedVerifier] struct instead.
#[derive(Debug)]
pub struct SP1CompressedVerifierRaw;

impl SP1CompressedVerifierRaw {
    /// Attempts to verify an SP1 "compressed" proof, as generated by the SP1 SDK.
    /// Returns `Ok` if the proof verifies or `Err` with the reason that verification failed.
    ///
    /// Due to technical limitations, this verifier rejects single-shard proofs, which roughly means
    /// proofs that come from very short programs. In this case, the verifier returns
    /// `Err(CompressedError::SingleShard)`.
    ///
    /// # Arguments
    ///
    /// * `proof` - The proof bytes.
    /// * `sp1_vkey_hash` - The SP1 vkey hash.
    ///
    /// The arguments may be generated in the following manner:
    ///
    /// ```ignore
    /// use sp1_sdk::{HashableKey, ProverClient, SP1Proof};
    /// let client = ProverClient::builder().cpu().build();
    /// let (pk, vk) = client.setup(ELF);
    /// let sp1_vkey_hash = bincode::serialize(&vk.hash_babybear()).unwrap();
    /// let proof = match client.prove(&pk, &stdin).compressed().run().unwrap().proof {
    ///     SP1Proof::Compressed(proof) => bincode::serialize(&proof),
    ///     _ => unreachable!("expected compressed proof"),
    /// };
    /// ```
    pub fn verify_with_public_values(
        proof: &[u8],
        sp1_public_inputs: &[u8],
        sp1_vkey_hash: &[u8],
    ) -> Result<(), CompressedError> {
        let recursion_proof: SP1Proof =
            bincode::deserialize(proof).map_err(CompressedError::DeserializeProof)?;
        let vkey_hash: [SP1Field; 8] =
            bincode::deserialize(sp1_vkey_hash).map_err(CompressedError::DeserializeVkeyHash)?;

        let verifier = SP1CompressedVerifier::new();

        if let SP1Proof::Compressed(proof) = recursion_proof {
            verifier.verify_compressed_with_public_values(&proof, sp1_public_inputs, &vkey_hash)?;
        } else {
            return Err(CompressedError::InvalidProofType(recursion_proof.discriminant()));
        }

        Ok(())
    }

    /// Attempts to verify an SP1 "compressed" proof, as generated by the SP1 SDK.
    /// Returns `Ok` if the proof verifies or `Err` with the reason that verification failed.
    ///
    /// Due to technical limitations, this verifier rejects single-shard proofs, which roughly means
    /// proofs that come from very short programs. In this case, the verifier returns
    /// `Err(CompressedError::SingleShard)`.
    ///
    /// # Arguments
    ///
    /// * `proof` - The proof bytes.
    /// * `sp1_vkey_hash` - The SP1 vkey hash.
    ///
    /// The arguments may be generated in the following manner:
    ///
    /// ```ignore
    /// use sp1_sdk::{HashableKey, ProverClient, SP1Proof};
    /// let client = ProverClient::builder().cpu().build();
    /// let (pk, vk) = client.setup(ELF);
    /// let sp1_vkey_hash = bincode::serialize(&vk.hash_babybear()).unwrap();
    /// let proof = match client.prove(&pk, &stdin).compressed().run().unwrap().proof {
    ///     SP1Proof::Compressed(proof) => bincode::serialize(&proof),
    ///     _ => unreachable!("expected compressed proof"),
    /// };
    /// ```
    pub fn verify(proof: &[u8], sp1_vkey_hash: &[u8]) -> Result<(), CompressedError> {
        let recursion_proof: SP1Proof =
            bincode::deserialize(proof).map_err(CompressedError::DeserializeProof)?;
        let vkey_hash: [SP1Field; 8] =
            bincode::deserialize(sp1_vkey_hash).map_err(CompressedError::DeserializeVkeyHash)?;

        let verifier = SP1CompressedVerifier::new();

        if let SP1Proof::Compressed(proof) = recursion_proof {
            verifier.verify_compressed(&proof, &vkey_hash)?;
        } else {
            return Err(CompressedError::InvalidProofType(recursion_proof.discriminant()));
        }

        Ok(())
    }
}
