use crate::decode;
use crate::hash::PerceptualHasher;
use crate::index::{BkTree, CandidateIndex};
use crate::matching::{self, CandidateEdge};
use crate::memory::MemoryBudget;
use crate::model::{Edge, ImageRecord, MatchOptions, Skipped};
use crate::pixel;
use crate::progress::Progress;
use crate::store::PixelStore;
use rayon::prelude::*;
use std::path::PathBuf;

/// Decode and hash every path, bounded by `budget` bytes of concurrent decoding.
/// Images whose decoded pair (`2·w·h·4`) exceeds the budget — or whose header is
/// unreadable — are skipped. Successful images become records with sequential
/// ids in input order; `progress` (if any) is advanced once per path.
pub fn ingest(
    paths: &[PathBuf],
    hasher: &dyn PerceptualHasher,
    budget: u64,
    progress: Option<&dyn Progress>,
) -> (Vec<ImageRecord>, Vec<Skipped>) {
    if let Some(p) = progress {
        p.begin("ingest", paths.len() as u64);
    }
    let mem = MemoryBudget::new(budget);

    let results: Vec<Result<(PathBuf, u32, u32, crate::hash::Hash64), Skipped>> = paths
        .par_iter()
        .map(|p| {
            let r = ingest_one(p, hasher, &mem, budget);
            if let Some(pr) = progress {
                pr.inc(1);
            }
            r
        })
        .collect();

    let mut records = Vec::new();
    let mut skipped = Vec::new();
    for r in results {
        match r {
            Ok((path, width, height, hash)) => {
                records.push(ImageRecord { id: records.len(), path, hash, width, height });
            }
            Err(s) => skipped.push(s),
        }
    }
    (records, skipped)
}

fn ingest_one(
    path: &std::path::Path,
    hasher: &dyn PerceptualHasher,
    mem: &MemoryBudget,
    budget: u64,
) -> Result<(PathBuf, u32, u32, crate::hash::Hash64), Skipped> {
    let (w, h) = match decode::dimensions(path) {
        Ok(d) => d,
        Err(e) => return Err(Skipped { path: path.to_path_buf(), reason: e.to_string() }),
    };
    let size = w as u64 * h as u64 * 4;
    // Decoding one image peaks at ~2x its resident RGBA8 size (the codec holds an
    // intermediate `DynamicImage` while `to_rgba8` allocates the final buffer).
    // Reserve that peak so concurrent decoding stays within budget; skip the
    // image when even its own peak does not fit.
    let peak = size.saturating_mul(2);
    if peak > budget {
        return Err(Skipped {
            path: path.to_path_buf(),
            reason: format!(
                "decoded pair ({}) exceeds memory budget ({})",
                crate::memory::human_bytes(peak),
                crate::memory::human_bytes(budget)
            ),
        });
    }

    mem.acquire(peak);
    let result = match decode::load(path) {
        Ok(img) => {
            let hash = hasher.hash(&img);
            Ok((path.to_path_buf(), img.width, img.height, hash))
        }
        Err(e) => Err(Skipped { path: path.to_path_buf(), reason: e.to_string() }),
    };
    mem.release(peak);
    result
}

/// Build a BK-tree over all record hashes.
fn build_index(records: &[ImageRecord]) -> BkTree {
    let mut tree = BkTree::new();
    for r in records {
        tree.insert(r.id, r.hash);
    }
    tree
}

/// The maximum differing-pixel count accepted for a valid match, given image
/// total pixels. `None` threshold means "no upper bound" (closest always wins).
fn accept_max(opts: &MatchOptions, total: u64) -> u64 {
    match opts.pixel_threshold {
        Some(t) => (t * total as f64).floor() as u64,
        None => u64::MAX,
    }
}

/// True if `a` is preferred over `b` for tie-breaking (stable: path, then id).
fn prefer(a: &ImageRecord, b: &ImageRecord) -> bool {
    (&a.path, a.id) < (&b.path, b.id)
}

/// Nearest-neighbour mode: for each image, its single closest same-dimension
/// match within the Hamming radius (and pixel threshold, if set).
pub fn match_nearest(
    records: &[ImageRecord],
    opts: &MatchOptions,
    store: &dyn PixelStore,
    progress: Option<&dyn Progress>,
) -> Vec<Edge> {
    let tree = build_index(records);
    if let Some(p) = progress {
        p.begin("match", records.len() as u64);
    }
    records
        .par_iter()
        .map(|rec| {
            let e = nearest_edge(rec, records, &tree, opts, store);
            if let Some(p) = progress {
                p.inc(1);
            }
            e
        })
        .collect()
}

fn nearest_edge(
    rec: &ImageRecord,
    records: &[ImageRecord],
    tree: &BkTree,
    opts: &MatchOptions,
    store: &dyn PixelStore,
) -> Edge {
    let total = rec.width as u64 * rec.height as u64;
    let amax = accept_max(opts, total);

    let mut best: Option<(&ImageRecord, u64)> = None;
    for cand_id in tree.query_within(rec.hash, opts.hamming_radius) {
        let cand = &records[cand_id];
        if cand.id == rec.id || cand.width != rec.width || cand.height != rec.height {
            continue;
        }
        let cap = best.map(|(_, d)| d).unwrap_or(amax).min(amax);
        let (primary, cand_img) = match store.get_pair(rec, cand) {
            Ok(pair) => pair,
            Err(_) => continue,
        };
        if let Some(diff) = pixel::diff_count(&primary, &cand_img, cap) {
            if diff <= amax {
                let take = match best {
                    None => true,
                    Some((bcand, bd)) => diff < bd || (diff == bd && prefer(cand, bcand)),
                };
                if take {
                    best = Some((cand, diff));
                }
            }
        }
    }

    match best {
        Some((cand, diff)) => Edge {
            image: rec.path.clone(),
            match_: Some(cand.path.clone()),
            pixel_distance: Some(diff as f64 / total as f64),
            hamming: Some(rec.hash.hamming(cand.hash)),
        },
        None => root_edge(rec),
    }
}

fn root_edge(rec: &ImageRecord) -> Edge {
    Edge { image: rec.path.clone(), match_: None, pixel_distance: None, hamming: None }
}

/// Forest mode: minimum spanning forest where each non-root image is stored as a
/// diff against its parent. Eligible edge iff `diff <= floor(threshold·w·h)` (or
/// `diff < w·h` with no threshold).
pub fn match_forest(
    records: &[ImageRecord],
    opts: &MatchOptions,
    store: &dyn PixelStore,
    progress: Option<&dyn Progress>,
) -> Vec<Edge> {
    let tree = build_index(records);
    if let Some(p) = progress {
        p.begin("match", records.len() as u64);
    }

    let edges: Vec<CandidateEdge> = records
        .par_iter()
        .flat_map_iter(|rec| {
            let v = candidate_edges_for(rec, records, &tree, opts, store);
            if let Some(p) = progress {
                p.inc(1);
            }
            v
        })
        .collect();

    let parents = matching::spanning_forest(records.len(), edges);

    records
        .iter()
        .map(|rec| match parents[rec.id] {
            Some((parent_id, diff)) => {
                let parent = &records[parent_id];
                let total = rec.width as f64 * rec.height as f64;
                Edge {
                    image: rec.path.clone(),
                    match_: Some(parent.path.clone()),
                    pixel_distance: Some(diff as f64 / total),
                    hamming: Some(rec.hash.hamming(parent.hash)),
                }
            }
            None => root_edge(rec),
        })
        .collect()
}

fn candidate_edges_for(
    rec: &ImageRecord,
    records: &[ImageRecord],
    tree: &BkTree,
    opts: &MatchOptions,
    store: &dyn PixelStore,
) -> Vec<CandidateEdge> {
    let total = rec.width as u64 * rec.height as u64;
    let amax = match opts.pixel_threshold {
        Some(t) => (t * total as f64).floor() as u64,
        None => total.saturating_sub(1),
    };

    let mut out = Vec::new();
    for cand_id in tree.query_within(rec.hash, opts.hamming_radius) {
        let cand = &records[cand_id];
        if cand.id <= rec.id || cand.width != rec.width || cand.height != rec.height {
            continue;
        }
        let (primary, cand_img) = match store.get_pair(rec, cand) {
            Ok(pair) => pair,
            Err(_) => continue,
        };
        if let Some(diff) = pixel::diff_count(&primary, &cand_img, amax) {
            if diff <= amax {
                out.push(CandidateEdge { a: rec.id, b: cand.id, diff });
            }
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hash::DHash;
    use crate::model::{MatchMode, MatchOptions};
    use crate::store::PooledPixelStore;
    use image::{Rgba, RgbaImage};
    use std::path::PathBuf;

    const BIG_BUDGET: u64 = 1 << 30; // 1 GiB — effectively unbounded for tiny test images

    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
    }

    #[test]
    fn ingest_records_good_and_skips_bad() {
        let dir = tempfile::tempdir().unwrap();
        let a = write(dir.path(), "a.png", 8, 8, 10);
        let bad = dir.path().join("bad.png");
        std::fs::write(&bad, b"not an image").unwrap();

        let (records, skipped) = ingest(&[a, bad], &DHash, BIG_BUDGET, None);
        assert_eq!(records.len(), 1);
        assert_eq!(records[0].id, 0);
        assert_eq!(skipped.len(), 1);
    }

    #[test]
    fn nearest_finds_closest_same_dimension_image() {
        let dir = tempfile::tempdir().unwrap();
        // Two near-identical 8x8 images (c=10 and c=11) and one different size.
        let a = write(dir.path(), "a.png", 8, 8, 10);
        let b = write(dir.path(), "b.png", 8, 8, 11);
        let c = write(dir.path(), "c.png", 4, 4, 10);

        let (records, _) = ingest(&[a, b, c], &DHash, BIG_BUDGET, None);
        let store = PooledPixelStore::new(BIG_BUDGET);
        let opts = MatchOptions { mode: MatchMode::Nearest, hamming_radius: 64, pixel_threshold: None };

        let edges = match_nearest(&records, &opts, &store, None);
        let edge_a = edges.iter().find(|e| e.image.ends_with("a.png")).unwrap();
        assert!(edge_a.match_.as_ref().unwrap().ends_with("b.png"));
        // c.png is a different size -> no same-dimension candidate -> root.
        let edge_c = edges.iter().find(|e| e.image.ends_with("c.png")).unwrap();
        assert_eq!(edge_c.match_, None);
    }

    #[test]
    fn nearest_threshold_rejects_far_matches() {
        let dir = tempfile::tempdir().unwrap();
        let a = write(dir.path(), "a.png", 8, 8, 0);
        let b = write(dir.path(), "b.png", 8, 8, 255); // every pixel differs
        let (records, _) = ingest(&[a, b], &DHash, BIG_BUDGET, None);
        let store = PooledPixelStore::new(BIG_BUDGET);
        let opts = MatchOptions { mode: MatchMode::Nearest, hamming_radius: 64, pixel_threshold: Some(0.1) };

        let edges = match_nearest(&records, &opts, &store, None);
        // 100% of pixels differ, above the 10% threshold -> both roots.
        assert!(edges.iter().all(|e| e.match_.is_none()));
    }

    #[test]
    fn forest_links_cluster_into_tree_with_single_root() {
        let dir = tempfile::tempdir().unwrap();
        // Three genuinely near-identical 8x8 images: a shared base, each
        // differing in only a pixel or two (diff_count compares pixels exactly,
        // so the fixtures must actually share most pixels).
        let make = |name: &str, dots: &[(u32, u32)]| {
            let mut img = RgbaImage::from_pixel(8, 8, Rgba([10, 10, 10, 255]));
            for &(x, y) in dots {
                img.put_pixel(x, y, Rgba([200, 200, 200, 255]));
            }
            let path = dir.path().join(name);
            img.save(&path).unwrap();
            path
        };
        let a = make("a.png", &[]);
        let b = make("b.png", &[(0, 0)]);
        let c = make("c.png", &[(0, 0), (1, 1)]);
        let (records, _) = ingest(&[a, b, c], &DHash, BIG_BUDGET, None);
        let store = PooledPixelStore::new(BIG_BUDGET);
        // At most 2/64 pixels differ between any pair — far under a 50% threshold.
        let opts = MatchOptions { mode: MatchMode::Forest, hamming_radius: 64, pixel_threshold: Some(0.5) };

        let edges = match_forest(&records, &opts, &store, None);
        // Exactly one root (match == None) across the connected cluster.
        let roots = edges.iter().filter(|e| e.match_.is_none()).count();
        assert_eq!(roots, 1);
    }

    #[test]
    fn forest_threshold_isolates_distinct_images() {
        let dir = tempfile::tempdir().unwrap();
        let a = write(dir.path(), "a.png", 8, 8, 0);
        let b = write(dir.path(), "b.png", 8, 8, 255); // all pixels differ
        let (records, _) = ingest(&[a, b], &DHash, BIG_BUDGET, None);
        let store = PooledPixelStore::new(BIG_BUDGET);
        let opts = MatchOptions { mode: MatchMode::Forest, hamming_radius: 64, pixel_threshold: Some(0.1) };

        let edges = match_forest(&records, &opts, &store, None);
        assert!(edges.iter().all(|e| e.match_.is_none())); // both roots
    }

    #[test]
    fn ingest_assigns_contiguous_ids_with_interleaved_failures() {
        // Failures interleaved among good files must NOT leave gaps in ids —
        // record ids must equal their index in the returned Vec.
        let dir = tempfile::tempdir().unwrap();
        let g0 = write(dir.path(), "g0.png", 8, 8, 10);
        let bad1 = dir.path().join("bad1.png");
        std::fs::write(&bad1, b"not an image").unwrap();
        let g1 = write(dir.path(), "g1.png", 8, 8, 11);
        let bad2 = dir.path().join("bad2.png");
        std::fs::write(&bad2, b"not an image").unwrap();
        let g2 = write(dir.path(), "g2.png", 8, 8, 12);

        let (records, skipped) = ingest(&[g0, bad1, g1, bad2, g2], &DHash, BIG_BUDGET, None);

        assert_eq!(records.len(), 3);
        assert_eq!(skipped.len(), 2);
        for (i, rec) in records.iter().enumerate() {
            assert_eq!(rec.id, i, "id must equal index");
        }
        assert!(records[0].path.ends_with("g0.png"));
        assert!(records[1].path.ends_with("g1.png"));
        assert!(records[2].path.ends_with("g2.png"));
    }

    #[test]
    fn nearest_tie_breaks_by_smaller_path() {
        // Target t differs from both a and z in exactly one pixel (equal
        // distance). The tie must resolve to the lexicographically smaller path.
        let dir = tempfile::tempdir().unwrap();
        let make = |name: &str, dot: (u32, u32)| {
            let mut img = RgbaImage::from_pixel(8, 8, Rgba([10, 10, 10, 255]));
            img.put_pixel(dot.0, dot.1, Rgba([200, 200, 200, 255]));
            let path = dir.path().join(name);
            img.save(&path).unwrap();
            path
        };
        // t is the plain base; a and z each differ from t in one (distinct) pixel.
        let base = dir.path().join("t.png");
        RgbaImage::from_pixel(8, 8, Rgba([10, 10, 10, 255])).save(&base).unwrap();
        let a = make("a.png", (0, 0));
        let z = make("z.png", (7, 7));

        let (records, _) = ingest(&[base, a, z], &DHash, BIG_BUDGET, None);
        let store = PooledPixelStore::new(BIG_BUDGET);
        let opts = MatchOptions { mode: MatchMode::Nearest, hamming_radius: 64, pixel_threshold: None };

        let edges = match_nearest(&records, &opts, &store, None);
        let edge_t = edges.iter().find(|e| e.image.ends_with("t.png")).unwrap();
        assert!(edge_t.match_.as_ref().unwrap().ends_with("a.png"));
    }

    #[test]
    fn ingest_skips_image_whose_pair_exceeds_budget() {
        let dir = tempfile::tempdir().unwrap();
        let small = write(dir.path(), "small.png", 4, 4, 10);
        let big = write(dir.path(), "big.png", 64, 64, 20); // 64*64*4 = 16384 bytes

        // Budget admits the 4x4 pair (128 B) but not the 64x64 pair (32768 B).
        let budget = 4 * 4 * 4 * 2 + 1; // 129 bytes
        let (records, skipped) = ingest(&[small, big], &DHash, budget, None);

        assert_eq!(records.len(), 1);
        assert!(records[0].path.ends_with("small.png"));
        assert_eq!(skipped.len(), 1);
        assert!(skipped[0].path.ends_with("big.png"));
        assert!(skipped[0].reason.contains("memory budget"));
    }
}
