#![cfg(feature = "prove")]

use k256::SecretKey;
use k256::elliptic_curve::sec1::ToEncodedPoint;
use risc0_zkvm::{ExecutorEnv, InnerReceipt, Receipt, default_executor};
use sha2::{Digest, Sha256};
use zkwid_zkvm::{Journal, ProveInput, ZKWID_GUEST_ELF, image_id, prove, verify};

const SHIM: &str = r#"
(module
  (import "env" "keygen" (func $keygen (param i32 i32) (result i32)))
  (memory (export "memory") 1)
  (func (export "input_ptr") (result i32) (i32.const 16))
  (func (export "output_ptr") (result i32) (i32.const 64))
  (func (export "run") (param i32) (result i32)
    (drop (call $keygen (i32.const 16) (i32.const 64)))
    (i32.const 33)))
"#;

fn input() -> ProveInput {
    ProveInput {
        stages: vec![wat::parse_str(SHIM).unwrap(), b"stage2 bytes".to_vec()],
        seed: [7; 32],
        nonce: [9; 32],
        aud: "did:web:zkwid.test".into(),
    }
}

fn expected(input: &ProveInput) -> Journal {
    let key = SecretKey::from_slice(&input.seed).unwrap().public_key();
    Journal {
        measurements: input
            .stages
            .iter()
            .map(|s| Sha256::digest(s).into())
            .collect(),
        key: key.to_encoded_point(true).as_bytes().to_vec(),
        nonce: input.nonce,
        aud: input.aud.clone(),
    }
}

fn reserialize(receipt: &Receipt) -> Vec<u8> {
    bincode::serialize(receipt).unwrap()
}

#[test]
fn roundtrip() {
    let input = input();
    let env = ExecutorEnv::builder()
        .write(&input)
        .unwrap()
        .build()
        .unwrap();
    let session = default_executor().execute(env, ZKWID_GUEST_ELF).unwrap();
    println!(
        "cycles {} segments {}",
        session.cycles(),
        session.segments.len()
    );
    for s in &session.segments {
        println!("segment po2 {} cycles {}", s.po2, s.cycles);
    }

    let bytes = prove(&input).unwrap();
    assert_eq!(verify(&bytes).unwrap(), expected(&input));

    let receipt: Receipt = bincode::deserialize(&bytes).unwrap();
    let mut other = image_id();
    other[0] ^= 1;
    assert!(receipt.verify(other).is_err());

    let mut tampered = receipt.clone();
    let mut journal = expected(&input);
    journal.aud = "did:web:evil.test".into();
    tampered.journal.bytes = risc0_zkvm::serde::to_vec(&journal)
        .unwrap()
        .iter()
        .flat_map(|w| w.to_le_bytes())
        .collect();
    assert!(verify(&reserialize(&tampered)).is_err());

    let mut flipped = bytes.clone();
    let last = flipped.len() - 1;
    flipped[last] ^= 1;
    assert!(verify(&flipped).is_err());

    assert!(verify(&bytes[..bytes.len() / 2]).is_err());
    assert!(verify(b"garbage").is_err());
}

#[test]
#[ignore]
fn real_proof() {
    let input = input();
    let start = std::time::Instant::now();
    let bytes = prove(&input).unwrap();
    println!("prove {:?} receipt {} bytes", start.elapsed(), bytes.len());
    let receipt: Receipt = bincode::deserialize(&bytes).unwrap();
    assert!(matches!(receipt.inner, InnerReceipt::Succinct(_)));
    assert_eq!(verify(&bytes).unwrap(), expected(&input));
}