mod cli;

use clap::Parser;
use std::collections::{HashMap, HashSet};

use dotenv::dotenv;
use image::{DynamicImage, GenericImageView, Rgb};

use crate::cli::Cli;

// TODO (medium priority): replace hashset with another hashmap<u8, Vec<(u32, u32)>>;
// that way we can track pixel locations for colours.
type ColourSet = HashMap<u8, HashMap<u8, HashSet<u8>>>;
type Colour = Rgb<u8>;

fn extract_colours(img: DynamicImage) -> ColourSet {
    let mut out: ColourSet = HashMap::new();
    for (_x, _y, p) in img.pixels() {
        out.entry(p[0])
            .or_insert_with(HashMap::new)
            .entry(p[1])
            .or_insert_with(HashSet::new)
            .insert(p[2]);
    }
    out
}

fn extract_colours_set(img: DynamicImage) -> HashSet<Rgb<u8>> {
    let mut out = HashSet::new();
    img.pixels().for_each(|(_x, _y, p)| {
        out.insert(p);
    });
    out.into_iter().map(|c| Rgb([c[0], c[1], c[2]])).collect()
}

fn closest_colour(colour: &Colour, all: &ColourSet) -> Option<Colour> {
    // use redmean to calculate the distance between colours
    let mut closest: Option<Colour> = None;
    let mut dist: f32 = f32::MAX;

    let (colour_r, colour_g, colour_b) = (colour[0] as f32, colour[1] as f32, colour[2] as f32);
    for (r, gbs) in all.iter() {
        let r_bar = 0.5 * (*r as f32 + colour_r);
        let r_squared = (*r as f32 - colour_r) * (*r as f32 - colour_r);
        for (g, bs) in gbs.iter() {
            let g_squared = (*g as f32 - colour_g) * (*g as f32 - colour_g);
            for b in bs {
                let b_squared = (*b as f32 - colour_b) * (*b as f32 - colour_b);
                let delta = (2.0 + r_bar / 256.0) * r_squared
                    + 4.0 * g_squared
                    + (2.0 + (255.0 - r_bar) / 256.0) * b_squared;
                if delta < dist {
                    dist = delta;
                    closest = Some(Rgb([*r, *g, *b]));
                }
            }
        }
    }
    return closest;
}

fn map_image_list(all_cols: &ColourSet, cols: Vec<Colour>) -> HashMap<Colour, Colour> {
    let mut output_cols: HashMap<Colour, Colour> = HashMap::new();

    for c in cols {
        if let Some(col) = closest_colour(&c, &all_cols) {
            output_cols.insert(c, col);
        }
    }
    output_cols
}

fn map_image_image(img_1_cols: &ColourSet, img_2: DynamicImage) -> HashMap<Colour, Colour> {
    let img_2_cols = extract_colours_set(img_2);

    let mut output_cols: HashMap<Colour, Colour> = HashMap::new();
    img_2_cols.iter().for_each(|c| {
        if let Some(col) = closest_colour(c, &img_1_cols) {
            output_cols.insert(*c, col);
        }
    });

    output_cols
}

fn encode_hex(c: &Colour) -> String {
    format!("#{:02x}{:02x}{:02x}", c[0], c[1], c[2])
}

fn pretty_print(h: HashMap<Colour, Colour>) {
    h.iter()
        .for_each(|(k, v)| println!("{}: {}", encode_hex(k), encode_hex(v)))
}

fn main() {
    dotenv().ok();
    pretty_env_logger::init();
    let cli = Cli::parse();

    let set = extract_colours(cli.image);

    if let Some(cs) = cli.colours {
        let colours = map_image_list(&set, cs);
        pretty_print(colours);
    }

    if let Some(i) = cli.second_image {
        let colours = map_image_image(&set, i);
        pretty_print(colours);
    }
}