use std::{f64, fmt::Display};

use aoc_runner_derive::{aoc, aoc_generator};
use itertools::Itertools;

#[derive(Eq, PartialEq, Clone, Debug, Hash)]
struct Junction {
    pos: (i64, i64, i64),
}

fn distance(a: &Junction, b: &Junction) -> f64 {
    if a.pos == b.pos {
        0f64
    } else {
        (((a.pos.0 - b.pos.0).pow(2) + (a.pos.1 - b.pos.1).pow(2) + (a.pos.2 - b.pos.2).pow(2))
            as f64)
            .sqrt()
    }
}

impl Junction {
    fn distance(&self, other: &Junction) -> f64 {
        distance(self, other)
    }
}

impl Display for Junction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!(
            "({},{},{})",
            self.pos.0, self.pos.1, self.pos.2
        ))
    }
}

impl From<&str> for Junction {
    fn from(value: &str) -> Self {
        Self {
            pos: value
                .split(',')
                .map(|v| v.parse().unwrap())
                .collect_tuple()
                .unwrap(),
        }
    }
}

#[derive(Clone)]
struct Circuits {
    junctions: Vec<Junction>,
    circuits: Vec<Vec<usize>>,
}

impl Circuits {
    fn find_circuit(&self, junction: usize) -> usize {
        self.circuits
            .iter()
            .enumerate()
            .find(|(_i, c)| c.contains(&junction))
            .map(|(i, _c)| i)
            .unwrap()
    }
    fn merge_circuits(&mut self, a: usize, b: usize) {
        if a == b {
            return;
        }
        let mut items = Vec::new();
        std::mem::swap(&mut items, &mut self.circuits[b]);
        self.circuits[a].append(&mut items);
        self.circuits.remove(b);
    }

    fn connect(&mut self, a: usize, b: usize) {
        let a_circuit = self.find_circuit(a);
        let b_circuit = self.find_circuit(b);

        self.merge_circuits(a_circuit, b_circuit); // both are already in circuits, merge them
    }
}

#[aoc_generator(day8)]
fn parse(input: &str) -> Circuits {
    let junctions = input.lines().map(|l| l.into()).collect_vec();

    Circuits {
        circuits: Vec::from_iter((0..junctions.len()).map(|i| vec![i])),
        junctions,
    }
}

fn part1_impl(input: &Circuits, n: usize) -> u64 {
    let mut circuits = input.clone();
    for (a, b, _d) in circuits
        .junctions
        .iter()
        .enumerate()
        .tuple_combinations()
        .map(|((a_pos, a), (b_pos, b))| (a_pos, b_pos, a.distance(b)))
        .sorted_by(|a, b| a.2.partial_cmp(&b.2).unwrap())
        .take(n)
    {
        circuits.connect(a, b)
    }

    circuits
        .circuits
        .iter()
        .sorted_by_key(|c| c.len())
        .rev()
        .take(3)
        .map(|c| c.len())
        .reduce(|a, b| a * b)
        .unwrap() as u64
}

#[aoc(day8, part1)]
fn part1(input: &Circuits) -> u64 {
    part1_impl(input, 1000)
}

#[aoc(day8, part2)]
fn part2(input: &Circuits) -> u64 {
    let mut circuits = input.clone();

    for (a, b, _d) in circuits
        .junctions
        .iter()
        .enumerate()
        .combinations(2)
        .map(|p| (p[0].0, p[1].0, p[0].1.distance(p[1].1)))
        .sorted_by(|a, b| a.2.partial_cmp(&b.2).unwrap())
    {
        circuits.connect(a, b);
        if circuits.circuits.len() == 1 {
            return (circuits.junctions[a].pos.0 * circuits.junctions[b].pos.0) as u64;
        }
    }

    panic!()
}

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

    const EXAMPLE: &str = "162,817,812
57,618,57
906,360,560
592,479,940
352,342,300
466,668,158
542,29,236
431,825,988
739,650,466
52,470,668
216,146,977
819,987,18
117,168,530
805,96,715
346,949,466
970,615,88
941,993,340
862,61,35
984,92,344
425,690,689";

    #[test]
    fn part1_example() {
        assert_eq!(part1_impl(&parse(EXAMPLE), 10), 40);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EXAMPLE)), 25272);
    }
}