#![cfg_attr(not(test), warn(unused_crate_dependencies))]

use crate::{
    compile::{Risc0Program, compile_risc0_program},
    compile_stock_rust::compile_risc0_program_stock_rust,
    error::Risc0Error,
    output::deserialize_from,
};
use borsh::{BorshDeserialize, BorshSerialize};
use risc0_zkvm::{
    DEFAULT_MAX_PO2, DefaultProver, ExecutorEnv, ExecutorEnvBuilder, ExternalProver, InnerReceipt,
    Journal, ProverOpts, Receipt, ReceiptClaim, SuccinctReceipt, default_executor, default_prover,
};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use std::{env, io::Read, ops::RangeInclusive, path::Path, rc::Rc, time::Instant};
use zkvm_interface::{
    Compiler, Input, InputItem, ProgramExecutionReport, ProgramProvingReport, Proof,
    ProverResourceType, PublicValues, zkVM, zkVMError,
};

include!(concat!(env!("OUT_DIR"), "/name_and_sdk_version.rs"));

mod compile;
mod compile_stock_rust;

mod error;
mod output;

/// Default logarithmic segment size from [`DEFAULT_SEGMENT_LIMIT_PO2`].
///
/// [`DEFAULT_SEGMENT_LIMIT_PO2`]: https://github.com/risc0/risc0/blob/v3.0.1/risc0/circuit/rv32im/src/execute/mod.rs#L39.
const DEFAULT_SEGMENT_PO2: usize = 20;

/// Supported range of logarithmic segment size.
///
/// The minimum is by [`MIN_LIFT_PO2`] to be lifted.
///
/// The maximum is by [`DEFAULT_MAX_PO2`], although the real maximum is `24`,
/// but it requires us to set the `control_ids` manually in the `ProverOpts`.
///
/// [`MIN_LIFT_PO2`]: https://github.com/risc0/risc0/blob/v3.0.1/risc0/circuit/recursion/src/control_id.rs#L19
/// [`DEFAULT_MAX_PO2`]: https://github.com/risc0/risc0/blob/v3.0.1/risc0/zkvm/src/receipt.rs#L884
const SEGMENT_PO2_RANGE: RangeInclusive<usize> = 14..=DEFAULT_MAX_PO2;

/// Default logarithmic keccak size from [`KECCAK_DEFAULT_PO2`].
///
/// [`KECCAK_DEFAULT_PO2`]: https://github.com/risc0/risc0/blob/v3.0.1/risc0/circuit/keccak/src/lib.rs#L27.
const DEFAULT_KECCAK_PO2: usize = 17;

/// Supported range of logarithmic keccak size from [`KECCAK_PO2_RANGE`].
///
/// [`KECCAK_PO2_RANGE`]: https://github.com/risc0/risc0/blob/v3.0.1/risc0/circuit/keccak/src/lib.rs#L29.
const KECCAK_PO2_RANGE: RangeInclusive<usize> = 14..=18;

#[allow(non_camel_case_types)]
pub struct RV32_IM_RISC0_ZKVM_ELF;

impl Compiler for RV32_IM_RISC0_ZKVM_ELF {
    type Error = Risc0Error;

    type Program = Risc0Program;

    fn compile(&self, guest_directory: &Path) -> Result<Self::Program, Self::Error> {
        let toolchain = env::var("ERE_GUEST_TOOLCHAIN").unwrap_or_else(|_error| "risc0".into());
        match toolchain.as_str() {
            "risc0" => Ok(compile_risc0_program(guest_directory)?),
            _ => Ok(compile_risc0_program_stock_rust(
                guest_directory,
                &toolchain,
            )?),
        }
    }
}

#[derive(Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct Risc0ProofWithPublicValues {
    /// The aggregated proof generated by the Risc0 zkVM.
    pub receipt: SuccinctReceipt<ReceiptClaim>,
    /// The public values generated by the Risc0 zkVM.
    pub journal: Journal,
}

impl From<Receipt> for Risc0ProofWithPublicValues {
    fn from(receipt: Receipt) -> Self {
        match receipt.inner {
            // We always use `ProverOpts::succinct()`.
            InnerReceipt::Succinct(inner) => Self {
                receipt: inner,
                journal: receipt.journal,
            },
            _ => unreachable!(),
        }
    }
}

impl From<Risc0ProofWithPublicValues> for Receipt {
    fn from(Risc0ProofWithPublicValues { receipt, journal }: Risc0ProofWithPublicValues) -> Self {
        Receipt::new(InnerReceipt::Succinct(receipt), journal.bytes)
    }
}

pub struct EreRisc0 {
    program: <RV32_IM_RISC0_ZKVM_ELF as Compiler>::Program,
    resource: ProverResourceType,
    segment_po2: usize,
    keccak_po2: usize,
}

impl EreRisc0 {
    pub fn new(
        program: <RV32_IM_RISC0_ZKVM_ELF as Compiler>::Program,
        resource: ProverResourceType,
    ) -> Result<Self, zkVMError> {
        if matches!(resource, ProverResourceType::Network(_)) {
            panic!(
                "Network proving not yet implemented for RISC Zero. Use CPU or GPU resource type."
            );
        }

        let [segment_po2, keccak_po2] = [
            ("RISC0_SEGMENT_PO2", DEFAULT_SEGMENT_PO2, SEGMENT_PO2_RANGE),
            ("RISC0_KECCAK_PO2", DEFAULT_KECCAK_PO2, KECCAK_PO2_RANGE),
        ]
        .map(|(key, default, range)| {
            let val = env::var(key)
                .ok()
                .and_then(|po2| po2.parse::<usize>().ok())
                .unwrap_or(default);
            if !range.contains(&val) {
                panic!("Unsupported po2 value {val} of {key}, expected in range {range:?}")
            }
            val
        });

        Ok(Self {
            program,
            resource,
            segment_po2,
            keccak_po2,
        })
    }
}

impl zkVM for EreRisc0 {
    fn execute(&self, inputs: &Input) -> Result<(PublicValues, ProgramExecutionReport), zkVMError> {
        let executor = default_executor();
        let mut env = ExecutorEnv::builder();
        serialize_inputs(&mut env, inputs).map_err(zkVMError::other)?;
        let env = env.build().map_err(zkVMError::other)?;

        let start = Instant::now();
        let session_info = executor
            .execute(env, &self.program.elf)
            .map_err(zkVMError::other)?;

        let public_values = session_info.journal.bytes.clone();

        Ok((
            public_values,
            ProgramExecutionReport {
                total_num_cycles: session_info.cycles() as u64,
                execution_duration: start.elapsed(),
                ..Default::default()
            },
        ))
    }

    fn prove(
        &self,
        inputs: &Input,
    ) -> Result<(PublicValues, Proof, ProgramProvingReport), zkVMError> {
        let prover = match self.resource {
            ProverResourceType::Cpu => Rc::new(ExternalProver::new("ipc", "r0vm")),
            ProverResourceType::Gpu => {
                if cfg!(feature = "metal") {
                    // When `metal` is enabled, we use the `LocalProver` to do
                    // proving. but it's not public so we use `default_prover`
                    // to instantiate it.
                    default_prover()
                } else {
                    // The `DefaultProver` uses `r0vm-cuda` to spawn multiple
                    // workers to do multi-gpu proving.
                    // It uses env `RISC0_DEFAULT_PROVER_NUM_GPUS` to determine
                    // how many available GPUs there are.
                    Rc::new(DefaultProver::new("r0vm-cuda").map_err(zkVMError::other)?)
                }
            }
            ProverResourceType::Network(_) => {
                panic!(
                    "Network proving not yet implemented for RISC Zero. Use CPU or GPU resource type."
                );
            }
        };

        let mut env = ExecutorEnv::builder();
        serialize_inputs(&mut env, inputs).map_err(zkVMError::other)?;
        let env = env
            .segment_limit_po2(self.segment_po2 as _)
            .keccak_max_po2(self.keccak_po2 as _)
            .map_err(zkVMError::other)?
            .build()
            .map_err(zkVMError::other)?;

        let now = std::time::Instant::now();
        let prove_info = prover
            .prove_with_opts(env, &self.program.elf, &ProverOpts::succinct())
            .map_err(zkVMError::other)?;
        let proving_time = now.elapsed();

        let public_values = prove_info.receipt.journal.bytes.clone();
        let proof = borsh::to_vec(&Risc0ProofWithPublicValues::from(prove_info.receipt))
            .map_err(zkVMError::other)?;

        Ok((
            public_values,
            proof,
            ProgramProvingReport::new(proving_time),
        ))
    }

    fn verify(&self, proof: &[u8]) -> Result<PublicValues, zkVMError> {
        let receipt: Receipt = borsh::from_slice::<Risc0ProofWithPublicValues>(proof)
            .map_err(zkVMError::other)?
            .into();

        receipt
            .verify(self.program.image_id)
            .map_err(zkVMError::other)?;

        let public_values = receipt.journal.bytes.clone();

        Ok(public_values)
    }

    fn name(&self) -> &'static str {
        NAME
    }

    fn sdk_version(&self) -> &'static str {
        SDK_VERSION
    }

    fn deserialize_from<R: Read, T: DeserializeOwned>(&self, reader: R) -> Result<T, zkVMError> {
        deserialize_from(reader)
    }
}

fn serialize_inputs(env: &mut ExecutorEnvBuilder, inputs: &Input) -> Result<(), anyhow::Error> {
    for input in inputs.iter() {
        match input {
            // Corresponding to `env.read::<T>()`.
            InputItem::Object(obj) => env.write(obj)?,
            // Corresponding to `env.read::<T>()`.
            //
            // Note that we call `write_slice` to append the bytes to the inputs
            // directly, to avoid double serailization.
            InputItem::SerializedObject(bytes) => env.write_slice(bytes),
            // Corresponding to `env.read_frame()`.
            //
            // Note that `write_frame` is different from `write_slice`, it
            // prepends the `bytes.len().to_le_bytes()`.
            InputItem::Bytes(bytes) => env.write_frame(bytes),
        };
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::OnceLock;
    use test_utils::host::{
        BasicProgramIo, run_zkvm_execute, run_zkvm_prove, testing_guest_directory,
    };

    static BASIC_PRORGAM: OnceLock<Risc0Program> = OnceLock::new();

    fn basic_program() -> Risc0Program {
        BASIC_PRORGAM
            .get_or_init(|| {
                RV32_IM_RISC0_ZKVM_ELF
                    .compile(&testing_guest_directory("risc0", "basic"))
                    .unwrap()
            })
            .clone()
    }

    #[test]
    fn test_execute() {
        let program = basic_program();
        let zkvm = EreRisc0::new(program, ProverResourceType::Cpu).unwrap();

        let io = BasicProgramIo::valid();
        run_zkvm_execute(&zkvm, &io);
    }

    #[test]
    fn test_execute_nightly() {
        let guest_directory = testing_guest_directory("risc0", "stock_nightly_no_std");
        let program =
            compile_risc0_program_stock_rust(&guest_directory, &"nightly".to_string()).unwrap();
        let zkvm = EreRisc0::new(program, ProverResourceType::Cpu).unwrap();

        let result = zkvm.execute(&BasicProgramIo::empty());
        assert!(result.is_ok(), "Risc0 execution failure");
    }

    #[test]
    fn test_execute_invalid_inputs() {
        let program = basic_program();
        let zkvm = EreRisc0::new(program, ProverResourceType::Cpu).unwrap();

        for inputs in [
            BasicProgramIo::empty(),
            BasicProgramIo::invalid_type(),
            BasicProgramIo::invalid_data(),
        ] {
            zkvm.execute(&inputs).unwrap_err();
        }
    }

    #[test]
    fn test_prove() {
        let program = basic_program();
        let zkvm = EreRisc0::new(program, ProverResourceType::Cpu).unwrap();

        let io = BasicProgramIo::valid();
        run_zkvm_prove(&zkvm, &io);
    }

    #[test]
    fn test_prove_invalid_inputs() {
        let program = basic_program();
        let zkvm = EreRisc0::new(program, ProverResourceType::Cpu).unwrap();

        for inputs in [
            BasicProgramIo::empty(),
            BasicProgramIo::invalid_type(),
            BasicProgramIo::invalid_data(),
        ] {
            zkvm.prove(&inputs).unwrap_err();
        }
    }
}
