from lib import *

EXAMPLE = """root: pppw + sjmn
dbpl: 5
cczh: sllz + lgvd
zczc: 2
ptdq: humn - dvpt
dvpt: 3
lfqf: 4
humn: 5
ljgn: 2
sjmn: drzm * dbpl
sllz: 4
pppw: cczh / lfqf
lgvd: ljgn * ptdq
drzm: hmdt - zczc
hmdt: 32
"""

def resolve(monkeys, s):
    equ = monkeys[s]
    try:
        return int(equ)
    except ValueError:
        pass
    l, op, r = equ.split(" ")
    l = resolve(monkeys, l)
    r = resolve(monkeys, r)
    res = eval(f"{l}{op}{r}")
    monkeys[s] = res
    return res


def resolve2(monkeys, s):
    if s == "humn":
        return tuple()

    equ = monkeys[s]
    try:
        return int(equ)
    except ValueError:
        pass

    l, op, r = equ.split(" ")
    if s == "root":
        op = "=="

    l = resolve2(monkeys, l)
    r = resolve2(monkeys, r)

    if isinstance(l, tuple) or isinstance(r, tuple):
        return (l, op, r)
    if op == "/":
        op = "//"
    res = eval(f"{l}{op}{r}")
    monkeys[s] = res
    return res


def simplify(equ, exp):
    if equ == tuple():
        return exp
    l, op, r = equ
    if op == '/' and isinstance(r, int):
        return simplify(l, exp * r)
    elif op == '+' and isinstance(l, int):
        return simplify(r, exp - l)
    elif op == '+' and isinstance(r, int):
        return simplify(l, exp - r)
    elif op == '-' and isinstance(l, int):
        return simplify(r, -(exp - l))
    elif op == '-' and isinstance(r, int):
        return simplify(l, exp + r)
    elif op == '*' and isinstance(r, int):
        return simplify(l, exp // r)
    elif op == '*' and isinstance(l, int):
        return simplify(r, exp // l)


def solve(input: Input, second=False):
    res = 0
    monkeys = {}
    for line in input.lines():
        l, r = line.split(": ")
        monkeys[l] = r

    if not second:
        return int(resolve(monkeys, "root"))
    else:
        equ = resolve2(monkeys, "root")
        assert equ[1] == '=='
        return simplify(equ[0], equ[2])


def main():
    DAY_INPUT = "i21.txt"
    print("Example  1:", solve(Input(EXAMPLE)))
    print("Solution 1:", solve(Input(DAY_INPUT)))
    print("Example  2:", solve(Input(EXAMPLE), True))
    print("Solution 2:", solve(Input(DAY_INPUT), True))
    assert solve(Input(DAY_INPUT), True) == 3243420789721


if __name__ == "__main__":
    main()