use crate::decode::DecodedImage;

/// A 64-bit perceptual hash.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct Hash64(pub u64);

impl Hash64 {
    /// Number of differing bits between two hashes (Hamming distance).
    pub fn hamming(self, other: Hash64) -> u32 {
        (self.0 ^ other.0).count_ones()
    }
}

/// Computes a 64-bit perceptual hash of an image. Implementations must be
/// thread-safe so hashing can run in parallel.
pub trait PerceptualHasher: Send + Sync {
    fn hash(&self, img: &DecodedImage) -> Hash64;
}

/// Nearest-neighbour downscale to `tw x th` grayscale (Rec.601 luminance).
fn downscale_gray(img: &DecodedImage, tw: usize, th: usize) -> Vec<u8> {
    let mut out = vec![0u8; tw * th];
    let sw = img.width.max(1) as usize;
    let sh = img.height.max(1) as usize;
    for ty in 0..th {
        let sy = (ty * sh) / th;
        for tx in 0..tw {
            let sx = (tx * sw) / tw;
            let idx = (sy * sw + sx) * 4;
            let r = img.rgba[idx] as u32;
            let g = img.rgba[idx + 1] as u32;
            let b = img.rgba[idx + 2] as u32;
            out[ty * tw + tx] = ((r * 30 + g * 59 + b * 11) / 100) as u8;
        }
    }
    out
}

/// Difference hash: downscale to 9x8 grayscale, set a bit where each pixel is
/// darker than its right-hand neighbour. 8 rows x 8 comparisons = 64 bits.
pub struct DHash;

impl PerceptualHasher for DHash {
    fn hash(&self, img: &DecodedImage) -> Hash64 {
        const W: usize = 9;
        const H: usize = 8;
        let small = downscale_gray(img, W, H);
        let mut bits: u64 = 0;
        let mut i = 0;
        for row in 0..H {
            for col in 0..(W - 1) {
                if small[row * W + col] < small[row * W + col + 1] {
                    bits |= 1u64 << i;
                }
                i += 1;
            }
        }
        Hash64(bits)
    }
}

/// Average hash: downscale to 8x8 grayscale, set a bit where each pixel is
/// brighter than the mean. 64 pixels = 64 bits.
pub struct AHash;

impl PerceptualHasher for AHash {
    fn hash(&self, img: &DecodedImage) -> Hash64 {
        const W: usize = 8;
        const H: usize = 8;
        let small = downscale_gray(img, W, H);
        let sum: u32 = small.iter().map(|&p| p as u32).sum();
        let mean = (sum / (W * H) as u32) as u8;
        let mut bits: u64 = 0;
        for (i, &p) in small.iter().enumerate() {
            if p > mean {
                bits |= 1u64 << i;
            }
        }
        Hash64(bits)
    }
}

/// One-dimensional DCT-II of `input` into `output` (same length).
fn dct_1d(input: &[f32], output: &mut [f32]) {
    let n = input.len();
    let factor = std::f32::consts::PI / n as f32;
    for (k, out) in output.iter_mut().enumerate() {
        let mut sum = 0.0f32;
        for (i, &x) in input.iter().enumerate() {
            sum += x * (factor * (i as f32 + 0.5) * k as f32).cos();
        }
        *out = sum;
    }
}

/// Perceptual hash (DCT-based): downscale to 32x32 grayscale, take the 2D DCT,
/// keep the top-left 8x8 low-frequency block, and set a bit where each
/// coefficient exceeds the median of that block (median computed excluding the
/// DC term). 64 coefficients = 64 bits.
pub struct PHash;

impl PerceptualHasher for PHash {
    fn hash(&self, img: &DecodedImage) -> Hash64 {
        const N: usize = 32;
        let gray = downscale_gray(img, N, N);
        let mut mat: Vec<f32> = gray.iter().map(|&p| p as f32).collect();

        // DCT over rows.
        let mut tmp = vec![0.0f32; N];
        for r in 0..N {
            dct_1d(&mat[r * N..r * N + N], &mut tmp);
            mat[r * N..r * N + N].copy_from_slice(&tmp);
        }
        // DCT over columns.
        let mut col_in = vec![0.0f32; N];
        let mut col_out = vec![0.0f32; N];
        for c in 0..N {
            for r in 0..N {
                col_in[r] = mat[r * N + c];
            }
            dct_1d(&col_in, &mut col_out);
            for r in 0..N {
                mat[r * N + c] = col_out[r];
            }
        }

        // Top-left 8x8 low-frequency coefficients.
        let mut coeffs = [0.0f32; 64];
        let mut idx = 0;
        for r in 0..8 {
            for c in 0..8 {
                coeffs[idx] = mat[r * N + c];
                idx += 1;
            }
        }

        // Median of all but the DC term (coeffs[0]).
        let mut rest: Vec<f32> = coeffs[1..].to_vec();
        rest.sort_by(|a, b| a.total_cmp(b));
        let median = rest[rest.len() / 2];

        let mut bits: u64 = 0;
        for (i, &c) in coeffs.iter().enumerate() {
            if c > median {
                bits |= 1u64 << i;
            }
        }
        Hash64(bits)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decode::DecodedImage;

    fn gradient(w: u32, h: u32) -> DecodedImage {
        let mut rgba = Vec::with_capacity((w * h * 4) as usize);
        for y in 0..h {
            for x in 0..w {
                let v = ((x * 255) / w.max(1)) as u8;
                let _ = y;
                rgba.extend_from_slice(&[v, v, v, 255]);
            }
        }
        DecodedImage { width: w, height: h, rgba }
    }

    fn solid(w: u32, h: u32, c: u8) -> DecodedImage {
        DecodedImage { width: w, height: h, rgba: [c, c, c, 255].repeat((w * h) as usize) }
    }

    #[test]
    fn hamming_counts_differing_bits() {
        assert_eq!(Hash64(0b0000).hamming(Hash64(0b0000)), 0);
        assert_eq!(Hash64(0b1011).hamming(Hash64(0b0001)), 2);
    }

    #[test]
    fn dhash_identical_images_have_zero_distance() {
        let a = gradient(40, 40);
        let b = gradient(40, 40);
        assert_eq!(DHash.hash(&a).hamming(DHash.hash(&b)), 0);
    }

    #[test]
    fn dhash_distinguishes_gradient_from_solid() {
        let g = gradient(40, 40);
        let s = solid(40, 40, 128);
        assert!(DHash.hash(&g).hamming(DHash.hash(&s)) > 5);
    }

    #[test]
    fn ahash_identical_images_have_zero_distance() {
        let a = gradient(40, 40);
        let b = gradient(40, 40);
        assert_eq!(AHash.hash(&a).hamming(AHash.hash(&b)), 0);
    }

    #[test]
    fn phash_identical_images_have_zero_distance() {
        let a = gradient(40, 40);
        let b = gradient(40, 40);
        assert_eq!(PHash.hash(&a).hamming(PHash.hash(&b)), 0);
    }

    #[test]
    fn phash_distinguishes_gradient_from_solid() {
        let g = gradient(40, 40);
        let s = solid(40, 40, 128);
        assert!(PHash.hash(&g).hamming(PHash.hash(&s)) > 0);
    }
}
