use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{line_ending, u64},
    combinator::{map, value},
    multi::separated_list1,
    sequence::{pair, preceded, tuple},
    IResult,
};
use std::fs;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_11.txt")?;
    let troop = MonkeyTroop::parser(&input).unwrap().1;

    {
        let mut troop_1 = troop.clone();
        for _ in 0..20 {
            troop_1.monkey_game_round(true);
        }
        dbg!(&troop_1.monkey_business_score());
    }

    {
        let mut troop_2 = troop.clone();
        for _ in 0..10000 {
            troop_2.monkey_game_round(false);
        }
        dbg!(&troop_2.monkey_business_score());
    }

    Ok(())
}

#[derive(Debug, Clone)]
struct MonkeyTroop {
    monkeys: Vec<Monkey>,
    gcd: u64,
}

#[derive(Debug, Clone)]
struct Monkey {
    items: Vec<u64>,
    operation: Operation,
    test_denominator: u64,
    true_target: usize,
    false_target: usize,
    inspection_count: usize,
}

#[derive(Debug, Clone)]
enum Operation {
    Add(u64),
    Multiply(u64),
    Square,
}

impl MonkeyTroop {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_list1(pair(line_ending, line_ending), Monkey::parser),
            |monkeys| MonkeyTroop {
                gcd: monkeys.iter().map(|m| m.test_denominator).product(),
                monkeys,
            },
        )(input)
    }

    fn monkey_game_round(&mut self, worry_reduction: bool) {
        for monkey_i in 0..self.monkeys.len() {
            self.monkeys[monkey_i].inspect_items(worry_reduction, self.gcd);

            let items = std::mem::take(&mut self.monkeys[monkey_i].items);
            let test_denominator = self.monkeys[monkey_i].test_denominator.clone();
            let true_target = self.monkeys[monkey_i].true_target.clone();
            let false_target = self.monkeys[monkey_i].false_target.clone();
            for item in items {
                let target_monkey = if item % test_denominator == 0 {
                    true_target
                } else {
                    false_target
                };
                self.monkeys[target_monkey].items.push(item);
            }
        }
    }

    fn monkey_business_score(&self) -> usize {
        let mut counts: Vec<usize> = self
            .monkeys
            .iter()
            .map(|m| m.inspection_count.clone())
            .collect();
        counts.sort();
        counts.reverse();
        counts[0] * counts[1]
    }
}

impl Monkey {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            tuple((
                preceded(
                    tuple((
                        tag("Monkey "),
                        u64,
                        tag(":"),
                        line_ending,
                        tag("  Starting items: "),
                    )),
                    separated_list1(tag(", "), u64),
                ),
                preceded(
                    pair(line_ending, tag("  Operation: new = old ")),
                    alt((
                        map(preceded(tag("+ "), u64), Operation::Add),
                        map(preceded(tag("* "), u64), Operation::Multiply),
                        value(Operation::Square, tag("* old")),
                    )),
                ),
                preceded(pair(line_ending, tag("  Test: divisible by ")), u64),
                preceded(pair(line_ending, tag("    If true: throw to monkey ")), u64),
                preceded(
                    pair(line_ending, tag("    If false: throw to monkey ")),
                    u64,
                ),
            )),
            |(items, operation, test_denominator, true_target, false_target)| Monkey {
                items,
                operation,
                test_denominator,
                true_target: true_target as usize,
                false_target: false_target as usize,
                inspection_count: 0,
            },
        )(input)
    }

    fn inspect_items(&mut self, worry_reduction: bool, gcd: u64) {
        for item in &mut self.items {
            match self.operation {
                Operation::Add(i) => *item += i,
                Operation::Multiply(i) => *item *= i,
                Operation::Square => *item = *item * *item,
            };
            if worry_reduction {
                *item /= 3;
            }
            *item = *item % gcd;
        }
        self.inspection_count += self.items.len()
    }
}