//! Day 8: Haunted Wasteland

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

use itertools::Itertools;

pub const DAY: u8 = 8;
pub const TITLE: &str = "Haunted Wasteland";

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Direction {
    Right,
    Left,
}

#[derive(Debug)]
struct Puzzle {
    directions: Vec<Direction>,
    nodes: HashMap<String, (String, String)>,
}

fn parse(input: Vec<String>) -> Puzzle {
    let mut lines = input.iter();
    let directions = lines
        .next()
        .unwrap()
        .chars()
        .filter_map(|c| match c {
            'R' => Some(Direction::Right),
            'L' => Some(Direction::Left),
            _ => None,
        })
        .collect();

    lines.next();

    let nodes = lines
        .map(|l| {
            let (def, assg) = l.split_once(" = (").unwrap();
            let (a, b) = assg.split_once(", ").unwrap();
            let b = b.strip_suffix(')').unwrap();
            (def.to_owned(), (a.to_owned(), b.to_owned()))
        })
        .collect();

    Puzzle { directions, nodes }
}

fn solve(parsed: &Puzzle, start: &str, end: fn(&str) -> bool) -> u64 {
    let mut steps = 0;
    let mut pos = start;
    let mut directions = parsed.directions.iter().cycle();

    while !end(pos) {
        let d = directions.next().unwrap();
        let ways = parsed.nodes.get(pos).unwrap();
        pos = match d {
            Direction::Right => ways.1.as_str(),
            Direction::Left => ways.0.as_str(),
        };

        steps += 1;
    }
    steps
}

pub fn part1(input: Vec<String>) -> String {
    let parsed = parse(input);
    solve(&parsed, "AAA", |n| n == "ZZZ").to_string()
}

pub fn part2(input: Vec<String>) -> String {
    let parsed = parse(input);

    let starts = parsed
        .nodes
        .keys()
        .filter(|k| k.ends_with('A'))
        .map(String::as_str)
        .collect_vec();
    let steps = starts
        .iter()
        .map(|start| solve(&parsed, start, |n| n.ends_with('Z')))
        .collect_vec();
    steps
        .into_iter()
        .reduce(num::integer::lcm)
        .unwrap()
        .to_string()
}

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

    #[test]
    fn t_example1() {
        assert_eq!(part1(crate::read_example(DAY, 1)), "2");
    }

    #[test]
    fn t_part1() {
        assert_eq!(part1(crate::read_input(DAY)), "18023");
    }

    #[test]
    fn t_example2() {
        assert_eq!(part2(crate::read_example(DAY, 2)), "6");
    }

    #[test]
    fn t_part2() {
        assert_eq!(part2(crate::read_input(DAY)), "14449445933179");
    }
}