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

#[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 From<&u8> for Op {
    fn from(value: &u8) -> Self {
        match value {
            b'+' => Op::Add,
            b'*' => 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, Spaghett)]
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, Spaghett)]
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>()
}

#[aoc_generator(day6, part2, StateMachine)]
fn parse2_walk(input: &str) -> Grid<u8> {
    Grid::from_str(input).unwrap()
}

enum NextPos {
    Advance((usize, usize)),
    NewCol((usize, usize)),
    Done,
}

#[aoc(day6, part2, StateMachine)]
fn part2_walk(g: &Grid<u8>) -> u64 {
    let mut accum = 0;

    // start at the bottom right
    let mut pos = (g.width() - 1, g.height() - 1);
    let mut digits = 0;
    let mut cur_value = 0;
    let mut cur_values = Vec::new();
    let mut cur_op = Op::Add;
    let mut last_col = false;

    let height = g.height() - 1;

    let next_pos = |(x, y)| -> NextPos {
        if y == 0 {
            if x == 0 {
                return NextPos::Done;
            }
            return NextPos::NewCol((x - 1, height));
        }
        NextPos::Advance((x, y - 1))
    };

    while let Some(cur_char) = g.get(&pos) {
        match cur_char {
            b'*' | b'+' => {
                last_col = true;
                cur_op = cur_char.into()
            }
            cur_char if cur_char.is_ascii_digit() => {
                cur_value += (*cur_char - b'0') as u64 * POW10[digits];
                digits += 1;
            }
            _ => {}
        }
        pos = match next_pos(pos) {
            NextPos::Advance(p) => p,
            NextPos::NewCol(p) => {
                // add the value
                if cur_value != 0 {
                    cur_values.push(cur_value);
                    cur_value = 0;
                }
                digits = 0;
                // this was a last col, so reduce and accumulate
                if last_col {
                    accum += cur_values
                        .iter()
                        .cloned()
                        .reduce(|l, r| cur_op.f(&l, &r))
                        .unwrap_or(0);
                    // clear the state
                    cur_values.clear();
                    last_col = false;
                }
                p
            }
            NextPos::Done => {
                if cur_value != 0 {
                    cur_values.push(cur_value);
                }
                // accumulate the final result
                accum += cur_values
                    .iter()
                    .cloned()
                    .reduce(|l, r| cur_op.f(&l, &r))
                    .unwrap_or(0);
                return accum;
            }
        }
    }

    unreachable!()
}

#[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);
    }

    #[test]
    fn part2_walk_example() {
        assert_eq!(part2_walk(&parse2_walk(EXAMPLE)), 3263827);
    }
}