module Main where

import           Utils
import           Text.ParserCombinators.Parsec
import           Text.Parsec.Char
import           Data.Either
import           Data.Char
import           Data.Strings
import           Numeric (showIntAtBase, readInt)
import           Data.Map (Map)
import qualified Data.Map as Map

data Stmt = Mask String | Mem Int Int deriving Show

parseMask :: Parser Stmt
parseMask = string "mask = " >> Mask <$> many anyChar

parseMem :: Parser Stmt
parseMem = do 
    one <- string "mem[" >> many1 digit
    two <- string "] = " >> many1 digit
    return $ Mem (read one) (read two)

parseLine :: Parser Stmt
parseLine = try parseMask <|> parseMem

applyMask :: Int -> String -> Int
applyMask v m = fst $ head $ readInt 2 (`elem` "01") digitToInt wm
    where bv = strPadLeft '0' 36 $ showIntAtBase 2 intToDigit v ""
          wm = zipWith fn bv m
          fn o 'X' = o
          fn _ '1' = '1'
          fn _ '0' = '0'

runProgram :: [Stmt] -> Int
runProgram ls = sum regs
    where (mask, regs) = foldl fn ("", Map.empty) ls
          fn (_, regs) (Mask s) = (s, regs)
          fn (m, regs) (Mem idx val) = (m, Map.insert idx nval regs)
              where nval = applyMask val m

floatings :: String -> [String]
floatings [] = [[]]
floatings ('X':xs) = floatings xs >>= (\b -> ['0':b, '1':b])
floatings (x:xs) = map (x:) $ floatings xs

genIdxs :: Int -> String -> [Int]
genIdxs v m = map (fst . head . readInt 2 (`elem` "01") digitToInt) (floatings wm)
    where bv = strPadLeft '0' 36 $ showIntAtBase 2 intToDigit v ""
          wm = zipWith fn bv m
          fn o '0' = o
          fn _ '1' = '1'
          fn _ 'X' = 'X'

v2chip :: [Stmt] -> Int
v2chip ls = sum regs
    where (mask, regs) = foldl fn ("", Map.empty) ls
          fn (_, regs) (Mask s) = (s, regs)
          fn (m, regs) (Mem idx val) = (m, nmap)
              where idxs = genIdxs idx m
                    nmap = flip Map.union regs $ Map.fromList $ zip idxs $ repeat val


main :: IO ()
main = do
    n <- rights . map (parse parseLine "main") . lines <$> readFile "input/14"
    print $ runProgram n
    print $ v2chip n