# -*- coding: utf-8 -*-
import dataclasses
from abc import ABC
from functools import reduce
from typing import List, Iterator, Any, Tuple

from aoc import BaseAssignment


@dataclasses.dataclass
class Monkey:
    items: List[int]
    operation: List[str]
    test_divisible: int
    test_true: int
    test_false: int

    inspections: int = 0

    @classmethod
    def from_input(cls, lines: List[str]):
        _, items = lines[1].split(":")

        return cls(
            items=[int(item) for item in items.split(",")],
            operation=lines[2].split(" ")[-2:],
            test_divisible=int(lines[3].split(" ")[-1]),
            test_true=int(lines[4].split(" ")[-1]),
            test_false=int(lines[5].split(" ")[-1]),
        )

    def new_worry_factor(self, item: int, worry_factor=1) -> int:
        operator, value = self.operation
        value = item if value == "old" else int(value)

        match operator:
            case "+":
                return item + value
            case "*":
                return item * value

    def throw_to(self, new_worry_factor: int) -> int:
        if new_worry_factor % self.test_divisible == 0:
            return self.test_true
        return self.test_false

    def inspect_items(self, worry_divisor: int) -> Iterator[Tuple[int, int]]:
        for item in self.items:
            self.inspections += 1
            new_worry_factor = self.new_worry_factor(item) // worry_divisor
            yield new_worry_factor, self.throw_to(new_worry_factor)

        self.items = []


class Assignment(BaseAssignment, ABC):
    def parse_monkeys(self, input: Iterator[str]):
        monkeys = []
        monkey = []
        for line in input:
            if line == "":
                monkeys.append(Monkey.from_input(monkey))
                monkey = []
                continue

            monkey.append(line)

        if len(monkey) != 0:
            monkeys.append(Monkey.from_input(monkey))

        return monkeys

    def calculate_monkey_business(
        self,
        monkeys: list[Monkey],
        rounds: int,
        worry_divisor: int = 1,
        common_devisor: int = None,
    ):

        for _ in range(rounds):
            for monkey in monkeys:
                for item, other_monkey in monkey.inspect_items(worry_divisor):
                    monkeys[other_monkey].items.append(
                        item % common_devisor if common_devisor is not None else item
                    )

        a, b = sorted([monkey.inspections for monkey in monkeys])[-2:]

        return a * b


class AssignmentOne(Assignment):
    example_result = 10605

    def run(self, input: Iterator) -> Any:
        monkeys = self.parse_monkeys(input)
        return self.calculate_monkey_business(monkeys, 20, worry_divisor=3)


class AssignmentTwo(Assignment):
    example_result = 2713310158

    def run(self, input: Iterator) -> Any:
        monkeys = self.parse_monkeys(input)
        common_divisor = reduce(lambda c, m: c * m.test_divisible, monkeys, 1)
        return self.calculate_monkey_business(
            monkeys, 10000, common_devisor=common_divisor
        )