use aoc_runner_derive::{aoc, aoc_generator};
use grid::Grid;
use itertools::Itertools;
use std::{
    fmt::{Display, Write},
    iter::repeat_n,
};

#[repr(u8)]
enum Op {
    Add = b'+',
    Mul = b'*',
}

impl From<&str> for Op {
    fn from(value: &str) -> Self {
        match value {
            "+" => Op::Add,
            "*" => Op::Mul,
            c => panic!("Invalid op `{c}`"),
        }
    }
}

impl Display for Op {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Op::Add => f.write_char('+'),
            Op::Mul => f.write_char('*'),
        }
    }
}

impl Op {
    fn f(&self, lhs: u64, rhs: u64) -> u64 {
        match self {
            Op::Add => lhs + rhs,
            Op::Mul => lhs * rhs,
        }
    }
}

#[aoc_generator(day6, part1)]
fn parse(input: &str) -> (Grid<u64>, Vec<Op>) {
    let mut rows = input
        .lines()
        .map(|l| l.split_ascii_whitespace().collect_vec())
        .collect_vec();

    let ops = rows
        .pop()
        .unwrap()
        .iter()
        .map(|op| Op::from(*op))
        .collect_vec();
    let mut grid = Grid::with_shape(rows[0].len(), rows.len(), 0);

    for (y, r) in rows.iter().enumerate() {
        for (x, v) in r.iter().enumerate() {
            grid.set(&(x, y), v.parse().unwrap());
        }
    }

    (grid, ops)
}

#[aoc(day6, part1)]
fn part1((grid, ops): &(Grid<u64>, Vec<Op>)) -> u64 {
    (0..grid.width())
        .map(|x| {
            grid.col_iter(x as i64)
                .unwrap()
                .cloned()
                .reduce(|l, r| ops[x].f(l, r))
                .unwrap()
        })
        .sum()
}

// fn split_digits(x: &u64) -> Vec<u8> {
//     let n_digits = x.ilog10() as usize;
//     (0..=n_digits)
//         .rev()
//         .map(|n| ((x / POW10[n]) % 10) as u8)
//         .collect()
// }

#[aoc_generator(day6, part2)]
fn parse2(input: &str) -> (Grid<Vec<char>>, Vec<Op>) {
    let mut rows = input.lines().collect_vec();
    let ops = rows.pop().unwrap();

    let col_starts = ops
        .chars()
        .enumerate()
        .filter(|(_i, c)| *c == '+' || *c == '*')
        .map(|(i, _c)| i)
        .collect_vec();

    let col_lengths = col_starts
        .iter()
        .tuple_windows()
        .map(|(l, r)| r - l - 1)
        .collect_vec();

    let split_rows = rows
        .iter()
        .map(|r| {
            col_starts
                .iter()
                .zip(col_lengths.iter())
                .map(|(s, l)| &r[*s..s + l])
                .chain(repeat_n(&r[*col_starts.last().unwrap()..], 1))
                .collect_vec()
        })
        .collect_vec();

    let ops = ops.split_ascii_whitespace().map(Op::from).collect_vec();
    let mut grid = Grid::with_shape(split_rows[0].len(), split_rows.len(), Vec::new());
    for (y, r) in split_rows.iter().enumerate() {
        for (x, v) in r.iter().enumerate() {
            grid.set(&(x, y), v.chars().collect_vec());
        }
    }

    (grid, ops)
}

#[aoc(day6, part2)]
fn part2((grid, ops): &(Grid<Vec<char>>, Vec<Op>)) -> u64 {
    let mut columns = Vec::new();
    for col in 0..grid.width() {
        let mut digit = 0;
        let mut col_values = Vec::new();
        loop {
            let val = grid
                .col_iter(col as i64)
                .unwrap()
                .filter_map(|s| s.get(digit))
                .filter(|c| !c.is_ascii_whitespace())
                .map(|v| v.to_digit(10).unwrap() as u64)
                .fold(0, |acc, n| acc * 10 + n);
            if val == 0 {
                columns.push(col_values);
                break;
            }
            col_values.push(val);
            digit += 1;
        }
    }
    columns
        .iter()
        .enumerate()
        .map(|(col, v)| v.iter().cloned().reduce(|l, r| ops[col].f(l, r)).unwrap())
        .sum::<u64>()
}

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

    const EXAMPLE: &str = r"123 328  51 64 
 45 64  387 23 
  6 98  215 314
*   +   *   +  ";

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

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