module Main where

import qualified Data.Map as M
import Control.Monad.Tardis
import qualified Data.ByteString.Char8 as B
import Data.Maybe
import Data.Monoid

type Monkeys = M.Map B.ByteString Int

runMonkeys :: [B.ByteString] -> Tardis Monkeys Monkeys Int
runMonkeys [] = getPast >>= return . fromJust . M.lookup (B.pack "root")

runMonkeys (x:xs) = case number of
  Just (number, _) -> do
    modifyBackwards (M.insert name number)
    modifyForwards (M.insert name number)
    runMonkeys xs
  Nothing -> do
    bw <- getFuture
    fw <- getPast
    -- union is REALLY slow
    let lookup k = First (M.lookup k bw) <> First (M.lookup k fw) 
    let v = case lookup lhs of
              First (Just lhs') -> case lookup rhs of
                First (Just rhs') -> fn lhs' rhs'
    modifyBackwards (M.insert name v)
    modifyForwards (M.insert name v)
    runMonkeys xs   
  where
    name = B.take 4 x
    -- readMaybe is slowwww
    number = B.readInt (B.drop 6 x)
    expression = fromJust $ x B.!? 11
    fn = case expression of
      '+' -> (+)
      '*' -> (*)
      '/' -> div
      '-' -> (-)
    lhs = B.take 4 $ B.drop 6 x
    rhs = B.take 4 $ B.drop 13 x



main = B.readFile "input.txt" >>= (print . (`evalTardis` (M.empty, M.empty)) . runMonkeys . B.lines)