from tools.aoc import AOCDay
from typing import Any


def has_humn(monkeys: dict, root: str) -> bool:
    if root == 'humn':
        return True

    if monkeys[root]['op'] is None:
        return False

    opr, _, opl = monkeys[root]['op']
    if opr == 'humn' or opl == 'humn':
        return True

    return has_humn(monkeys, opl) or has_humn(monkeys, opr)


def find_humn(monkeys: dict, root: str, target: int) -> int:
    if root == 'humn':
        return target

    opl, op, opr = monkeys[root]['op']
    if has_humn(monkeys, opl):
        value = get_monkey(monkeys, opr)
        match op:
            case '+':
                return find_humn(monkeys, opl, target - value)
            case '-':
                return find_humn(monkeys, opl, target + value)
            case '*':
                return find_humn(monkeys, opl, target // value)
            case '/':
                return find_humn(monkeys, opl, target * value)
    else:
        value = get_monkey(monkeys, opl)
        match op:
            case '+':
                return find_humn(monkeys, opr, target - value)
            case '-':
                return find_humn(monkeys, opr, value - target)
            case '*':
                return find_humn(monkeys, opr, target // value)
            case '/':
                return find_humn(monkeys, opr, target // value)


def get_monkey(monkeys: dict, monkey: str, finalize: bool = False) -> int | None:
    if monkeys[monkey]['op'] is None:
        return monkeys[monkey]['value']

    opl, op, opr = monkeys[monkey]['op']
    opl, opr = get_monkey(monkeys, opl, finalize), get_monkey(monkeys, opr, finalize)
    if opl is None or opr is None:
        return None
    match op:
        case "+":
            monkeys[monkey]['value'] = opl + opr
        case "-":
            monkeys[monkey]['value'] = opl - opr
        case "*":
            monkeys[monkey]['value'] = opl * opr
        case "/":
            monkeys[monkey]['value'] = opl // opr
        case _:
            print("Unknown op:", op)

    if finalize:
        monkeys[monkey]['op'] = None

    return monkeys[monkey]['value']


class Day(AOCDay):
    inputs = [
        [
            (152, "input21_test"),
            (21120928600114, "input21_dennis"),
            (41857219607906, "input21"),
        ],
        [
            (301, "input21_test"),
            (3453748220116, "input21_dennis"),
            (3916936880448, "input21"),
        ]
    ]

    def get_monkey_dict(self) -> dict:
        monkeys = {}
        for line in self.getInput():
            a, b = line.split(": ")
            try:
                monkeys[a] = {
                    'op': None,
                    'value': int(b)
                }
            except ValueError:
                monkeys[a] = {
                    'op': b.split(" "),
                    'value': None
                }

        return monkeys

    def part1(self) -> Any:
        return get_monkey(self.get_monkey_dict(), "root")

    def part2(self) -> Any:
        monkeys = self.get_monkey_dict()
        opl, _, opr = monkeys['root']['op']
        if has_humn(monkeys, opl):
            return find_humn(monkeys, opl, get_monkey(monkeys, opr))
        else:
            return find_humn(monkeys, opr, get_monkey(monkeys, opl))


if __name__ == '__main__':
    day = Day(2022, 21)
    day.run(verbose=True)