import Aoc
import Data.Char ( isDigit )
import Data.Bits

data Op = AND | OR | SHIFTL | SHIFTR | NOT | CONST
    deriving (Show, Eq)

parseFile :: [String] -> [(String, (Op, String, String))] -- [(out, (op, arg1, arg2))]
parseFile = map (\line ->
    let (expr, ["->", out]) = lastSplit 2 $ words line in
        (out, case expr of
            [arg1, "AND", arg2] -> (AND, arg1, arg2)
            [arg1, "OR", arg2] -> (OR, arg1, arg2)
            [arg1, "LSHIFT", arg2] -> (SHIFTL, arg1, arg2)
            [arg1, "RSHIFT", arg2] -> (SHIFTR, arg1, arg2)
            ["NOT", arg] -> (NOT, arg, "")
            [arg] -> (CONST, arg, "")))

eval :: Op -> Int -> Int -> Int
eval OR = (.|.)
eval AND = (.&.)
eval SHIFTR = shiftR
eval SHIFTL = shiftL
eval CONST = const
eval NOT = \a _-> complement a

eval' :: Op -> Int -> Int -> Int
eval' op arg1 arg2 = (.&. 0xFFFF) $ eval op arg1 arg2

eval'' :: [(String, (Op, String, String))] -> [(String, Int)] -> String -> (Int, [(String, Int)])
eval'' wires memos label = case lookup label memos of
    Just value -> (value, memos)
    Nothing -> case lookup label wires of
        Just (op, arg1, arg2) ->
            let (val1, memos1) = eval'' wires memos arg1 in
            let (val2, memos2) = eval'' wires memos1 arg2 in
            let val = eval' op val1 val2 in
            (val, (label, val):memos2)
        Nothing -> (if all isDigit label then read label else undefined, memos) -- if null label then 0 else 

part1 :: [(String, (Op, String, String))] -> Int
part1 wires = fst $ eval'' wires [] "a"

part2 :: [(String, (Op, String, String))] -> Int -> Int
part2 wires part1Answer = fst $ eval'' wires [("b", part1Answer)] "a"

main :: IO ()
main = aocMain' parseFile (dup . part1) part2