use std::collections::HashMap;

fn blink(stone: u64) -> Vec<u64> {
    let n_s = stone.to_string();
    if stone == 0 {
        vec![1]
    } else if n_s.len() % 2 == 0 {
        let mut l = n_s;
        let r = l.split_off(l.len() / 2);
        vec![l.parse().unwrap(), r.parse().unwrap()]
    } else {
        vec![stone * 2024]
    }
}

#[derive(Debug)]
struct VecStones(Vec<u64>);

impl VecStones {
    fn parse(input: &str) -> Self {
        VecStones(
            input
                .trim()
                .split(" ")
                .map(|n| n.parse().unwrap())
                .collect(),
        )
    }

    fn step(&self) -> Self {
        let VecStones(ref s) = self;
        VecStones(s.iter().flat_map(|&n| blink(n)).collect())
    }

    fn count(mut self, iterations: u8) -> u64 {
        for _i in 0..iterations {
            self = self.step();
        }

        self.0.len() as u64
    }

    fn count_cached_v(&self, iterations: u8, cache: &mut Cache) -> u64 {
        self.0
            .iter()
            .map(|&stone| count_cached_s(stone, iterations, cache))
            .sum()
    }

    fn count_cached(&self, iterations: u8) -> u64 {
        let mut cache: Cache = HashMap::new();
        self.count_cached_v(iterations, &mut cache)
    }
}

fn p1(input: &str) -> String {
    let stones = VecStones::parse(input);
    println!("VecStones: {stones:?}");

    let result = stones.count(25);
    result.to_string()
}

// (stone, iteration) -> count
type Cache = HashMap<(u64, u8), u64>;

fn count_cached_s(stone: u64, iterations: u8, cache: &mut Cache) -> u64 {
    if iterations == 0 {
        return 1;
    }

    match cache.get(&(stone, iterations)) {
        Some(&count) => count,
        None => {
            let nexts = VecStones(blink(stone));
            let count = nexts.count_cached_v(iterations - 1, cache);
            cache.insert((stone, iterations), count);
            count
        }
    }
}

fn p2(input: &str) -> String {
    let stones = VecStones::parse(input);

    let result = stones.count_cached(75);
    result.to_string()
}

fn main() {
    aoc2024::run_day("11", p1, p2);
}