use k256::SecretKey;
use k256::elliptic_curve::sec1::ToEncodedPoint;
use risc0_zkvm::guest::env;
use risc0_zkvm::sha::{Impl, Sha256};
use wasmi::{Caller, Engine, Extern, Linker, Memory, Module, Store};
use zkwid_zkvm_core::{Journal, ProveInput};

fn memory(caller: &Caller<'_, ()>) -> Memory {
    match caller.get_export("memory") {
        Some(Extern::Memory(m)) => m,
        _ => panic!("no memory"),
    }
}

fn keygen(mut caller: Caller<'_, ()>, seed_ptr: i32, out_ptr: i32) -> i32 {
    let mem = memory(&caller);
    let mut seed = [0u8; 32];
    mem.read(&caller, seed_ptr as u32 as usize, &mut seed)
        .expect("seed");
    let key = SecretKey::from_slice(&seed).expect("secret key");
    let point = key.public_key().to_encoded_point(true);
    let bytes = point.as_bytes();
    mem.write(&mut caller, out_ptr as u32 as usize, bytes)
        .expect("out");
    bytes.len() as i32
}

fn run_shim(wasm: &[u8], seed: &[u8; 32]) -> Vec<u8> {
    let engine = Engine::default();
    let module = Module::new(&engine, wasm).expect("module");
    let mut store = Store::new(&engine, ());
    let mut linker = <Linker<()>>::new(&engine);
    linker.func_wrap("env", "keygen", keygen).expect("link");
    let instance = linker
        .instantiate(&mut store, &module)
        .expect("instantiate")
        .start(&mut store)
        .expect("start");
    let mem = instance.get_memory(&store, "memory").expect("memory");
    let input_ptr = instance
        .get_typed_func::<(), i32>(&store, "input_ptr")
        .expect("input_ptr");
    let output_ptr = instance
        .get_typed_func::<(), i32>(&store, "output_ptr")
        .expect("output_ptr");
    let run = instance
        .get_typed_func::<i32, i32>(&store, "run")
        .expect("run");
    let inp = input_ptr.call(&mut store, ()).expect("input_ptr call") as u32 as usize;
    mem.write(&mut store, inp, seed).expect("write seed");
    let len = run.call(&mut store, seed.len() as i32).expect("run call");
    assert_eq!(len, 33, "output length");
    let out = output_ptr.call(&mut store, ()).expect("output_ptr call") as u32 as usize;
    let mut key = vec![0u8; 33];
    mem.read(&store, out, &mut key).expect("read key");
    key
}

fn main() {
    let input: ProveInput = env::read();
    assert!(!input.stages.is_empty(), "no stages");
    let measurements = input
        .stages
        .iter()
        .map(|s| Impl::hash_bytes(s).as_bytes().try_into().expect("digest"))
        .collect();
    let key = run_shim(&input.stages[0], &input.seed);
    env::commit(&Journal {
        measurements,
        key,
        nonce: input.nonce,
        aud: input.aud,
    });
}