use std::collections::HashMap;

use anyhow::Context;

use crate::{Problem, Solution};

pub struct Day08;

impl Problem for Day08 {
    const DAY: u8 = 8;

    const INPUT: &'static str = include_str!("../input/day_08.txt");
}

impl Solution for Day08 {
    type Answer1 = usize;

    type Answer2 = usize;

    fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
        let (mut instructions, nodes) = parse_input(input)?;
        let mut step = 0;
        let mut node = "AAA";
        while node != "ZZZ" {
            step += 1;
            let (left, right) = nodes
                .get(node)
                .with_context(|| format!("No node found: {node}"))?;
            node = match instructions.next() {
                Some('L') => left,
                Some('R') => right,
                Some(c) => anyhow::bail!("Invalid instruction found: {c}"),
                None => anyhow::bail!("Failed to get next instruction (this should not happen)"),
            }
        }

        Ok(step)
    }

    /// Calculate the least common multiple of each ghost's step count
    fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
        let (mut instructions, nodes) = parse_input(input)?;

        nodes
            .keys()
            .filter(|s| s.ends_with('A'))
            .map(|mut node| {
                let mut step = 0;
                while !node.ends_with('Z') {
                    step += 1;
                    let (left, right) = nodes.get(node).expect("Invalid node");
                    node = match instructions.next() {
                        Some('L') => left,
                        Some('R') => right,
                        Some(c) => panic!("Invalid instruction found: {c}"),
                        None => {
                            panic!("Failed to get next instruction (this should not happen)")
                        }
                    }
                }
                step
            })
            .reduce(lcd)
            .ok_or_else(|| anyhow::anyhow!("Failed to find any ghosts"))
    }
}

fn parse_input(
    input: &str,
) -> anyhow::Result<(impl Iterator<Item = char> + '_, HashMap<&str, (&str, &str)>)> {
    let (instructions, rest) = input
        .trim()
        .split_once("\n\n")
        .context("Invalid input format")?;

    let nodes = rest
        .trim()
        .lines()
        .map(|l| {
            let (name, paths) = l
                .split_once(" = ")
                .with_context(|| format!("Invalid node format: {l}"))?;

            let paths = paths
                .trim_matches(['(', ')'].as_slice())
                .split_once(", ")
                .with_context(|| format!("Invalid node format: {l}"))?;

            anyhow::Ok((name, paths))
        })
        .try_collect::<HashMap<_, _>>()?;

    Ok((instructions.trim().chars().cycle(), nodes))
}

fn lcd(n: usize, m: usize) -> usize {
    n * (m / gcd(n, m))
}

/// Greatest common divisor using [Stein's algorithm]
///
/// [Stein's algorithm]: https://en.wikipedia.org/wiki/Binary_GCD_algorithm
fn gcd(mut n: usize, mut m: usize) -> usize {
    if m == 0 || n == 0 {
        return m | n;
    }

    // get common factor of 2
    let shift = (m | n).trailing_zeros();

    // shift n and m until odd
    m >>= m.trailing_zeros();
    n >>= n.trailing_zeros();

    // reduce greater value until equal
    while m != n {
        if n < m {
            core::mem::swap(&mut n, &mut m);
        }
        n -= m;
        n >>= n.trailing_zeros();
    }
    m << shift
}

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

    #[test]
    fn test_part_1() -> anyhow::Result<()> {
        const INPUT: &str = indoc::indoc! {"
            LLR

            AAA = (BBB, BBB)
            BBB = (AAA, ZZZ)
            ZZZ = (ZZZ, ZZZ)
        "};

        Ok(assert_eq!(6, Day08::part_1(INPUT)?))
    }

    #[test]
    #[ignore]
    fn test_part_2() -> anyhow::Result<()> {
        const INPUT: &str = indoc::indoc! {"
            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)
        "};

        Ok(assert_eq!(6, Day08::part_2(INPUT)?))
    }
}