//! Conformance corpus for the avatars extension (`avatar-conformance/`).
//! Pose-equivalence, a special case of replay-equivalence: the same
//! signed animation op log replays to the identical per-bone-transform
//! state hash on every sink — hash equality, no tolerance. Sealed
//! vectors carry checkpoint hashes; live vectors carry an adversarial
//! arrival order + the final pose hash every sink must reach.
//!
//!   wai_avatar_vectors gen <dir>  |  verify <dir>

use std::path::Path;
use wai::world::replay::{run_sealed_dispatch, Ingest, LiveWorld};
use wai::world::skeletal::{pack_bone, pack_fx, skel_mutation as m};
use wai::world::{
    mint_node_id, CausalOp, Fx, Skeletal0, SimContract, WorldContainer, WorldState, Xoshiro256pp,
    CAP_SIM_SKELETAL0, NUMERIC_FIXED64, RESOLVE_LWW,
};
use ed25519_dalek::SigningKey;

const PUBLISHER_SEED: u8 = 9;
const IDENT: [Fx; 4] = [Fx::ZERO, Fx::ZERO, Fx::ZERO, Fx::ONE];

fn main() {
    let a: Vec<String> = std::env::args().collect();
    match (a.get(1).map(String::as_str), a.get(2)) {
        (Some("gen"), Some(d)) => generate(Path::new(d)),
        (Some("verify"), Some(d)) => verify(Path::new(d)),
        _ => {
            eprintln!("usage: wai_avatar_vectors <gen|verify> <corpus-dir>");
            std::process::exit(2);
        }
    }
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}
fn key(b: u8) -> SigningKey {
    SigningKey::from_bytes(&[b; 32])
}
fn fxf(v: f64) -> Fx {
    Fx::from_bits((v * 4294967296.0).round() as i64)
}
fn v3(x: f64, y: f64, z: f64) -> [Fx; 3] {
    [fxf(x), fxf(y), fxf(z)]
}
const ONE3: [Fx; 3] = [Fx::ONE, Fx::ONE, Fx::ONE];

/// Authoring-only axis-angle → unit quaternion (sin/cos here, off-line —
/// the op carries the `Fx` quaternion; the floor never sees trig).
fn quat(axis: [f64; 3], radians: f64) -> [Fx; 4] {
    let h = radians * 0.5;
    let s = h.sin();
    [fxf(axis[0] * s), fxf(axis[1] * s), fxf(axis[2] * s), fxf(h.cos())]
}

fn contract() -> SimContract {
    SimContract {
        tick_rate: 30,
        numeric: NUMERIC_FIXED64.into(),
        rng_alg: "xoshiro256++".into(),
        rng_seed: [0u8; 32],
        sim: CAP_SIM_SKELETAL0.into(),
        resolve_default: RESOLVE_LWW.into(),
        checkpoint_every: 1,
    }
}

fn scene() -> WorldContainer {
    let c = contract();
    WorldContainer {
        contract_bytes: c.to_canonical_json(),
        contract: c,
        graph_bytes: WorldState::<wai::world::Bone>::new().canonical_bytes(&Skeletal0),
        assets: vec![],
        ops: vec![],
        seal: None,
    }
}

// ---------------------------------------------------------------------------
// Vector op sets
// ---------------------------------------------------------------------------

/// A 3-bone arm (shoulder → elbow → wrist) waving: the shoulder and
/// elbow rotate about z over several ticks. The wrist rides along by FK.
fn ops_arm_wave() -> Vec<CausalOp> {
    let k = key(1);
    let pk = k.verifying_key().to_bytes();
    let shoulder = mint_node_id(&pk, 1);
    let elbow = mint_node_id(&pk, 2);
    let wrist = mint_node_id(&pk, 3);
    let z16 = [0u8; 16];
    let zero3 = [Fx::ZERO; 3];

    let mut ops = vec![
        CausalOp::sign(&k, 1, 1, vec![], shoulder, m::CREATE, pack_bone(&z16, &zero3, &IDENT, &ONE3)),
        CausalOp::sign(&k, 2, 1, vec![], elbow, m::CREATE, pack_bone(&shoulder, &v3(0.0, 1.0, 0.0), &IDENT, &ONE3)),
        CausalOp::sign(&k, 3, 1, vec![], wrist, m::CREATE, pack_bone(&elbow, &v3(0.0, 1.0, 0.0), &IDENT, &ONE3)),
    ];
    // a 4-frame wave: shoulder ±, elbow bends
    let mut lam = 4u64;
    let frames = [(2u64, 0.3, 0.5), (4, -0.2, 0.9), (6, 0.4, 0.3), (8, 0.0, 0.6)];
    for (tick, sh, el) in frames {
        ops.push(CausalOp::sign(&k, lam, tick, vec![], shoulder, m::SET_ROTATION, pack_fx(&quat([0.0, 0.0, 1.0], sh))));
        lam += 1;
        ops.push(CausalOp::sign(&k, lam, tick, vec![], elbow, m::SET_ROTATION, pack_fx(&quat([0.0, 0.0, 1.0], el))));
        lam += 1;
    }
    ops
}

/// A 5-bone humanoid stub (pelvis → spine → head; spine → l_arm, r_arm)
/// struck into a single expression pose, then the head turns.
fn ops_humanoid() -> Vec<CausalOp> {
    let k = key(2);
    let pk = k.verifying_key().to_bytes();
    let pelvis = mint_node_id(&pk, 1);
    let spine = mint_node_id(&pk, 2);
    let head = mint_node_id(&pk, 3);
    let l_arm = mint_node_id(&pk, 4);
    let r_arm = mint_node_id(&pk, 5);
    let z16 = [0u8; 16];

    vec![
        CausalOp::sign(&k, 1, 1, vec![], pelvis, m::CREATE, pack_bone(&z16, &v3(0.0, 1.0, 0.0), &IDENT, &ONE3)),
        CausalOp::sign(&k, 2, 1, vec![], spine, m::CREATE, pack_bone(&pelvis, &v3(0.0, 0.6, 0.0), &IDENT, &ONE3)),
        CausalOp::sign(&k, 3, 1, vec![], head, m::CREATE, pack_bone(&spine, &v3(0.0, 0.5, 0.0), &IDENT, &ONE3)),
        CausalOp::sign(&k, 4, 1, vec![], l_arm, m::CREATE, pack_bone(&spine, &v3(-0.4, 0.4, 0.0), &IDENT, &ONE3)),
        CausalOp::sign(&k, 5, 1, vec![], r_arm, m::CREATE, pack_bone(&spine, &v3(0.4, 0.4, 0.0), &IDENT, &ONE3)),
        // raise the right arm (about z) and turn the head (about y)
        CausalOp::sign(&k, 6, 3, vec![], r_arm, m::SET_ROTATION, pack_fx(&quat([0.0, 0.0, 1.0], -1.2))),
        CausalOp::sign(&k, 7, 5, vec![], head, m::SET_ROTATION, pack_fx(&quat([0.0, 1.0, 0.0], 0.6))),
    ]
}

fn seal_from_live(mut c: WorldContainer, ops: &[CausalOp], final_tick: u64) -> WorldContainer {
    let mut live = LiveWorld::open(Skeletal0, c.clone()).expect("open");
    for op in ops {
        let outcome = live.ingest(op.clone()).expect("ingest");
        assert!(matches!(outcome, Ingest::Applied | Ingest::Buffered));
    }
    let final_hash = live.advance_to(final_tick).expect("advance");
    c.ops = live.canonical_ops();
    c.seal(&key(PUBLISHER_SEED), final_tick, final_hash);
    c
}

fn write_ops(path: &Path, ops: &[CausalOp]) {
    let bytes: Vec<u8> = ops.iter().flat_map(|o| o.to_bytes()).collect();
    std::fs::write(path, bytes).expect("write ops");
}

fn shuffles(ops: &[CausalOp], n: usize) -> Vec<Vec<CausalOp>> {
    let mut rng = Xoshiro256pp::from_seed([11u8; 32]);
    (0..n)
        .map(|_| {
            let mut s = ops.to_vec();
            for i in (1..s.len()).rev() {
                let j = (rng.next_u64() % (i as u64 + 1)) as usize;
                s.swap(i, j);
            }
            s
        })
        .collect()
}

// ---------------------------------------------------------------------------

fn generate(dir: &Path) {
    let vectors = dir.join("vectors");
    std::fs::create_dir_all(&vectors).expect("mkdir");
    let mut manifest = serde_json::Map::new();

    let sealed: Vec<(&str, WorldContainer)> = vec![
        ("arm_wave", seal_from_live(scene(), &ops_arm_wave(), 8)),
        ("humanoid", seal_from_live(scene(), &ops_humanoid(), 5)),
    ];
    for (name, c) in &sealed {
        let file = format!("{name}.wwld");
        std::fs::write(vectors.join(&file), c.to_bytes()).expect("write wwld");
        let cps = run_sealed_dispatch(c).expect("replay");
        manifest.insert(
            file,
            serde_json::json!({
                "type": "sealed",
                "sim": CAP_SIM_SKELETAL0,
                "publisher": hex(&key(PUBLISHER_SEED).verifying_key().to_bytes()),
                "checkpoints": cps.iter()
                    .map(|cp| serde_json::json!({"tick": cp.tick, "hash": hex(&cp.hash)}))
                    .collect::<Vec<_>>(),
            }),
        );
    }

    let scene_file = "scene.wwld";
    std::fs::write(vectors.join(scene_file), scene().to_bytes()).expect("write scene");
    for (i, s) in shuffles(&ops_arm_wave(), 3).into_iter().enumerate() {
        let name = format!("arm_wave_shuffled_{}", i + 1);
        let file = format!("{name}.ops");
        write_ops(&vectors.join(&file), &s);
        let mut live = LiveWorld::open(Skeletal0, scene()).expect("open");
        for op in &s {
            live.ingest(op.clone()).expect("ingest");
        }
        let h = live.advance_to(8).expect("advance");
        manifest.insert(
            file,
            serde_json::json!({ "type": "live", "scene": scene_file, "final_tick": 8, "final_hash": hex(&h) }),
        );
    }

    std::fs::write(
        dir.join("expected_hashes.json"),
        serde_json::to_string_pretty(&serde_json::Value::Object(manifest)).unwrap() + "\n",
    )
    .expect("manifest");
    println!("generated avatar vectors into {}", dir.display());
}

fn verify(dir: &Path) {
    let m: serde_json::Value =
        serde_json::from_slice(&std::fs::read(dir.join("expected_hashes.json")).expect("manifest")).expect("parse");
    let vectors = dir.join("vectors");
    let mut fails = 0;
    for (file, spec) in m.as_object().expect("obj") {
        let kind = spec["type"].as_str().unwrap();
        let bytes = std::fs::read(vectors.join(file)).expect("vector file");
        let ok = match kind {
            "sealed" => {
                let c = WorldContainer::from_bytes(&bytes).expect("parse wwld");
                match run_sealed_dispatch(&c) {
                    Ok(cps) => {
                        let exp = spec["checkpoints"].as_array().unwrap();
                        cps.len() == exp.len()
                            && exp.iter().zip(&cps).all(|(e, got)| {
                                e["tick"].as_u64() == Some(got.tick) && e["hash"].as_str() == Some(&hex(&got.hash))
                            })
                    }
                    Err(_) => false,
                }
            }
            "live" => {
                let scene_bytes = std::fs::read(vectors.join(spec["scene"].as_str().unwrap())).expect("scene");
                let c = WorldContainer::from_bytes(&scene_bytes).expect("parse scene");
                let ops = CausalOp::parse_log(&bytes).expect("ops");
                let mut live = LiveWorld::open(Skeletal0, c).expect("open");
                for op in ops {
                    live.ingest(op).expect("ingest");
                }
                let h = live.advance_to(spec["final_tick"].as_u64().unwrap()).expect("advance");
                hex(&h) == spec["final_hash"].as_str().unwrap()
            }
            other => {
                eprintln!("unknown vector type {other}");
                false
            }
        };
        println!("{} {file}", if ok { "PASS" } else { "FAIL" });
        if !ok {
            fails += 1;
        }
    }
    if fails > 0 {
        eprintln!("{fails} failed — NOT conformant");
        std::process::exit(1);
    }
    println!("all vectors pass: pose-equivalence holds for avatars");
}
