{-# language RecordWildCards #-}

import Control.Arrow ((&&&))
import Control.Foldl (Fold(Fold), fold)
import Data.Bits (complement, zeroBits, shift, setBit, clearBit, (.|.), (.&.))
import Data.Char (isDigit)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM
import Data.Maybe (listToMaybe)

data Mask = MkMask { clearBits :: Int, setBits :: Int } deriving (Eq, Show)
data Write = MkWrite { location :: Int, value :: Int } deriving (Eq, Show)
data Instruction = WriteInstr Write | MaskInstr Mask deriving (Eq, Show)

readMay :: Read a => String -> Maybe a
readMay = fmap fst . listToMaybe . reads

parseMask :: String -> Mask
parseMask s = MkMask{..}
 where
  (clearBits, setBits) = fold ((,) <$> clearFold <*> setFold) s
  clearFold = Fold updClear (complement zeroBits) id
  setFold = Fold updSet zeroBits id
  updClear m c = shiftIn ('0' /= c) m
  updSet m c = shiftIn ('1' == c) m
  shiftIn True = flip setBit 0 . flip shift 1
  shiftIn False = flip shift 1

parseLine :: String -> Maybe Instruction
parseLine ('m':'a':_:_:_:_:_:maskStr) = pure . MaskInstr $ parseMask maskStr
parseLine ('m':'e':_:_:memStr) = WriteInstr <$> do
  location <- readMay locStr
  value <- readMay valStr
  return MkWrite{..}
 where
  (locStr, rest) = span isDigit memStr
  valStr = drop 4 rest
parseLine _ = Nothing

common :: String -> [String]
common = lines

data ProgramState = MkProgState { mask :: Mask, mem :: IntMap Int } deriving (Eq, Show)

initProgState :: ProgramState
initProgState =
  MkProgState
  { mask = MkMask { clearBits = complement zeroBits, setBits = zeroBits }
  , mem = IM.empty
  }

setMask :: Mask -> ProgramState -> ProgramState
setMask m s = s{ mask = m }

writeMem :: Int -> Int -> IntMap Int -> IntMap Int
writeMem loc 0 = IM.update (const Nothing) loc
writeMem loc val = IM.insert loc val

write :: Write -> ProgramState -> ProgramState
write MkWrite{..} s =
  s{ mem = writeMem location (setBits curMask .|. (clearBits curMask .&. value)) (mem s) }
 where curMask = mask s

applyInstruction :: Instruction -> ProgramState -> ProgramState
applyInstruction (WriteInstr w) = write w
applyInstruction (MaskInstr m) = setMask m

runProgram :: [Instruction] -> ProgramState
runProgram = fold (Fold (flip applyInstruction) initProgState id)

part1 :: [String] -> Maybe Int
part1 = fmap (sum . mem . runProgram) . traverse parseLine

data MaskBit = Unchanged | Overwritten | Floating deriving (Eq, Show)
type Mask2 = [MaskBit]
data Instruction2 = WriteInstr2 Write | MaskInstr2 Mask2 deriving (Eq, Show)

parseMaskBit :: Char -> Maybe MaskBit
parseMaskBit '0' = pure Unchanged
parseMaskBit '1' = pure Overwritten
parseMaskBit 'X' = pure Floating
parseMaskBit _ = Nothing

parseMask2 :: String -> Maybe Mask2
parseMask2 = traverse parseMaskBit

parseLine2 :: String -> Maybe Instruction2
parseLine2 ('m':'a':_:_:_:_:_:maskStr) = MaskInstr2 <$> parseMask2 maskStr
parseLine2 ('m':'e':_:_:memStr) = WriteInstr2 <$> do
  location <- readMay locStr
  value <- readMay valStr
  return MkWrite{..}
 where
  (locStr, rest) = span isDigit memStr
  valStr = drop 4 rest
parseLine2 _ = Nothing

data ProgramState2 = MkProgState2 { mask2 :: Mask2, mem2 :: IntMap Int } deriving (Eq, Show)

initProgState2 :: ProgramState2
initProgState2 =
  MkProgState2
  { mask2 = []
  , mem2 = IM.empty
  }

decodeMem :: Mask2 -> Int -> [Int]
decodeMem mbs i = snd $ foldr upd (0,[i]) mbs
 where
  upd Unchanged (n, is) = (succ n, is)
  upd Overwritten (n, is) = (succ n, map (flip setBit n) is)
  upd Floating (n, is) = (succ n, map (flip setBit n) is ++ map (flip clearBit n) is)

setMask2 :: Mask2 -> ProgramState2 -> ProgramState2
setMask2 m s = s{ mask2 = m }

write2 :: Write -> ProgramState2 -> ProgramState2
write2 MkWrite{..} s =
  s{ mem2 = foldr (\l -> writeMem l value) (mem2 s) $ decodeMem (mask2 s) location }

applyInstruction2 :: Instruction2 -> ProgramState2 -> ProgramState2
applyInstruction2 (WriteInstr2 w) = write2 w
applyInstruction2 (MaskInstr2 m) = setMask2 m

runProgram2 :: [Instruction2] -> ProgramState2
runProgram2 = fold (Fold (flip applyInstruction2) initProgState2 id)

part2 :: [String] -> Maybe Int
part2 = fmap (sum . mem2 . runProgram2) . traverse parseLine2

interactive :: Show a => (String -> a) -> IO ()
interactive f = getContents >>= print . f

main :: IO ()
main = interactive ((part1 &&& part2) . common)