use anyhow::Result;
use searchi_semantic::{Device, GemmaEmbedder};
use std::path::Path;

#[test]
#[ignore = "Requires explicitly provisioned Google media model and CPU worker"]
fn google_cpu_pixels_match_digup_reference() -> Result<()> {
    let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set model path");
    let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set worker path");
    let mut embedder = GemmaEmbedder::load_with_device(model, worker, 1, Device::Cpu)?;
    let fixtures = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
    let reference: serde_json::Value =
        serde_json::from_str(include_str!("fixtures/digup-media-reference.json"))?;
    for (name, file, budget) in [
        ("screen-280", "digup-screen.png", 280),
        ("photo-140", "digup-photo.jpg", 140),
        ("frame-70", "digup-photo.jpg", 70),
    ] {
        let path = fixtures.join(file);
        let vector = embedder
            .embed_images_with_budget(&[path.as_path()], budget)?
            .remove(0);
        let cosine: f64 = vector
            .iter()
            .zip(reference["vectors"][name].as_array().unwrap())
            .map(|(a, b)| f64::from(*a) * b.as_f64().unwrap())
            .sum();
        println!("Google image {name} reference cosine={cosine:.8}");
        assert!(cosine >= 0.999, "Image parity below0.999: {name} {cosine}");
    }
    assert_eq!(embedder.info.device, "cpu");
    let wav = include_bytes!("fixtures/digup-speech.wav");
    let mut offset = 12;
    let mut samples = Vec::new();
    while offset + 8 <= wav.len() {
        let length = u32::from_le_bytes(wav[offset + 4..offset + 8].try_into()?) as usize;
        if &wav[offset..offset + 4] == b"data" {
            for pair in wav[offset + 8..offset + 8 + length].as_chunks::<2>().0 {
                samples.push(f32::from(i16::from_le_bytes(*pair)) / 32768.0);
            }
            break;
        }
        offset += 8 + length + (length % 2);
    }
    let vector = embedder.embed_audio_pcm(&samples)?;
    let cosine: f64 = vector
        .iter()
        .zip(reference["vectors"]["speech"].as_array().unwrap())
        .map(|(a, b)| f64::from(*a) * b.as_f64().unwrap())
        .sum();
    println!("Google audio reference cosine={cosine:.8}");
    assert!(
        cosine >= 0.999,
        "Audio reference parity below0.999: {cosine}"
    );
    Ok(())
}