use aoc_runner_derive::{aoc, aoc_generator};

type Input = Vec<usize>;

#[aoc_generator(day9)]
fn parse(input: &str) -> Input {
    input
        .chars()
        .map(|c| c.to_digit(10).unwrap() as usize)
        .collect()
}

#[aoc(day9, part1)]
fn part1(input: &Input) -> usize {
    let mut sum = 0;
    let mut len = input.chunks(2).map(|c| c[0]).sum::<usize>();
    let mut l_idx = 0;
    let mut l_pos = 0;
    let mut r_idx = input.len() + 1;
    let mut r_left = 0;
    while len > 0 {
        if l_idx % 2 == 0 {
            for _ in 0..input[l_idx] {
                // println!("l {}*{}", l_pos, l_idx / 2);
                // print!("{}", l_idx / 2);
                len -= 1;
                sum += (l_idx / 2) * l_pos;
                l_pos += 1;
                if len == 0 {
                    break;
                }
            }
            l_idx += 1;
        } else {
            for _ in 0..input[l_idx] {
                while r_left == 0 {
                    r_idx -= 2;
                    r_left = input[r_idx];
                }
                // println!("r {}*{}", l_pos, r_idx / 2);
                // print!("{}", r_idx / 2);
                len -= 1;
                sum += (r_idx / 2) * l_pos;
                l_pos += 1;
                r_left -= 1;
                if len == 0 {
                    break;
                }
            }
            l_idx += 1;
        }
    }
    sum
}

#[aoc(day9, part2)]
fn part2(input: &Input) -> usize {
    let mut input = input
        .iter()
        .map(|c| (0, *c, *c))
        .collect::<Vec<(usize, usize, usize)>>();
    for i in 1..input.len() {
        input[i].0 = input[i - 1].0 + input[i - 1].1;
    }
    let mut min_idx = [1; 10];
    // println!("{:?}", &input);
    let mut sum = 0;
    for r_idx in (0..input.len()).step_by(2).rev() {
        let val = r_idx / 2;
        let count = input[r_idx].2;
        // println!("attempt: {:?}", input[r_idx]);
        for l_idx in (min_idx[count]..r_idx).step_by(2) {
            if input[l_idx].1 >= input[r_idx].2 {
                // print!("{:?}", input[l_idx]);
                input[r_idx].0 = input[l_idx].0;
                input[l_idx].1 -= input[r_idx].2;
                input[l_idx].0 += input[r_idx].2;
                input[l_idx - 1].1 += input[r_idx].2;

                for i in count..10 {
                    if min_idx[i] > min_idx[count] {
                        break;
                    } else {
                        while min_idx[i] <= r_idx && input[min_idx[i]].1 < i {
                            min_idx[i] += 2;
                        }
                    }
                }
                // println!(" -> {:?}", input[l_idx]);
                break;
            }
        }
        for i in 0..count {
            // println!("{}*{}", (input[r_idx].0 + i), val);
            sum += val * (input[r_idx].0 + i);
            // input[r_idx].0 += 1;
        }
    }
    sum
}

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

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse("2333133121414131402")), 1928);
        assert_eq!(
            part1(&parse(include_str!("../input/2024/day9.txt").trim_end())),
            6337921897505
        );
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse("2333133121414131402")), 2858);
        assert_eq!(
            part2(&parse(include_str!("../input/2024/day9.txt").trim_end())),
            6362722604045
        );
    }
}