use std::collections::HashMap;

use regex::Regex;

pub fn run() {
    part1();
    part2();
}

fn part1() {
    let input = include_str!("../input/day8.txt");
    // let input = r#"LLR

    // AAA = (BBB, BBB)
    // BBB = (AAA, ZZZ)
    // ZZZ = (ZZZ, ZZZ)
    // "#;

    let (instructions, rest) = input.split_once("\n\n").unwrap();
    let links = parse_links(rest);
    dbg!(count_steps("AAA", "ZZZ", instructions, &links));
}

fn part2() {
    let input = include_str!("../input/day8.txt");
    // let input = r#"LR

    // 11A = (11B, XXX)
    // 11B = (XXX, 11Z)
    // 11Z = (11B, XXX)
    // 22A = (22B, XXX)
    // 22B = (22C, 22C)
    // 22C = (22Z, 22Z)
    // 22Z = (22B, 22B)
    // XXX = (XXX, XXX)
    // "#;

    let (instructions, rest) = input.split_once("\n\n").unwrap();
    let links = parse_links(rest);
    dbg!(count_ghost_steps(instructions, &links));
}

fn parse_links(s: &str) -> HashMap<&str, (&str, &str)> {
    let mut map = HashMap::new();
    let re = Regex::new(r#"^(\w{3}) = \((\w{3}), (\w{3})\)$"#).unwrap();
    for l in s.lines() {
        let l = l.trim();
        if l.is_empty() {
            continue;
        }
        let captures = re.captures(l).unwrap();
        let (_, [node, left, right]) = captures.extract::<3>();
        assert!(!map.contains_key(node));
        map.insert(node, (left, right));
    }
    map
}

fn count_steps(
    start: &str,
    end: &str,
    instructions: &str,
    links: &HashMap<&str, (&str, &str)>,
) -> usize {
    let mut steps = 0;
    let mut cur = start;
    for c in instructions.chars().cycle() {
        if cur == end {
            break;
        }

        steps += 1;

        let link = links[cur];
        if c == 'L' {
            cur = link.0;
        } else if c == 'R' {
            cur = link.1;
        } else {
            panic!("invalid char {}", c);
        }
    }
    steps
}

fn ghost_starts<'a>(links: &HashMap<&'a str, (&str, &str)>) -> Vec<&'a str> {
    links
        .keys()
        .filter(|k| k.ends_with("A"))
        .map(|k| *k)
        .collect()
}

fn count_ghost_steps<'a>(
    instructions: &'a str,
    links: &HashMap<&'a str, (&'a str, &'a str)>,
) -> usize {
    let starts = ghost_starts(links);
    let steps_until_end = starts
        .iter()
        .map(|start| count_steps_until_end(start, instructions, links));
    lcm(steps_until_end)
}

fn count_steps_until_end(
    start: &str,
    instructions: &str,
    links: &HashMap<&str, (&str, &str)>,
) -> usize {
    let mut steps = 0;
    let mut cur = start;
    for c in instructions.chars().cycle() {
        if cur.ends_with("Z") {
            break;
        }

        steps += 1;

        let link = links[cur];
        if c == 'L' {
            cur = link.0;
        } else if c == 'R' {
            cur = link.1;
        } else {
            panic!();
        }
    }
    steps
}

pub fn lcm(xs: impl Iterator<Item = usize>) -> usize {
    let mut a = 1;
    for x in xs {
        a = a * x / gcd(a, x);
    }
    a
}

fn gcd(mut a: usize, mut b: usize) -> usize {
    loop {
        let r = a % b;
        if r == 0 {
            return b;
        } else {
            a = b;
            b = r;
        }
    }
}

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

    #[test]
    fn test_lcm() {
        assert_eq!(lcm([4, 6].into_iter()), 12);
    }

    #[test]
    fn test_gcd() {
        assert_eq!(gcd(10, 50), 10);
        assert_eq!(gcd(252, 105), 21);
    }
}