use crate::{
    compile::compile_zisk_program,
    error::{ExecuteError, ProveError, VerifyError, ZiskError},
};
use serde::{Deserialize, Serialize};
use std::{
    fs,
    io::{self, Write},
    path::{Path, PathBuf},
    process::{Command, Stdio},
    time,
};
use tempfile::{TempDir, tempdir};
use zkvm_interface::{
    Compiler, Input, ProgramExecutionReport, ProgramProvingReport, ProverResourceType, zkVM,
    zkVMError,
};

mod compile;
mod error;

#[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(path_to_program: &Path) -> Result<Self::Program, Self::Error> {
        compile_zisk_program(path_to_program).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 EreZisk {}

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

        let input_bytes = input
            .iter()
            .try_fold(Vec::new(), |mut acc, item| {
                acc.extend(item.as_bytes().map_err(ExecuteError::SerializeInput)?);
                Ok(acc)
            })
            .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("--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() {
            return Err(ZiskError::Execute(ExecuteError::ZiskemuFailed {
                status: output.status,
            })
            .into());
        }
        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))?;

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

    fn prove(&self, input: &Input) -> Result<(Vec<u8>, ProgramProvingReport), zkVMError> {
        // Write ELF and serialized input to file.

        let input_bytes = input
            .iter()
            .try_fold(Vec::new(), |mut acc, item| {
                acc.extend(item.as_bytes().map_err(ProveError::SerializeInput)?);
                Ok(acc)
            })
            .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)))?;

        // Setup ROM.

        // FIXME: This currently uses global build directory `${HOME}/.zisk/zisk/emulator-asm`
        //        which causes `compile_zisk_program` to panic if ran in parallel.
        //        We should create a temporary directory and copy only necessary
        //        data to setup each ELF.
        let status = Command::new("cargo-zisk")
            .arg("rom-setup")
            .arg("--elf")
            .arg(tempdir.elf_path())
            .arg("--output-dir")
            .arg(tempdir.rom_dir_path())
            .status()
            .map_err(|e| ZiskError::Prove(ProveError::CargoZiskRomSetup { source: e }))?;

        if !status.success() {
            return Err(ZiskError::Prove(ProveError::CargoZiskRomSetupFailed { status }).into());
        }

        // Prove.

        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("--asm")
                    .arg(tempdir.asm_path())
                    .arg("--input")
                    .arg(tempdir.input_path())
                    .arg("--output-dir")
                    .arg(tempdir.output_dir_path())
                    // FIXME: Not sure why if we don't set the flag `aggregation`
                    //        there will be no proof, but ideally we should
                    //        only generate the core proof as other zkVMs.
                    .arg("--aggregation")
                    .status()
                    .map_err(|e| ZiskError::Prove(ProveError::CargoZiskProve { source: e }))?;

                if !status.success() {
                    return Err(
                        ZiskError::Prove(ProveError::CargoZiskProveFailed { status }).into(),
                    );
                }
            }
            ProverResourceType::Gpu => {
                // TODO: Need to install another version of `cargo-zisk` with
                //       `features = gpu` and call it here.
                unimplemented!()
            }
            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 proof_with_public_values = ZiskProofWithPublicValues {
            proof: tempdir
                .read_proof()
                .map_err(|e| ZiskError::Prove(ProveError::TempDir(e)))?,
            public_values: tempdir
                .read_public_values()
                .map_err(|e| ZiskError::Prove(ProveError::TempDir(e)))?,
        };
        let bytes = bincode::serialize(&proof_with_public_values)
            .map_err(|err| ZiskError::Prove(ProveError::Bincode(err)))?;

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

    fn verify(&self, bytes: &[u8]) -> Result<(), zkVMError> {
        // 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() {
            return Err(ZiskError::Verify(VerifyError::InvalidProof(
                String::from_utf8_lossy(&output.stderr).to_string(),
            ))
            .into());
        }

        Ok(())
    }
}

struct ZiskTempDir {
    tempdir: TempDir,
    elf_hash: Option<String>,
}

impl ZiskTempDir {
    fn new() -> io::Result<Self> {
        let tempdir = Self {
            tempdir: tempdir()?,
            elf_hash: None,
        };
        fs::create_dir(tempdir.rom_dir_path())?;
        fs::create_dir_all(tempdir.proof_dir_path())?;
        Ok(tempdir)
    }

    fn write_elf(&mut self, elf: &[u8]) -> io::Result<()> {
        self.elf_hash = Some(blake3::hash(elf).to_hex().to_string());
        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_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("elf")
    }

    fn rom_dir_path(&self) -> PathBuf {
        self.tempdir.path().join("rom")
    }

    fn asm_path(&self) -> PathBuf {
        let elf_path = self.elf_path();
        let stem = elf_path
            .file_stem()
            .expect("ELF file has name")
            .to_str()
            .expect("ELF file name is valid UTF-8");
        let hash = self.elf_hash.as_deref().expect("setup has been called");
        self.rom_dir_path().join(format!("{stem}-{hash}-mt.bin"))
    }

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

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

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

    fn proof_dir_path(&self) -> PathBuf {
        self.output_dir_path().join("proofs")
    }

    fn proof_path(&self) -> PathBuf {
        self.proof_dir_path().join("vadcop_final_proof.json")
    }
}

#[cfg(test)]
mod execute_tests {
    use super::*;

    fn get_compiled_test_zisk_elf() -> Result<Vec<u8>, ZiskError> {
        let test_guest_path = get_execute_test_guest_program_path();
        RV64_IMA_ZISK_ZKVM_ELF::compile(&test_guest_path)
    }

    fn get_execute_test_guest_program_path() -> PathBuf {
        let workspace_dir = env!("CARGO_WORKSPACE_DIR");
        PathBuf::from(workspace_dir)
            .join("tests")
            .join("zisk")
            .join("execute")
            .join("basic")
            .canonicalize()
            .expect("Failed to find or canonicalize test guest program at <CARGO_WORKSPACE_DIR>/tests/execute/zisk")
    }

    #[test]
    fn test_execute_zisk_dummy_input() {
        let elf_bytes = get_compiled_test_zisk_elf()
            .expect("Failed to compile test ZisK guest for execution test");

        let mut input_builder = Input::new();
        let n: u32 = 42;
        let a: u16 = 42;
        input_builder.write(n);
        input_builder.write(a);

        let zkvm = EreZisk::new(elf_bytes, ProverResourceType::Cpu);

        let result = zkvm.execute(&input_builder);

        if let Err(e) = &result {
            panic!("Execution error: {e:?}");
        }
    }

    #[test]
    fn test_execute_zisk_no_input_for_guest_expecting_input() {
        let elf_bytes = get_compiled_test_zisk_elf()
            .expect("Failed to compile test ZisK guest for execution test");

        let empty_input = Input::new();

        let zkvm = EreZisk::new(elf_bytes, ProverResourceType::Cpu);
        assert!(zkvm.execute(&empty_input).is_err());
    }
}

#[cfg(test)]
mod prove_tests {
    use std::path::PathBuf;

    use super::*;
    use zkvm_interface::Input;

    fn get_prove_test_guest_program_path() -> PathBuf {
        let workspace_dir = env!("CARGO_WORKSPACE_DIR");
        PathBuf::from(workspace_dir)
            .join("tests")
            .join("zisk")
            .join("prove")
            .join("basic")
            .canonicalize()
            .expect("Failed to find or canonicalize test guest program at <CARGO_WORKSPACE_DIR>/tests/execute/zisk")
    }

    fn get_compiled_test_zisk_elf_for_prove() -> Result<Vec<u8>, ZiskError> {
        let test_guest_path = get_prove_test_guest_program_path();
        RV64_IMA_ZISK_ZKVM_ELF::compile(&test_guest_path)
    }

    #[test]
    fn test_prove_zisk_dummy_input() {
        let elf_path = get_compiled_test_zisk_elf_for_prove()
            .expect("Failed to compile test ZisK guest for proving test");

        let mut input_builder = Input::new();
        let n: u32 = 42;
        let a: u16 = 42;
        input_builder.write(n);
        input_builder.write(a);

        let zkvm = EreZisk::new(elf_path, ProverResourceType::Cpu);

        let proof_bytes = match zkvm.prove(&input_builder) {
            Ok((prove_result, _)) => prove_result,
            Err(err) => {
                panic!("Proving error in test: {err:?}");
            }
        };

        assert!(!proof_bytes.is_empty(), "Proof bytes should not be empty.");

        assert!(zkvm.verify(&proof_bytes).is_ok());

        let invalid_proof_bytes = {
            let mut invalid_proof: ZiskProofWithPublicValues =
                bincode::deserialize(&proof_bytes).unwrap();
            // alter the first digit of `evals[0][0]`
            invalid_proof.proof[40] = invalid_proof.proof[40].overflowing_add(1).0;
            bincode::serialize(&invalid_proof).unwrap()
        };
        assert!(zkvm.verify(&invalid_proof_bytes).is_err());

        // TODO: Check public inputs
    }

    #[test]
    fn test_prove_zisk_fails_on_bad_input_causing_execution_failure() {
        let elf_path = get_compiled_test_zisk_elf_for_prove()
            .expect("Failed to compile test ZisK guest for proving test");

        let empty_input = Input::new();

        let zkvm = EreZisk::new(elf_path, ProverResourceType::Cpu);
        assert!(zkvm.prove(&empty_input).is_err());
    }
}
