use std::collections::BTreeSet;
use std::collections::HashMap;

#[derive(Debug)]
struct OrbitMap {
    ids: Vec<usize>,
}

impl OrbitMap {
    fn new(n: usize) -> Self {
        Self {
            ids: (0..n).collect::<Vec<usize>>(),
        }
    }

    fn orbit_around(&mut self, base: usize, orbiter: usize) {
        self.ids[orbiter] = base;
        //println!("{:?}", self.ids);
    }

    fn range(&self, idx: usize) -> usize {
        let mut i = idx;
        let mut cnt = 0;
        while i != self.ids[i] {
            cnt += 1;
            i = self.ids[i];
        }
        cnt
    }

    fn range_to(&self, idx: usize, limit: usize) -> usize {
        let mut i = idx;
        let mut cnt = 0;
        while i != limit {
            i = self.ids[i];
            cnt += 1;
        }
        cnt - 1
    }

    fn path(&self, idx: usize) -> BTreeSet<usize> {
        let mut i = idx;
        let mut res = BTreeSet::new();
        while i != self.ids[i] {
            res.insert(i);
            i = self.ids[i];
        }
        res
    }

    fn distance(&self, p: usize, q: usize) -> usize {
        let ppath = self.path(p);
        let qpath = self.path(q);
        let common_roots = ppath.intersection(&qpath);

        common_roots
            .map(|cr| self.range_to(p, *cr) + self.range_to(q, *cr))
            .min()
            .unwrap()
    }
}

fn get_pairs(input: &str) -> Vec<Vec<String>> {
    input
        .lines()
        .map(|line| line.split(")").map(String::from).collect::<Vec<String>>())
        .collect()
}

fn get_space_objects(input: &str) -> HashMap<String, usize> {
    let items = input
        .lines()
        .map(|line| line.split(")").map(String::from).collect::<Vec<String>>())
        .flatten();

    let mut space_objects: HashMap<String, usize> = HashMap::new();
    let mut n = 0;
    items.for_each(|x| {
        if !space_objects.contains_key(&x) {
            space_objects.insert(x, n);
            n += 1;
        }
    });
    //println!("{:?}", space_objects);

    space_objects
}

fn calculate_map(input: &str) -> usize {
    let space_objects = get_space_objects(input);
    let pairs = get_pairs(input);
    let mut orbit_map = OrbitMap::new(space_objects.len());

    for pair in pairs {
        //println!("{:?}", pair);
        let base = space_objects.get(&pair[0]).unwrap();
        let orbiter = space_objects.get(&pair[1]).unwrap();
        orbit_map.orbit_around(*base, *orbiter);
    }
    //println!("{:?}", orbit_map.ids);
    (0..orbit_map.ids.len()).map(|i| orbit_map.range(i)).sum()
}

fn calculate_distance(input: &str) -> usize {
    let space_objects = get_space_objects(input);
    let pairs = get_pairs(input);
    let mut orbit_map = OrbitMap::new(space_objects.len());

    for pair in pairs {
        //println!("{:?}", pair);
        let base = space_objects.get(&pair[0]).unwrap();
        let orbiter = space_objects.get(&pair[1]).unwrap();
        orbit_map.orbit_around(*base, *orbiter);
    }
    //println!("{:?}", orbit_map.ids);
    let me = *space_objects.get("YOU").unwrap();
    let santa = *space_objects.get("SAN").unwrap();
    orbit_map.distance(me, santa)
}

pub fn main() -> std::io::Result<()> {
    let input = std::include_str!("../day6-input.txt").trim();
    println!("part1 = {:?}", calculate_map(input));
    println!("part2 = {:?}", calculate_distance(input));
    Ok(())
}

#[test]
fn test_day6_part1() {
    let input = "COM)B\n\
                 B)C\n\
                 C)D\n\
                 D)E\n\
                 E)F\n\
                 B)G\n\
                 G)H\n\
                 D)I\n\
                 E)J\n\
                 J)K\n\
                 K)L";
    assert_eq!(calculate_map(input), 42);
}

#[test]
fn test_day6_part2() {
    let input = "COM)B\n\
                 B)C\n\
                 C)D\n\
                 D)E\n\
                 E)F\n\
                 B)G\n\
                 G)H\n\
                 D)I\n\
                 E)J\n\
                 J)K\n\
                 K)L\n\
                 K)YOU\n\
                 I)SAN";
    assert_eq!(calculate_distance(input), 4);
}