# -*- coding: utf-8 -*-
import dataclasses
from abc import ABC
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
    worry_divider: int = 1

    inspections: int = 0

    @classmethod
    def from_input(cls, lines: List[str], worry_divider: int):
        _, 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]),
            worry_divider=worry_divider,
        )

    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) -> Iterator[Tuple[int, int]]:
        for item in self.items:
            self.inspections += 1
            new_worry_factor = self.new_worry_factor(item) // self.worry_divider
            yield new_worry_factor, self.throw_to(new_worry_factor)

        self.items = []


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

            monkey.append(line)

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

        return monkeys

    def calculate_monkey_business(
        self, input: Iterator[str], rounds: int, worry_divider: int = 1
    ):
        monkeys = self.parse_monkeys(input, worry_divider)

        for _ in range(rounds):
            for monkey in monkeys:
                for item, other_monkey in monkey.inspect_items():
                    monkeys[other_monkey].items.append(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:
        return self.calculate_monkey_business(input, 20, 3)


class AssignmentTwo(Assignment):
    example_result = 2713310158

    def run(self, input: Iterator) -> Any:
        return self.calculate_monkey_business(input, 10000)