use std::collections::HashMap;

use itertools::Itertools;

pub fn process_part1(input: &str) -> u64 {
    let stones = input
        .split_whitespace()
        .map(|stone| stone.parse::<u64>().unwrap())
        .collect_vec();
    let mut result = stones;
    for _ in 0..25 {
        result = blink(result);
    }
    let mut test = HashMap::new();
    for stone in result.iter() {
        test.entry(stone)
            .and_modify(|count| *count += 1)
            .or_insert(1);
    }
    result.len() as u64
}

pub fn process_part2(input: &str, blinks: u64) -> u64 {
    let mut stones: HashMap<u64, u64> = HashMap::new();
    input.split_whitespace().for_each(|stone| {
        let stone = stone.parse::<u64>().unwrap();
        stones
            .entry(stone)
            .and_modify(|count| *count += 1)
            .or_insert(1);
    });
    let mut stones_after: HashMap<u64, u64> = HashMap::new();
    for _idx in 0..blinks {
        for (stone, count) in stones.iter() {
            if *stone == 0 {
                stones_after
                    .entry(1)
                    .and_modify(|inner_count| *inner_count += count)
                    .or_insert(*count);
            } else if (((*stone as f64).log10() + 1.0).floor() as u64) % 2 == 0 {
                let num_digits = ((*stone as f64).log10() + 1.0).floor() as u64;
                let digits_per_side = num_digits / 2;
                let left = stone / 10_u64.pow(digits_per_side as u32);
                let right = stone % 10_u64.pow(digits_per_side as u32);
                stones_after
                    .entry(left)
                    .and_modify(|inner_count| *inner_count += count)
                    .or_insert(*count);
                stones_after
                    .entry(right)
                    .and_modify(|inner_count| *inner_count += count)
                    .or_insert(*count);
                //stones_after.push(left);
                //stones_after.push(right);
            } else {
                let new_stone = stone * 2024;
                stones_after
                    .entry(new_stone)
                    .and_modify(|inner_count| *inner_count += count)
                    .or_insert(*count);
                //stones_after.push(stone * 2024);
            }
        }
        stones = stones_after.clone();
        stones_after.clear();
    }
    stones
        .into_values()
        .reduce(|acc, count| acc + count)
        .unwrap()
}

fn blink(stones: Vec<u64>) -> Vec<u64> {
    let mut stones_after = Vec::new();
    for stone in stones {
        if stone == 0 {
            stones_after.push(1);
            continue;
        }
        let num_digits = ((stone as f64).log10() + 1.0).floor() as u64;
        if num_digits % 2 == 0 {
            let digits_per_side = num_digits / 2;
            let left = stone / 10_u64.pow(digits_per_side as u32);
            let right = stone % 10_u64.pow(digits_per_side as u32);
            stones_after.push(left);
            stones_after.push(right);
        } else {
            stones_after.push(stone * 2024);
        }
    }
    stones_after
}

#[cfg(test)]
mod tests {
    use itertools::Itertools;

    use super::*;

    const INPUT_1: &str = "0 1 10 99 999";
    const INPUT_2: &str = "125 17";

    #[test]
    fn test_blink_1() {
        let stones = INPUT_1
            .split_whitespace()
            .map(|stone| stone.parse::<u64>().unwrap())
            .collect_vec();
        let stones = blink(stones);
        assert_eq!(stones, vec![1, 2024, 1, 0, 9, 9, 2021976]);
    }

    #[test]
    fn test_blink_2() {
        let stones = INPUT_2
            .split_whitespace()
            .map(|stone| stone.parse::<u64>().unwrap())
            .collect_vec();
        let mut result = stones;
        for _ in 0..6 {
            result = blink(result);
        }
        assert_eq!(
            result,
            vec![
                2097446912, 14168, 4048, 2, 0, 2, 4, 40, 48, 2024, 40, 48, 80, 96, 2, 8, 6, 7, 6,
                0, 3, 2
            ]
        );
    }

    #[test]
    fn part1() {
        let result = process_part1(INPUT_2);
        assert_eq!(result, 55312);
    }

    #[test]
    fn part2_1() {
        let result = process_part2(INPUT_1, 1);
        assert_eq!(result, 7);
    }

    #[test]
    fn part2_2() {
        let result = process_part2(INPUT_2, 6);
        assert_eq!(result, 22);
    }

    #[test]
    fn part2_3() {
        let result = process_part2(INPUT_2, 25);
        assert_eq!(result, 55312);
    }
}