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

use crate::{
    compile::compile_zisk_program,
    error::{CommonError, ExecuteError, ProveError, VerifyError, ZiskError},
};
use blake3::Hash;
use dashmap::{DashMap, Entry};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use std::{
    fs,
    io::{self, BufRead, Read, Write},
    os::unix::fs::symlink,
    path::{Path, PathBuf},
    process::{Command, Stdio},
    sync::{LazyLock, Mutex},
    time,
};
use tempfile::{TempDir, tempdir};
use tracing::info;
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 error;

/// Lock for the command `cargo-zisk check-setup` to avoid multiple runs.
static SETUP_LOCK: Mutex<bool> = Mutex::new(false);

/// Mapping from ELF hash to ROM digest. It uses `blake3` for ELF hash as ZisK.
static ROM_DIGEST_MAP: LazyLock<DashMap<Hash, RomDigest>> = LazyLock::new(DashMap::new);

/// It panics if `EreZisk::prove` is called concurrently, so we need a lock here
/// to avoid that.
static PROVE_LOCK: Mutex<()> = Mutex::new(());

/// Merkle root of ROM trace generated by `cargo-zisk rom-setup`.
pub type RomDigest = [u64; 4];

#[allow(non_camel_case_types)]
pub struct RV64_IMA_ZISK_ZKVM_ELF;

impl Compiler for RV64_IMA_ZISK_ZKVM_ELF {
    type Error = ZiskError;

    type Program = Vec<u8>;

    fn compile(&self, guest_directory: &Path) -> Result<Self::Program, Self::Error> {
        compile_zisk_program(guest_directory).map_err(ZiskError::Compile)
    }
}

#[derive(Serialize, Deserialize)]
pub struct ZiskProofWithPublicValues {
    /// The raw aggregated proof generated by the ZisK zkVM.
    pub proof: Vec<u8>,
    /// The public values generated by the ZisK zkVM.
    pub public_values: Vec<u8>,
}

pub struct EreZisk {
    elf: Vec<u8>,
    resource: ProverResourceType,
}

impl EreZisk {
    pub fn new(elf: Vec<u8>, resource: ProverResourceType) -> Self {
        Self { elf, resource }
    }
}

impl zkVM for EreZisk {
    fn execute(&self, inputs: &Input) -> Result<(PublicValues, ProgramExecutionReport), zkVMError> {
        // Write ELF and serialized input to file.

        let input_bytes = serialize_inputs(inputs)
            .map_err(|err| ExecuteError::SerializeInput(err.into()))
            .map_err(ZiskError::Execute)?;

        let mut tempdir =
            ZiskTempDir::new().map_err(|e| ZiskError::Execute(ExecuteError::TempDir(e)))?;
        tempdir
            .write_elf(&self.elf)
            .map_err(|e| ZiskError::Execute(ExecuteError::TempDir(e)))?;
        tempdir
            .write_input(&input_bytes)
            .map_err(|e| ZiskError::Execute(ExecuteError::TempDir(e)))?;

        // Execute.

        let start = time::Instant::now();
        let output = Command::new("ziskemu")
            .arg("--elf")
            .arg(tempdir.elf_path())
            .arg("--inputs")
            .arg(tempdir.input_path())
            .arg("--output")
            .arg(tempdir.execution_output_path())
            .arg("--stats") // NOTE: enable stats in order to get total steps.
            .stderr(Stdio::inherit())
            .output()
            .map_err(|e| ZiskError::Execute(ExecuteError::Ziskemu { source: e }))?;

        if !output.status.success() {
            Err(ZiskError::Execute(ExecuteError::ZiskemuFailed {
                status: output.status,
            }))?
        }
        let execution_duration = start.elapsed();

        // Extract cycle count from the stdout.

        let total_num_cycles = String::from_utf8_lossy(&output.stdout)
            .split_once("total steps = ")
            .and_then(|(_, stats)| {
                stats
                    .split_whitespace()
                    .next()
                    .and_then(|steps| steps.parse::<u64>().ok())
            })
            .ok_or(ZiskError::Execute(ExecuteError::TotalStepsNotFound))?;

        let public_values = tempdir
            .read_execution_output()
            .map_err(|e| ZiskError::Execute(ExecuteError::TempDir(e)))?;

        Ok((
            public_values,
            ProgramExecutionReport {
                total_num_cycles,
                execution_duration,
                ..Default::default()
            },
        ))
    }

    fn prove(
        &self,
        inputs: &Input,
    ) -> Result<(PublicValues, Proof, ProgramProvingReport), zkVMError> {
        // Make sure proving key setup is done.
        check_setup()?;

        // Run ELF specific setup
        let preprocessed_rom_digest = rom_setup(&self.elf)?;

        // Obtain the prove lock to make sure proving can't be called concurrently.
        let _guard = PROVE_LOCK
            .lock()
            .map_err(|_| ProveError::ProveLockPoisoned)
            .map_err(ZiskError::Prove)?;

        // Write ELF and serialized input to file.

        let input_bytes = serialize_inputs(inputs)
            .map_err(|err| ProveError::SerializeInput(err.into()))
            .map_err(ZiskError::Prove)?;

        let mut tempdir =
            ZiskTempDir::new().map_err(|e| ZiskError::Prove(ProveError::TempDir(e)))?;
        tempdir
            .write_elf(&self.elf)
            .map_err(|e| ZiskError::Prove(ProveError::TempDir(e)))?;
        tempdir
            .write_input(&input_bytes)
            .map_err(|e| ZiskError::Prove(ProveError::TempDir(e)))?;

        // Prove.

        // TODO: Use `mpirun --np {num_processes} cargo-zisk prove ...` to
        //       utilize multiple CPUs, probably need the `ProverResourceType`
        //       to specify the number of available CPUs.

        let start = time::Instant::now();
        match self.resource {
            ProverResourceType::Cpu => {
                let status = Command::new("cargo-zisk")
                    .arg("prove")
                    .arg("--elf")
                    .arg(tempdir.elf_path())
                    .arg("--input")
                    .arg(tempdir.input_path())
                    .arg("--output-dir")
                    .arg(tempdir.output_dir_path())
                    .args([
                        "--aggregation",
                        "--verify-proofs",
                        "--save-proofs",
                        // Uncomment this if locked memory is not enough.
                        // "--unlock-mapped-memory",
                    ])
                    .status()
                    .map_err(|e| ZiskError::Prove(ProveError::CargoZiskProve { source: e }))?;

                if !status.success() {
                    Err(ZiskError::Prove(ProveError::CargoZiskProveFailed {
                        status,
                    }))?
                }
            }
            ProverResourceType::Gpu => {
                // TODO: Set env `CUDA_VISIBLE_DEVICES = {0..num_devices}` to
                //       control how many GPUs to use, probably need the `ProverResourceType`
                //       to specify the number of available GPUs.
                let witness_lib_path = dot_zisk_dir_path()
                    .join("bin")
                    .join("libzisk_witness_gpu.so");
                let status = Command::new("cargo-zisk-gpu")
                    .arg("prove")
                    .arg("--witness-lib")
                    .arg(witness_lib_path)
                    .arg("--elf")
                    .arg(tempdir.elf_path())
                    .arg("--input")
                    .arg(tempdir.input_path())
                    .arg("--output-dir")
                    .arg(tempdir.output_dir_path())
                    .args([
                        "--aggregation",
                        "--verify-proofs",
                        "--save-proofs",
                        // Comment out this if GPU RAM is not enough.
                        "--preallocate",
                        // Uncomment this if locked memory is not enough.
                        // "--unlock-mapped-memory",
                    ])
                    .status()
                    .map_err(|e| ZiskError::Prove(ProveError::CargoZiskProve { source: e }))?;

                if !status.success() {
                    Err(ZiskError::Prove(ProveError::CargoZiskProveFailed {
                        status,
                    }))?
                }
            }
            ProverResourceType::Network(_) => {
                panic!(
                    "Network proving not yet implemented for ZisK. Use CPU or GPU resource type."
                );
            }
        }
        let proving_time = start.elapsed();

        // Read proof and public values.

        let public_values = tempdir
            .read_public_values()
            .map_err(|e| ZiskError::Prove(ProveError::TempDir(e)))?;

        let proof_with_public_values = ZiskProofWithPublicValues {
            proof: tempdir
                .read_proof()
                .map_err(|e| ZiskError::Prove(ProveError::TempDir(e)))?,
            public_values,
        };
        let bytes = bincode::serialize(&proof_with_public_values)
            .map_err(|err| ZiskError::Prove(ProveError::Bincode(err)))?;

        // Deserialize public values.
        let (proved_rom_digest, public_values) =
            deserialize_public_values(&proof_with_public_values.public_values)?;

        // The proved ROM digest should be equal to preprocessed one.
        if proved_rom_digest != preprocessed_rom_digest {
            return Err(CommonError::UnexpectedRomDigest {
                preprocessed: preprocessed_rom_digest,
                proved: proved_rom_digest,
            })?;
        }

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

    fn verify(&self, bytes: &[u8]) -> Result<PublicValues, zkVMError> {
        // Run ELF specific setup
        let preprocessed_rom_digest = rom_setup(&self.elf)?;

        // Write proof and public values to file.

        let proof_with_public_values: ZiskProofWithPublicValues = bincode::deserialize(bytes)
            .map_err(|err| ZiskError::Verify(VerifyError::Bincode(err)))?;

        let mut tempdir =
            ZiskTempDir::new().map_err(|e| ZiskError::Verify(VerifyError::TempDir(e)))?;
        tempdir
            .write_proof(&proof_with_public_values.proof)
            .map_err(|e| ZiskError::Verify(VerifyError::TempDir(e)))?;
        tempdir
            .write_public_values(&proof_with_public_values.public_values)
            .map_err(|e| ZiskError::Verify(VerifyError::TempDir(e)))?;

        // Verify.

        let output = Command::new("cargo-zisk")
            .arg("verify")
            .arg("--proof")
            .arg(tempdir.proof_path())
            .arg("--public-inputs")
            .arg(tempdir.public_values_path())
            .output()
            .map_err(|e| ZiskError::Verify(VerifyError::CargoZiskVerify { source: e }))?;

        if !output.status.success() {
            Err(ZiskError::Verify(VerifyError::InvalidProof(
                String::from_utf8_lossy(&output.stderr).to_string(),
            )))?
        }

        // Deserialize public values.
        let (proved_rom_digest, public_values) =
            deserialize_public_values(&proof_with_public_values.public_values)?;

        // The proved ROM digest should be equal to preprocessed one.
        if proved_rom_digest != preprocessed_rom_digest {
            return Err(CommonError::UnexpectedRomDigest {
                preprocessed: preprocessed_rom_digest,
                proved: proved_rom_digest,
            })?;
        }

        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, _: R) -> Result<T, zkVMError> {
        unimplemented!("no native serialization in this platform")
    }
}

/// Serialize `Input` into sequence of bytes.
///
/// Because ZisK doesn't provide stdin API so we need to handle multiple inputs,
/// the current approach naively serializes each `InputItem` individually, then
/// concat them into single `Vec<u8>`.
fn serialize_inputs(inputs: &Input) -> Result<Vec<u8>, bincode::Error> {
    inputs.iter().try_fold(Vec::new(), |mut acc, item| {
        match item {
            InputItem::Object(obj) => bincode::serialize_into(&mut acc, &**obj)?,
            InputItem::SerializedObject(bytes) => acc.extend(bytes),
            InputItem::Bytes(bytes) => bincode::serialize_into(&mut acc, bytes)?,
        };
        Ok(acc)
    })
}

/// Deserialize public values as json string sequence, and parse the `RomDigest`
/// and user set public values as `Vec<u8>`.
fn deserialize_public_values(public_values: &[u8]) -> Result<(RomDigest, Vec<u8>), CommonError> {
    let public_values = serde_json::from_slice::<Vec<String>>(public_values)
        .map_err(CommonError::DeserializePublicValues)?
        .into_iter()
        .map(|v| v.parse())
        .collect::<Result<Vec<u64>, _>>()
        .map_err(CommonError::ParsePublicValue)?;

    // The public values contain at least the `RomDigest` and the number of user
    // set public values.
    if public_values.len() < 5 {
        return Err(CommonError::InvalidPublicValuesLength(public_values.len()));
    }

    // The first 4 elements of public values should be ROM digest.
    let rom_digest = public_values[..4].try_into().unwrap();

    // The 5-th element should be the number of user set public values.
    let num_user_public_values = public_values[4] as usize;

    // The rest elements should be user set public values and should be `u32`.
    let public_values = public_values[5..]
        .iter()
        .map(|v| Some(u32::try_from(*v).ok()?.to_le_bytes()))
        .take(num_user_public_values)
        .collect::<Option<Vec<_>>>()
        .ok_or(CommonError::InvalidPublicValue)?
        .into_iter()
        .flatten()
        .collect();

    Ok((rom_digest, public_values))
}

fn dot_zisk_dir_path() -> PathBuf {
    PathBuf::from(std::env::var("HOME").expect("env `$HOME` should be set")).join(".zisk")
}

/// Run ELF independent proving key setup.
fn check_setup() -> Result<(), zkVMError> {
    let mut setup = SETUP_LOCK
        .lock()
        .map_err(|_| zkVMError::Other("Setup lock is poisoned".into()))?;

    if !*setup {
        info!("Running command `cargo-zisk check-setup --aggregation`...");

        let status = Command::new("cargo-zisk")
            .args(["check-setup", "--aggregation"])
            .status()
            .map_err(|e| CommonError::CargoZiskCheckSetup { source: e })?;

        if !status.success() {
            Err(CommonError::CargoZiskCheckSetupFailed { status })?
        }

        info!("Command `cargo-zisk check-setup --aggregation` succeeded");

        *setup = true;
    }

    Ok(())
}

/// Run ELF specific setup and returns digest of ROM.
fn rom_setup(elf: &[u8]) -> Result<RomDigest, zkVMError> {
    let mut tempdir = ZiskTempDir::new().map_err(CommonError::TempDir)?;
    tempdir.create_zisk_dir().map_err(CommonError::TempDir)?;
    tempdir.write_elf(elf).map_err(CommonError::TempDir)?;

    let rom_digest = match ROM_DIGEST_MAP.entry(blake3::hash(elf)) {
        Entry::Occupied(entry) => *entry.get(),
        Entry::Vacant(entry) => {
            info!("Running command `cargo-zisk rom-setup` ...");

            let output = Command::new("cargo-zisk")
                .arg("rom-setup")
                .arg("--elf")
                .arg(tempdir.elf_path())
                .arg("--zisk-path")
                .arg(tempdir.zisk_dir_path())
                .output()
                .map_err(|e| CommonError::CargoZiskRomSetup { source: e })?;

            if !output.status.success() {
                Err(CommonError::CargoZiskRomSetupFailed {
                    status: output.status,
                })?;
            }

            let rom_digest = output
                .stdout
                .lines()
                .find_map(|line| {
                    let line = line.ok()?;
                    let line = line.split_once("Root hash: [")?.1;
                    let line = line.strip_suffix("]")?;
                    line.split(", ")
                        .filter_map(|word| word.parse::<u64>().ok())
                        .collect::<Vec<_>>()
                        .try_into()
                        .ok()
                })
                .ok_or(CommonError::RomDigestNotFound)?;

            info!("Command `cargo-zisk rom-setup` succeeded");

            *entry.insert(rom_digest)
        }
    };

    Ok(rom_digest)
}

struct ZiskTempDir {
    tempdir: TempDir,
}

impl ZiskTempDir {
    /// Create temporary directories for:
    /// - `guest.elf` - ELF compiled from guest program.
    /// - `input.bin` - Input of execution or proving.
    /// - `output/vadcop_final_proof.json` - Aggregated proof generated by proving.
    /// - `output/publics.json` - Public values generated by proving.
    fn new() -> io::Result<Self> {
        let tempdir = Self {
            tempdir: tempdir()?,
        };

        fs::create_dir(tempdir.output_dir_path())?;

        Ok(tempdir)
    }

    /// Create temporary directory `zisk` for `rom-setup`.
    fn create_zisk_dir(&mut self) -> io::Result<()> {
        fs::create_dir_all(self.zisk_dir_path())?;

        // Check the global zisk directory exists.
        let global_zisk_dir_path = dot_zisk_dir_path().join("zisk");
        if !global_zisk_dir_path.exists() {
            Err(io::Error::new(
                io::ErrorKind::NotFound,
                format!(
                    "Global .zisk/zisk directory not found at: {}",
                    global_zisk_dir_path.display()
                ),
            ))?
        }

        // Symlink necessary files for `make` command of `cargo-zisk rom-setup`.
        // The `Makefile` can be found https://github.com/0xPolygonHermez/zisk/blob/v0.10.0/emulator-asm/Makefile.
        symlink(
            dot_zisk_dir_path().join("bin"),
            self.dot_zisk_dir_path().join("bin"),
        )?;
        let temp_zisk_dir_path = self.zisk_dir_path();
        fs::create_dir_all(temp_zisk_dir_path.join("emulator-asm").join("build"))?;
        symlink(
            global_zisk_dir_path.join("emulator-asm").join("Makefile"),
            temp_zisk_dir_path.join("emulator-asm").join("Makefile"),
        )?;
        symlink(
            global_zisk_dir_path.join("emulator-asm").join("src"),
            temp_zisk_dir_path.join("emulator-asm").join("src"),
        )?;
        symlink(
            global_zisk_dir_path.join("lib-c"),
            temp_zisk_dir_path.join("lib-c"),
        )?;

        Ok(())
    }

    fn write_elf(&mut self, elf: &[u8]) -> io::Result<()> {
        fs::write(self.elf_path(), elf)
    }

    fn write_input(&mut self, input: &[u8]) -> io::Result<()> {
        fs::File::create(self.input_path()).and_then(|mut file| file.write_all(input))
    }

    fn read_execution_output(&self) -> io::Result<Vec<u8>> {
        fs::read(self.execution_output_path())
    }

    fn read_proof(&self) -> io::Result<Vec<u8>> {
        fs::read(self.proof_path())
    }

    fn write_proof(&mut self, proof: &[u8]) -> io::Result<()> {
        fs::File::create(self.proof_path()).and_then(|mut file| file.write_all(proof))
    }

    fn read_public_values(&self) -> io::Result<Vec<u8>> {
        fs::read(self.public_values_path())
    }

    fn write_public_values(&mut self, public_values: &[u8]) -> io::Result<()> {
        fs::File::create(self.public_values_path())
            .and_then(|mut file| file.write_all(public_values))
    }

    fn elf_path(&self) -> PathBuf {
        self.tempdir.path().join("guest.elf")
    }

    fn dot_zisk_dir_path(&self) -> PathBuf {
        self.tempdir.path().join(".zisk")
    }

    fn zisk_dir_path(&self) -> PathBuf {
        self.dot_zisk_dir_path().join("zisk")
    }

    fn input_path(&self) -> PathBuf {
        self.tempdir.path().join("input.bin")
    }

    fn output_dir_path(&self) -> PathBuf {
        self.tempdir.path().join("output")
    }

    fn execution_output_path(&self) -> PathBuf {
        self.output_dir_path().join("execution_output.bin")
    }

    fn public_values_path(&self) -> PathBuf {
        self.output_dir_path().join("publics.json")
    }

    fn proof_path(&self) -> PathBuf {
        self.output_dir_path().join("vadcop_final_proof.bin")
    }
}

#[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<Vec<u8>> = OnceLock::new();

    fn basic_program() -> Vec<u8> {
        BASIC_PRORGAM
            .get_or_init(|| {
                RV64_IMA_ZISK_ZKVM_ELF
                    .compile(&testing_guest_directory("zisk", "basic"))
                    .unwrap()
            })
            .to_vec()
    }

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

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

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

        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 = EreZisk::new(program, ProverResourceType::Cpu);

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

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

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