use image::{Rgba, RgbaImage};
use macchikane_core::{engine, DHash, MatchMode, MatchOptions, PooledPixelStore};
use std::path::PathBuf;

fn write(dir: &std::path::Path, name: &str, w: u32, h: u32, c: u8) -> PathBuf {
    let mut img = RgbaImage::new(w, h);
    for p in img.pixels_mut() {
        *p = Rgba([c, c, c, 255]);
    }
    let path = dir.join(name);
    img.save(&path).unwrap();
    path
}

fn run_in_pool<F: FnOnce() -> R + Send, R: Send>(threads: usize, f: F) -> R {
    rayon::ThreadPoolBuilder::new()
        .num_threads(threads)
        .build()
        .unwrap()
        .install(f)
}

fn signature(edges: &[macchikane_core::Edge]) -> Vec<(String, Option<String>)> {
    edges
        .iter()
        .map(|e| {
            (
                e.image.display().to_string(),
                e.match_.as_ref().map(|p| p.display().to_string()),
            )
        })
        .collect()
}

#[test]
fn forest_output_is_identical_across_thread_counts() {
    let dir = tempfile::tempdir().unwrap();
    // A cluster of near-identical images plus a couple of outliers.
    let paths: Vec<PathBuf> = (0..12)
        .map(|i| write(dir.path(), &format!("img{i:02}.png"), 16, 16, (i * 3) as u8))
        .collect();

    let opts = MatchOptions {
        mode: MatchMode::Forest,
        hamming_radius: 64,
        pixel_threshold: Some(0.9),
    };

    let single = run_in_pool(1, || {
        let (records, _) = engine::ingest(&paths, &DHash, 1 << 30, None);
        let store = PooledPixelStore::new(1 << 30);
        engine::match_forest(&records, &opts, &store, None)
    });

    let many = run_in_pool(8, || {
        let (records, _) = engine::ingest(&paths, &DHash, 1 << 30, None);
        let store = PooledPixelStore::new(1 << 30);
        engine::match_forest(&records, &opts, &store, None)
    });

    assert_eq!(signature(&single), signature(&many));
}

#[test]
fn nearest_output_is_identical_across_thread_counts() {
    let dir = tempfile::tempdir().unwrap();
    let paths: Vec<PathBuf> = (0..12)
        .map(|i| write(dir.path(), &format!("n{i:02}.png"), 16, 16, (i * 3) as u8))
        .collect();

    let opts = MatchOptions {
        mode: MatchMode::Nearest,
        hamming_radius: 64,
        pixel_threshold: None,
    };

    let single = run_in_pool(1, || {
        let (records, _) = engine::ingest(&paths, &DHash, 1 << 30, None);
        let store = PooledPixelStore::new(1 << 30);
        engine::match_nearest(&records, &opts, &store, None)
    });
    let many = run_in_pool(8, || {
        let (records, _) = engine::ingest(&paths, &DHash, 1 << 30, None);
        let store = PooledPixelStore::new(1 << 30);
        engine::match_nearest(&records, &opts, &store, None)
    });

    assert_eq!(signature(&single), signature(&many));
}

#[test]
fn forest_output_is_identical_under_tight_budget() {
    let dir = tempfile::tempdir().unwrap();
    let paths: Vec<PathBuf> = (0..8)
        .map(|i| write(dir.path(), &format!("t{i:02}.png"), 16, 16, (i * 5) as u8))
        .collect();
    let opts = MatchOptions { mode: MatchMode::Forest, hamming_radius: 64, pixel_threshold: Some(0.9) };

    let (records, _) = engine::ingest(&paths, &DHash, 1 << 30, None);

    let roomy = {
        let store = PooledPixelStore::new(1 << 30);
        engine::match_forest(&records, &opts, &store, None)
    };
    let tight = {
        let store = PooledPixelStore::new(2048); // exactly one 16x16 pair
        engine::match_forest(&records, &opts, &store, None)
    };
    assert_eq!(signature(&roomy), signature(&tight));
}
