use std::collections::HashMap;

static DAY: u8 = 11;

fn main() {
    let input = advent::read_lines(DAY);
    println!("{DAY}a: {}", blink_stones(&input[0], 25));
    println!("{DAY}b: {}", blink_stones(&input[0], 75));
}

fn next_stones(number: u64) -> Vec<u64> {
    /* rule 1 */
    if number == 0 {
        return [1].to_vec();
    }

    /* rule 2 */
    let number_str = number.to_string();
    if number_str.len() % 2 == 0 {
        let left = number_str.chars().take(number_str.len() / 2).collect::<String>();
        let right = number_str.chars().skip(number_str.len() / 2).collect::<String>();

        let left = left.parse::<u64>().unwrap();
        let right = right.parse::<u64>().unwrap();

        return [left, right].to_vec();
    }

    /* rule 3 */
    [2024 * number].to_vec()
}

fn count_stones(number: u64, rounds: u32, cache: &mut HashMap<(u64, u32), u64>) -> u64 {
    if let Some(n) = cache.get(&(number, rounds)) {
        return *n;
    }

    if rounds == 0 {
        return 1;
    }

    let mut count = 0;
    for n in next_stones(number) {
        count += count_stones(n, rounds - 1, cache);
        cache.insert((number, rounds), count);
    }
    count
}

fn blink_stones(input: &str, rounds: u32) -> u64 {
    let mut cache = HashMap::new();
    input.split(" ")
         .map(|val| val.parse::<u64>().unwrap())
         .map(|val| count_stones(val, rounds, &mut cache))
         .sum()
}

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

    #[test]
    fn test() {
        let input = "125 17";
        assert_eq!(blink_stones(input, 25), 55312);
    }
}