import Data.Bits
import Data.Int
import Data.List
import Debug.Trace
import Data.List.Split
import Data.Maybe
import qualified Data.Map.Strict as Map

type State = (Int64, [Int], Map.Map Int64 Int64)

parseMask :: (Int64, [Int]) -> (Int, Char) -> (Int64, [Int])
parseMask (oMask, flucBits) (nth, ch) = 
    let oMask2 = if ch == '1' then oMask .|. (shift 1 nth) else oMask
        flucBits2 = if ch == 'X' then nth : flucBits else flucBits in
        (oMask2, flucBits2)

applyMask :: Bits a => a -> (Int, Int) -> a
applyMask n (i, c) = if c == 0 then clearBit n i else setBit n i

nextState :: State -> String -> State
nextState (oMask, flucBits, intMap) origLine =
    if isPrefixOf "mask" origLine then
        let line = fromJust $ stripPrefix "mask = " origLine in
        let len = length line - 1 in
        let enumLine = zip [len, len-1..0] line in
        let (oMask, flucBits) = foldl parseMask (0, []) $ enumLine in
        (oMask, flucBits, intMap)
    else
        let line = fromJust $ stripPrefix "mem[" origLine in
        let [addr, value] = map read $ splitOn "] = " line in
        let flucIter = sequence $ take (length flucBits) $ repeat [0, 1] in
        let enumBits lines = zip flucBits lines in
        let f = foldl applyMask (addr .|. oMask) . enumBits in
        let addrs = map f flucIter in
        let intMap2 = foldl (\acc el -> Map.insert el value acc) intMap addrs in
        (oMask, flucBits, intMap2)


solve2 :: [String] -> Int64
solve2 = sumElements . foldl nextState (0, [], Map.empty)
    where sumElements (_, _, map) = Map.foldl (+) 0 map

main :: IO ()
main = do
    l <- lines <$> readFile "14.txt"
    print $ solve2 l