import Common
import qualified Data.Map as M
import Data.Function.Memoize (Memoizable(memoize))

type Monkey = String

data Operation = Const Int | Add Monkey Monkey | Sub Monkey Monkey | Mult Monkey Monkey | Div Monkey Monkey
  deriving (Show, Eq)

charToOp '+' = Add
charToOp '-' = Sub
charToOp '*' = Mult
charToOp '/' = Div
charToOp c = error $ "unknown op " ++ [c]

parseOp :: String -> Operation
parseOp str
  | length s == 1 = Const $ read str
  | length s == 3 = op (s !! 0) (s !! 2)
  | otherwise = error $ "can't parse op " ++ str
  where
    s = words str
    op = charToOp $ head $ s !! 1

parseLine :: String -> (Monkey, Operation)
parseLine s = (name, op)
  where
    op = parseOp $ trim $ tail opS
    (name, opS) = splitAt 4 s

rootMonkey :: Monkey
rootMonkey = "root"

eval :: M.Map Monkey Operation -> Int
eval monkeys = evalMonkey' rootMonkey
  where
    evalMonkey' = memoize evalMonkey

    evalMonkey :: String -> Int
    evalMonkey m = case op of
      Const n -> n
      Add a b -> evalMonkey' a + evalMonkey' b
      Sub a b -> evalMonkey' a - evalMonkey' b
      Mult a b -> evalMonkey' a * evalMonkey' b
      Div a b -> evalMonkey' a `div` evalMonkey' b
      where
        op = (M.!) monkeys m

main = interact $
  show
  . eval
  . M.fromList . map parseLine . lines