use super::array::Dtype;
use super::{ops, qsa_attention};

#[test]
#[ignore = "Metal and exhaustive pinned gate fixtures; external supervisor required"]
fn issue109_gate_phase_and_dtype_laws() {
    super::configure_sources().unwrap();
    let _context = crate::engine::metal::gpu::Context::open_qwen(0).unwrap();
    let streams = super::stream::Streams::new(Some(super::allocator::Allocator::new().unwrap()));
    let gpu = streams.default_stream(super::stream::Device::Gpu).unwrap();
    let root = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE109_GATE").unwrap());
    for (name, dtype) in [
        ("bf16", Dtype::BF16),
        ("f16", Dtype::F16),
        ("f32", Dtype::F32),
    ] {
        let load = |part: &str| {
            let bytes = std::fs::read(root.join(format!("{name}-{part}.f32"))).unwrap();
            let array = super::array::Array::new(
                &[(bytes.len() / 4) as i32],
                Dtype::F32,
                super::scalar_buffer(&bytes).unwrap(),
            )
            .unwrap();
            ops::astype(&array, dtype, false, gpu).unwrap()
        };
        let gate = load("gate");
        let output = load("output");
        for phase in ["prefill", "ar_decode", "decode_verify"] {
            let actual =
                qsa_attention::attention_gate(&output, &gate, phase == "decode_verify", gpu)
                    .unwrap();
            let actual = ops::read_f32_as_f64(
                &ops::astype(&actual, Dtype::F32, false, gpu).unwrap(),
                &streams,
                gpu,
            )
            .unwrap()
            .into_iter()
            .flat_map(|v| (v as f32).to_le_bytes())
            .collect::<Vec<_>>();
            let expected = std::fs::read(root.join(format!("{name}-{phase}.f32"))).unwrap();
            assert_eq!(actual.len(), expected.len());
            let different = actual
                .chunks_exact(4)
                .zip(expected.chunks_exact(4))
                .enumerate()
                .filter_map(|(i, (a, b))| (a != b).then_some(i))
                .collect::<Vec<_>>();
            assert!(
                different.is_empty(),
                "{name}/{phase}: {} different rows; first {:?}",
                different.len(),
                different.first()
            );
            println!(
                "{}",
                serde_json::json!({"case":"gate","dtype":name,"phase":phase,"values":actual.len()/4,"exact":true})
            );
            if dtype == Dtype::BF16 {
                let inject =
                    super::hyper::inject_gate(&gate, phase == "decode_verify", gpu).unwrap();
                let actual = ops::read_f32_as_f64(
                    &ops::astype(&inject, Dtype::F32, false, gpu).unwrap(),
                    &streams,
                    gpu,
                )
                .unwrap()
                .into_iter()
                .flat_map(|v| (v as f32).to_le_bytes())
                .collect::<Vec<_>>();
                assert!(
                    actual == std::fs::read(root.join(format!("inject-{phase}.f32"))).unwrap(),
                    "inject/{phase}"
                );
                println!(
                    "{}",
                    serde_json::json!({"case":"inject","phase":phase,"values":actual.len()/4,"exact":true})
                );
            }
        }
    }
}

thread_local! {
    static RECEIPTS: std::cell::RefCell<Option<(std::path::PathBuf, usize)>> = const { std::cell::RefCell::new(None) };
}

pub(super) fn capture(directory: Option<std::path::PathBuf>) {
    RECEIPTS.with_borrow_mut(|slot| {
        *slot = directory.map(|d| {
            std::fs::create_dir_all(&d).unwrap();
            (d, 0)
        })
    });
}

pub(super) fn record(
    execution: &super::execution::Execution,
    ids: &super::array::Array,
    output: &super::model::Output,
    verify: bool,
) -> Result<(), String> {
    RECEIPTS.with_borrow_mut(|slot| {
        let Some((directory, index)) = slot else { return Ok(()); };
        let ids = ops::read_i64(ids, &execution.streams, execution.stream)?;
        let mut record=serde_json::json!({"index":*index,"ids":ids,"verify":verify});
        for (part, a) in [("logits",&output.logits),("hidden",&output.hidden)] {
            let a=a.as_ref().ok_or("missing #109 output")?;
            let bytes=execution.values_f32(a)?.into_iter().flat_map(f32::to_le_bytes).collect::<Vec<_>>();
            std::fs::write(directory.join(format!("{index}-{part}.f32")),&bytes).map_err(|e|e.to_string())?;
            record[part]=serde_json::json!({"shape":a.layout().shape(),"sha256":crate::engine::metal::qwen_mtplx_tests::hex_digest(&bytes)});
        }
        std::fs::write(directory.join(format!("{index}.json")),serde_json::to_vec(&record).unwrap()).map_err(|e|e.to_string())?;
        println!("{}",serde_json::json!({"case":"verify_forward","record":record}));
        *index+=1;
        Ok(())
    })
}