use itertools::Itertools;

use crate::{spatial::Point3D, union_find::UnionFind};

pub fn day8_part1(input: &str, connection_limit: usize) -> String {
    let (box_count, distanceless_pairs) = get_pairs(input, connection_limit);

    let mut ufo = UnionFind::default();

    for (l, r) in distanceless_pairs {
        let l = ufo.make_set(l);
        let r = ufo.make_set(r);
        ufo.union(l, r);
    }

    let mut rss = ufo.root_sizes();
    rss.sort();
    let product = rss.pop().unwrap() * rss.pop().unwrap() * rss.pop().unwrap();
    product.to_string()
}
pub fn day8_part2(input: &str, connection_limit: usize) -> String {
    let (box_count, distanceless_pairs) = get_pairs(input, connection_limit);
    println!("{}", distanceless_pairs.len());
    let mut ufo = UnionFind::default();
    let mut index = 0;

    let (l, r) = loop {
        let (l, r) = distanceless_pairs[index];

        let lindex = ufo.make_set(l);
        let rindex = ufo.make_set(r);
        ufo.union(lindex, rindex);

        let rss = ufo.root_sizes();
        println!("{:?}", rss);
        if rss.len() > 1 || rss[0] < box_count {
            index += 1;
        } else {
            break dbg!((l, r));
        }
    };
    let product = l.x * r.x;
    product.to_string()
}

fn get_pairs(input: &str, connection_limit: usize) -> (usize, Vec<(Point3D, Point3D)>) {
    let boxes = parse(input);
    let mut pairs = boxes
        .iter()
        .combinations(2)
        .map(|pair| (pair[0], pair[1]))
        .map(|(left, right)| (left.euclidean_distance(right), left, right))
        .collect_vec();
    pairs.sort_by(|(ld, _, _), (rd, _, _)| ld.partial_cmp(rd).unwrap());
    (
        boxes.len(),
        pairs
            .into_iter()
            .map(|(_, l, r)| (*l, *r))
            .take(connection_limit)
            .collect_vec(),
    )
}

fn parse(input: &str) -> Vec<Point3D> {
    input
        .lines()
        .map(|line| {
            line.split(',')
                .map(|coord| coord.parse::<usize>().unwrap())
                .collect_tuple()
                .unwrap()
        })
        .map(|(x, y, z)| Point3D { x, y, z })
        .collect()
}