use itertools::Itertools;

fn decompress(input: &str, recursive: bool) -> usize {
    let mut stack = input
        .chars()
        .filter(|x| !x.is_whitespace())
        .rev()
        .collect_vec();

    let mut drained = vec![];
    let mut count = 0;

    while let Some(c) = stack.pop() {
        if c == '(' {
            let mut n = 0;
            let sublen = loop {
                let l = stack.pop().unwrap();
                if l == 'x' {
                    break n;
                }
                n = n * 10 + l.to_digit(10).unwrap() as usize;
            };

            let mut n = 0;
            let repeat = loop {
                let l = stack.pop().unwrap();
                if l == ')' {
                    break n;
                }
                n = n * 10 + l.to_digit(10).unwrap() as usize;
            };

            drained.clear();

            for _ in 0..sublen {
                drained.push(stack.pop().unwrap());
            }

            if recursive {
                for _ in 0..repeat {
                    for d in drained.iter().copied().rev() {
                        stack.push(d);
                    }
                }
            } else {
                count += drained.len() * repeat
            }
        } else {
            count += 1
        }
    }

    count
}

pub fn part_one(input: &str) -> Option<usize> {
    Some(decompress(input, false))
}
pub fn part_two(input: &str) -> Option<usize> {
    Some(decompress(input, true))
}

aoc::solution!(9);

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_part_one() {
        let input = aoc::template::read_file("examples", 9);
        assert_eq!(part_one(&input), Some(238));
    }
    #[test]
    fn test_part_two() {
        let input = aoc::template::read_file("examples", 9);
        assert_eq!(part_two(&input), Some(445));
    }
}