mod utils;
extern crate nom;
use nom::{
    branch::alt,
    bytes::complete::{tag, take_until},
    character::complete::{digit0, u32},
    multi::separated_list0,
    IResult,
};
use utils::{max_n, read_input};

fn main() {
    let input = read_input();

    let (_, monkeys_items) = separated_list0(tag("\n\n"), Monkey::parse)(&input).unwrap();
    let (monkeys, items): (Vec<_>, Vec<_>) = monkeys_items.into_iter().unzip();

    println!(
        "Part 1: {}",
        monkey_business::<3, 20>(&monkeys, items.clone(), usize::MAX),
    );
    println!(
        "Part 2: {}",
        monkey_business::<1, 10000>(&monkeys, items, monkeys.iter().map(|x| x.test).product())
    );
}

fn monkey_business<const REDUCER: usize, const ROUNDS: usize>(
    monkeys: &[Monkey],
    mut items: Vec<Vec<usize>>,
    additional_reducer: usize,
) -> usize {
    let mut inspections = vec![0; monkeys.len()];
    for _ in 0..ROUNDS {
        for monkey_idx in 0..monkeys.len() {
            let curr_monkey = &monkeys[monkey_idx];
            while let Some(worry) = items[monkey_idx].pop() {
                let new_worry = (curr_monkey.do_op(worry) / REDUCER) % additional_reducer;
                inspections[monkey_idx] += 1;
                let throw_to = curr_monkey.throw_to(new_worry);

                items[throw_to].push(new_worry);
            }
        }
    }

    max_n(&mut inspections, 2).iter().product::<usize>()
}

#[derive(Debug)]
struct Monkey {
    op: Operation,
    test: usize,
    next_monkey: (usize, usize),
}

impl Monkey {
    fn parse(i: &str) -> IResult<&str, (Monkey, Vec<usize>)> {
        let (i, _) = tag("Monkey ")(i)?;
        let (i, _) = digit0(i)?;
        let (i, _) = tag(":\n  Starting items: ")(i)?;
        let (i, items) = separated_list0(tag(", "), u32)(i)?;
        let (i, _) = tag("\n  Operation: new = ")(i)?;
        let (i, op) = Operation::parse(i)?;
        let (i, _) = take_until("\n")(i)?;
        let (i, _) = tag("\n  Test: divisible by ")(i)?;
        let (i, test) = u32(i)?;
        let (i, _) = tag("\n    If true: throw to monkey ")(i)?;
        let (i, monkey_t) = u32(i)?;
        let (i, _) = tag("\n    If false: throw to monkey ")(i)?;
        let (i, monkey_f) = u32(i)?;

        Ok((
            i,
            (
                Monkey {
                    op,
                    test: test as usize,
                    next_monkey: (monkey_t as usize, monkey_f as usize),
                },
                items.into_iter().map(|x| x as usize).collect(),
            ),
        ))
    }

    fn do_op(&self, x: usize) -> usize {
        match self.op {
            Operation::Square => x * x,
            Operation::Add(c) => x + c,
            Operation::Mult(c) => x * c,
        }
    }
    fn throw_to(&self, x: usize) -> usize {
        if (x % self.test) == 0 {
            self.next_monkey.0
        } else {
            self.next_monkey.1
        }
    }
}

#[derive(Debug)]
enum Operation {
    Square,
    Add(usize),
    Mult(usize),
}

impl Operation {
    fn parse(i: &str) -> IResult<&str, Operation> {
        let (i, _) = tag("old ")(i)?;
        let (i, op) = alt((tag("* "), tag("+ ")))(i)?;
        if let Ok((i, _)) = tag::<&str, &str, (&str, nom::error::ErrorKind)>("old")(i) {
            Ok((i, Operation::Square))
        } else {
            let (i, x) = u32(i)?;
            Ok((
                i,
                match op {
                    "* " => Operation::Mult(x as usize),
                    "+ " => Operation::Add(x as usize),
                    _ => unreachable!(),
                },
            ))
        }
    }
}