from functools import cache
from operator import add, sub, mul, truediv, eq
from subprocess import run, PIPE
from utils import open_day
from z3 import Int, ArithRef, solve

monkeys = {}
with open_day(21) as f:
    for line in f:
        target, expr = line.rstrip().split(': ')
        expr = expr.split(' ')
        if len(expr) == 1:
            expr = int(expr[0])
        monkeys[target] = expr

ops = { '+': add, '-': sub, '*': mul, '/': truediv, '==': eq }

def eval_monkey(m):
    expr = monkeys[m]
    if isinstance(expr, int) or isinstance(expr, ArithRef):
        return expr
    return ops[expr[1]](eval_monkey(expr[0]), eval_monkey(expr[2]))

print(int(eval_monkey('root')))
monkeys['humn'] = Int('humn')
monkeys['root'][1] = '=='
solve(eval_monkey('root'))