use aoc_runner_derive::aoc;
use itertools::Itertools;
use rustc_hash::FxHashMap;
use std::iter::repeat;

type IntType = u64;
type CacheType = FxHashMap<Stone, IntType>;

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
struct Stone(IntType);
struct Stones(Vec<Stone>);

enum BlinkResult {
    One(Stone),
    Two(Stone, Stone),
}

impl From<&str> for Stones {
    fn from(input: &str) -> Self {
        Stones(
            input
                .split_ascii_whitespace()
                .map(|v| Stone(v.parse().unwrap()))
                .collect_vec(),
        )
    }
}

fn parse(input: &str) -> Stones {
    Stones::from(input)
}

impl Stone {
    fn blink_once(&self) -> BlinkResult {
        let n_digits = if self.0 == 0 { 1 } else { self.0.ilog10() + 1 };
        if self.0 == 0 {
            BlinkResult::One(Stone(1))
        } else if n_digits % 2 == 0 {
            let split_factor = (10 as IntType).pow(n_digits / 2);
            let parts = (self.0 / split_factor, self.0 % split_factor);
            BlinkResult::Two(Stone(parts.0), Stone(parts.1))
        } else {
            BlinkResult::One(Stone(&self.0 * 2024))
        }
    }
}

fn count_blinks(stone: &Stone, blink: usize, cache: &mut Vec<CacheType>) -> IntType {
    if cache[blink].contains_key(stone) {
        return cache[blink][stone];
    }
    let stones = stone.blink_once();
    let result = if blink == 0 {
        match stones {
            BlinkResult::One(_) => 1,
            BlinkResult::Two(_, _) => 2,
        }
    } else {
        match stones {
            BlinkResult::One(s) => count_blinks(&s, blink - 1, cache),
            BlinkResult::Two(s1, s2) => count_blinks(&s1, blink - 1, cache) + count_blinks(&s2, blink - 1, cache),
        }
    };
    cache[blink].insert(stone.clone(), result);
    cache[blink][stone]
}

fn blink_stones(stones: Stones, blinks: usize) -> IntType {
    let mut cache = Vec::from_iter(repeat(CacheType::default()).take(blinks));
    stones
        .0
        .iter()
        .map(|stone| count_blinks(stone, blinks - 1, &mut cache))
        .sum()
}

#[aoc(day11, part1)]
pub fn part1(input: &str) -> IntType {
    let stones = parse(input);
    blink_stones(stones, 25)
}

#[aoc(day11, part2)]
pub fn part2(input: &str) -> IntType {
    let stones = parse(input);
    blink_stones(stones, 75)
}

#[cfg(test)]
mod tests {
    use super::*;
    pub const EXAMPLE: &str = "125 17";

    #[test]
    fn part1_example() {
        assert_eq!(part1(EXAMPLE), 55312);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(EXAMPLE), 65601038650482);
    }
}