use alloc::{boxed::Box, string::String, vec::Vec};
use core::fmt::{self, Debug, Display, Formatter};

use serde::{Deserialize, Serialize};
use sp1_hypercube::{SP1PcsProofInner, SP1RecursionProof, ShardProof};
use sp1_primitives::{io::SP1PublicValues, SP1GlobalContext};
use strum::{EnumDiscriminants, EnumTryAs, IntoDiscriminant};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ProofBn254 {
    Plonk(PlonkBn254Proof),
    Groth16(Groth16Bn254Proof),
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PlonkBn254Proof {
    pub public_inputs: [String; 5],
    pub encoded_proof: String,
    pub raw_proof: String,
    pub plonk_vkey_hash: [u8; 32],
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Groth16Bn254Proof {
    pub public_inputs: [String; 5],
    pub encoded_proof: String,
    pub raw_proof: String,
    pub groth16_vkey_hash: [u8; 32],
}

/// A proof generated by the SP1 RISC-V zkVM.
#[derive(Clone, Serialize, Deserialize, EnumDiscriminants, EnumTryAs)]
#[strum_discriminants(derive(Default, Hash, PartialOrd, Ord))]
#[strum_discriminants(name(SP1ProofMode))]
pub enum SP1Proof {
    /// A proof generated by the core proof mode.
    ///
    /// The proof size scales linearly with the number of cycles.
    #[strum_discriminants(default)]
    Core(Vec<ShardProof<SP1GlobalContext, SP1PcsProofInner>>),
    /// A proof generated by the compress proof mode.
    ///
    /// The proof size is constant, regardless of the number of cycles.
    Compressed(Box<SP1RecursionProof<SP1GlobalContext, SP1PcsProofInner>>),
    /// A proof generated by the Plonk proof mode.
    Plonk(PlonkBn254Proof),
    /// A proof generated by the Groth16 proof mode.
    Groth16(Groth16Bn254Proof),
}

impl Debug for SP1Proof {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            SP1Proof::Core(_) => write!(f, "Core"),
            SP1Proof::Compressed(_) => write!(f, "Compressed"),
            SP1Proof::Plonk(_) => write!(f, "Plonk"),
            SP1Proof::Groth16(_) => write!(f, "Groth16"),
        }
    }
}

impl Display for SP1Proof {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            SP1Proof::Core(_) => write!(f, "Core"),
            SP1Proof::Compressed(_) => write!(f, "Compressed"),
            SP1Proof::Plonk(_) => write!(f, "Plonk"),
            SP1Proof::Groth16(_) => write!(f, "Groth16"),
        }
    }
}

/// The proof generated by the prover network.
///
/// Since [`bincode`] is not self describing, it cannot handle "nullable" optional values.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProofFromNetwork {
    pub proof: SP1Proof,
    pub public_values: SP1PublicValues,
    pub sp1_version: String,
}

impl ProofFromNetwork {
    #[inline]
    pub fn mode(&self) -> SP1ProofMode {
        self.proof.discriminant()
    }
}
