module Main where

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

data Stmt = Mask String | Mem Int Int deriving Show

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

parseNumber :: Parser Int
parseNumber = read <$> many1 digit

parseMem :: Parser Stmt
parseMem = Mem <$ string "mem[" <*> parseNumber <* string "] = " <*> parseNumber

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