//! Palette grouping and simultaneous replacement using original RGB values.
use super::Swatch;
use std::collections::HashMap;

#[derive(Clone, Debug)]
pub struct ColorGroup {
    pub rgb: [u8; 3],
    pub count: u64,
    pub members: Vec<[u8; 3]>,
}

/// Most frequent colors seed groups. Every member is within `tolerance` of
/// its seed in each RGB channel; nearby shades cannot form unbounded chains.
/// Spatial buckets keep large palettes from requiring an all-pairs search.
pub fn group_palette(palette: &[Swatch], tolerance: u8) -> Vec<ColorGroup> {
    if tolerance == 0 {
        return palette
            .iter()
            .map(|s| ColorGroup {
                rgb: s.rgb,
                count: s.count,
                members: vec![s.rgb],
            })
            .collect();
    }
    let mut sorted: Vec<_> = palette.iter().collect();
    sorted.sort_by_key(|s| (std::cmp::Reverse(s.count), s.rgb));
    let mut groups: Vec<ColorGroup> = Vec::new();
    let mut buckets = HashMap::<[i16; 3], usize>::new();
    let side = i16::from(tolerance) + 1;
    for swatch in sorted {
        let bucket = swatch.rgb.map(|v| i16::from(v) / side);
        let mut best = None;
        for x in -1..=1 {
            for y in -1..=1 {
                for z in -1..=1 {
                    if let Some(&index) =
                        buckets.get(&[bucket[0] + x, bucket[1] + y, bucket[2] + z])
                    {
                        let seed = groups[index].rgb;
                        let delta =
                            std::array::from_fn::<_, 3, _>(|i| seed[i].abs_diff(swatch.rgb[i]));
                        if delta.iter().all(|&d| d <= tolerance) {
                            let distance: u32 = delta.iter().map(|&d| u32::from(d).pow(2)).sum();
                            let candidate = (distance, index);
                            if best.is_none_or(|old| candidate < old) {
                                best = Some(candidate);
                            }
                        }
                    }
                }
            }
        }
        if let Some((_, index)) = best {
            groups[index].count += swatch.count;
            groups[index].members.push(swatch.rgb);
        } else {
            buckets.insert(bucket, groups.len());
            groups.push(ColorGroup {
                rgb: swatch.rgb,
                count: swatch.count,
                members: vec![swatch.rgb],
            });
        }
    }
    groups.sort_by_key(|g| (std::cmp::Reverse(g.count), g.rgb));
    groups
}

impl ColorGroup {
    pub fn highlight_rgba(&self, rgba: &[u8]) -> Vec<u8> {
        let members: std::collections::HashSet<_> = self.members.iter().copied().collect();
        let mut output = rgba.to_vec();
        for pixel in output.as_chunks_mut::<4>().0 {
            if pixel[3] == 0 {
                continue;
            }
            if members.contains(&[pixel[0], pixel[1], pixel[2]]) {
                pixel[..3].copy_from_slice(&[255, 0, 255]);
            } else {
                for channel in &mut pixel[..3] {
                    *channel /= 4;
                }
            }
        }
        output
    }
}

pub struct ColorMapping {
    pub source: [u8; 3],
    pub target: [u8; 3],
    pub members: Vec<[u8; 3]>,
}

#[derive(Default)]
pub struct ColorMappings(pub Vec<ColorMapping>);

impl ColorMappings {
    pub fn insert_group(&mut self, group: &ColorGroup) {
        let members: std::collections::HashSet<_> = group.members.iter().copied().collect();
        // The latest explicit mapping owns overlapping colors. Keep unaffected
        // members of earlier mappings, and remove entries that become empty.
        for entry in &mut self.0 {
            entry.members.retain(|rgb| !members.contains(rgb));
        }
        self.0.retain(|entry| !entry.members.is_empty());
        self.0.push(ColorMapping {
            source: group.rgb,
            target: group.rgb,
            members: group.members.clone(),
        });
    }

    #[cfg(test)]
    pub fn insert(&mut self, source: [u8; 3], target: [u8; 3]) {
        self.insert_group(&ColorGroup {
            rgb: source,
            count: 1,
            members: vec![source],
        });
        self.0.last_mut().unwrap().target = target;
    }

    pub fn lookup(&self) -> HashMap<[u8; 3], [u8; 3]> {
        self.0
            .iter()
            .flat_map(|entry| entry.members.iter().map(move |&rgb| (rgb, entry.target)))
            .collect()
    }

    pub fn apply_rgba(&self, rgba: &[u8]) -> Vec<u8> {
        let lookup = self.lookup();
        let mut output = rgba.to_vec();
        for pixel in output.as_chunks_mut::<4>().0 {
            if pixel[3] != 0
                && let Some(target) = lookup.get(&[pixel[0], pixel[1], pixel[2]])
            {
                pixel[..3].copy_from_slice(target);
            }
        }
        output
    }
}

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

    fn swatch(rgb: [u8; 3], count: u64) -> Swatch {
        Swatch {
            rgb,
            count,
            opaque_count: count,
        }
    }

    #[test]
    fn grouping_combines_counts_without_chaining_and_keeps_exact_mode() {
        let palette = vec![
            swatch([100, 0, 0], 100),
            swatch([105, 0, 0], 3),
            swatch([110, 0, 0], 2),
            swatch([0, 255, 0], 1),
        ];
        let groups = group_palette(&palette, 5);
        assert_eq!(groups.len(), 3);
        assert_eq!(groups[0].members, [[100, 0, 0], [105, 0, 0]]);
        assert_eq!(groups[0].count, 103);
        assert_eq!(groups.iter().map(|g| g.count).sum::<u64>(), 106);
        assert_eq!(group_palette(&palette, 0).len(), 4);
        let mut reversed = palette;
        reversed.reverse();
        assert_eq!(group_palette(&reversed, 5)[0].members, groups[0].members);
    }

    #[test]
    fn group_mapping_uses_fixed_members_and_latest_mapping_owns_overlap() {
        let groups = group_palette(&[swatch([255, 0, 0], 10), swatch([252, 2, 1], 1)], 4);
        let mut mappings = ColorMappings::default();
        mappings.insert_group(&groups[0]);
        mappings.0[0].target = [0, 0, 255];
        assert_eq!(mappings.apply_rgba(&[252, 2, 1, 128]), [0, 0, 255, 128]);
        mappings.insert([252, 2, 1], [0, 255, 0]);
        assert_eq!(mappings.0[0].members, [[255, 0, 0]]);
        assert_eq!(
            mappings.apply_rgba(&[252, 2, 1, 255, 255, 0, 0, 255]),
            [0, 255, 0, 255, 0, 0, 255, 255]
        );
    }

    #[test]
    fn highlight_matches_group_members_and_preserves_transparency() {
        let group = ColorGroup {
            rgb: [250, 0, 0],
            count: 2,
            members: vec![[250, 0, 0], [253, 2, 0]],
        };
        assert_eq!(
            group.highlight_rgba(&[253, 2, 0, 128, 0, 0, 252, 255, 250, 0, 0, 0]),
            [255, 0, 255, 128, 0, 0, 63, 255, 250, 0, 0, 0]
        );
    }

    #[test]
    fn dense_palette_preserves_every_color_and_respects_tolerance() {
        let palette: Vec<_> = (0..=255)
            .flat_map(|r| (0..=255).step_by(8).map(move |g| swatch([r, g, 128], 1)))
            .collect();
        let groups = group_palette(&palette, 8);
        let members: std::collections::HashSet<_> = groups
            .iter()
            .flat_map(|g| g.members.iter().copied())
            .collect();
        assert_eq!(members.len(), palette.len());
        assert_eq!(
            groups.iter().map(|g| g.count).sum::<u64>(),
            palette.len() as u64
        );
        for group in groups {
            assert!(
                group
                    .members
                    .iter()
                    .all(|rgb| (0..3).all(|i| rgb[i].abs_diff(group.rgb[i]) <= 8))
            );
        }
    }

    #[test]
    fn replacements_are_exact_simultaneous_and_preserve_alpha() {
        let mut mappings = ColorMappings::default();
        mappings.insert([255, 0, 0], [0, 255, 0]);
        mappings.insert([0, 255, 0], [255, 0, 0]);
        let source = [255, 0, 0, 255, 0, 255, 0, 128, 254, 0, 0, 255, 255, 0, 0, 0];
        assert_eq!(
            mappings.apply_rgba(&source),
            [0, 255, 0, 255, 255, 0, 0, 128, 254, 0, 0, 255, 255, 0, 0, 0]
        );
        assert_eq!(ColorMappings::default().apply_rgba(&source), source);
    }
}