use std::path::PathBuf;
use std::process::Command;

use anyhow::{Context, bail, ensure};
use atproto_identity::key::{KeyType, generate_key};
use wasmi::core::Trap;
use wasmi::{Caller, Engine, Linker, Module, Store};

pub fn seed() -> [u8; 32] {
    generate_key(KeyType::K256Private)
        .expect("k256 key generation")
        .bytes()
        .try_into()
        .expect("32-byte k256 scalar")
}

pub fn stage_bytes(name: &str) -> anyhow::Result<Vec<u8>> {
    let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("../../stages")
        .join(name);
    let manifest = dir.join("Cargo.toml");
    ensure!(manifest.is_file(), "no stage {name}");
    let status = Command::new(std::env::var_os("CARGO").unwrap_or_else(|| "cargo".into()))
        .args([
            "build",
            "--quiet",
            "--release",
            "--target",
            "wasm32-unknown-unknown",
            "--manifest-path",
        ])
        .arg(&manifest)
        .env_remove("CARGO_TARGET_DIR")
        .env_remove("RUSTFLAGS")
        .env_remove("CARGO_ENCODED_RUSTFLAGS")
        .status()
        .context("running cargo")?;
    ensure!(status.success(), "building stage {name}: {status}");
    let wasm = dir.join(format!("target/wasm32-unknown-unknown/release/{name}.wasm"));
    std::fs::read(&wasm).with_context(|| format!("reading {}", wasm.display()))
}

pub fn run_stage(wasm: &[u8], input: &[u8]) -> anyhow::Result<Vec<u8>> {
    let engine = Engine::default();
    let module = Module::new(&engine, wasm).map_err(anyhow::Error::msg)?;
    let mut store = Store::new(&engine, ());
    let mut linker = Linker::<()>::new(&engine);
    linker
        .func_wrap("env", "keygen", keygen)
        .map_err(anyhow::Error::msg)?;
    let instance = linker
        .instantiate(&mut store, &module)
        .and_then(|pre| pre.start(&mut store))
        .map_err(anyhow::Error::msg)?;
    let memory = instance
        .get_memory(&store, "memory")
        .context("no memory export")?;
    let input_ptr = instance
        .get_typed_func::<(), i32>(&store, "input_ptr")
        .map_err(anyhow::Error::msg)?
        .call(&mut store, ())
        .map_err(anyhow::Error::msg)?;
    let output_ptr = instance
        .get_typed_func::<(), i32>(&store, "output_ptr")
        .map_err(anyhow::Error::msg)?
        .call(&mut store, ())
        .map_err(anyhow::Error::msg)?;
    memory
        .write(&mut store, input_ptr as u32 as usize, input)
        .map_err(anyhow::Error::msg)?;
    let len = instance
        .get_typed_func::<i32, i32>(&store, "run")
        .map_err(anyhow::Error::msg)?
        .call(&mut store, i32::try_from(input.len())?)
        .map_err(anyhow::Error::msg)?;
    let Ok(len) = usize::try_from(len) else {
        bail!("negative output length {len}");
    };
    let mut output = vec![0; len];
    memory
        .read(&store, output_ptr as u32 as usize, &mut output)
        .map_err(anyhow::Error::msg)?;
    Ok(output)
}

fn keygen(mut caller: Caller<'_, ()>, seed_ptr: i32, out_ptr: i32) -> Result<i32, Trap> {
    let memory = caller
        .get_export("memory")
        .and_then(|export| export.into_memory())
        .ok_or_else(|| Trap::new("no memory export"))?;
    let mut seed = [0u8; 32];
    memory
        .read(&caller, seed_ptr as u32 as usize, &mut seed)
        .map_err(|e| Trap::new(e.to_string()))?;
    let key = public_key(&seed).map_err(Trap::new)?;
    memory
        .write(&mut caller, out_ptr as u32 as usize, &key)
        .map_err(|e| Trap::new(e.to_string()))?;
    Ok(33)
}

pub(crate) fn public_key(seed: &[u8; 32]) -> Result<[u8; 33], String> {
    let secret = k256::SecretKey::from_slice(seed).map_err(|e| e.to_string())?;
    let point =
        k256::elliptic_curve::sec1::ToEncodedPoint::to_encoded_point(&secret.public_key(), true);
    point
        .as_bytes()
        .try_into()
        .map_err(|_| "bad point".to_string())
}

#[cfg(test)]
mod tests {
    use super::*;
    use atproto_identity::key::{KeyData, to_public};

    #[test]
    fn seed_is_k256_scalar() {
        let a = seed();
        assert_ne!(a, seed());
        let key = to_public(&KeyData::new(KeyType::K256Private, a.to_vec())).unwrap();
        assert_eq!(key.bytes(), public_key(&a).unwrap());
    }

    #[test]
    fn invalid_seed_fails() {
        assert!(public_key(&[0; 32]).is_err());
        assert!(public_key(&[0xff; 32]).is_err());
    }

    #[test]
    fn rejects_garbage_wasm() {
        assert!(run_stage(b"not wasm", b"").is_err());
    }
}