{-# OPTIONS_GHC -Wno-simplifiable-class-constraints #-}

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 = const' complement

evalM :: [(String, (Op, String, String))] -> MemoF String Int
evalM wires = memoise $ \values label ->
    case lookup label wires of
        Just (op, arg1, arg2) ->
            let 
                (arg1', values') = evalM wires values arg1 
                (arg2', values'') = evalM wires values' arg2
            in
                (eval op arg1' arg2' .&. 0xFFFF, values'')
        Nothing ->
            if all isDigit label 
                then (read label, values)
                else error label

part1 :: [(String, (Op, String, String))] -> Int
part1 wires = unmemo (evalM wires) "a"

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

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