module Main (main) where

import Lib.Parser (Parser (parse), separated, char, parens, string, number, satisfy)
import Control.Applicative (Alternative (..))
import Data.Maybe (fromJust, catMaybes)

-- | Part 1
data MulOper = MulOper Int Int deriving (Show, Eq)

perform :: MulOper -> Int
perform (MulOper a b) = a * b

mulOper :: Parser MulOper
mulOper = (\_ [a, b] -> MulOper a b) <$> string "mul" <*> parens (char ',' `separated` number)

manyMaybeMulOper :: Parser [Maybe MulOper]
manyMaybeMulOper = many $ Just <$> mulOper <|> Nothing <$ satisfy (const True)

extractMulOps :: Maybe ([Maybe MulOper], String) -> [MulOper]
extractMulOps = catMaybes . fst . fromJust

part1 :: String -> IO ()
part1 s = print $ sum $ map perform $ extractMulOps $ parse manyMaybeMulOper s

-- | Part 2
data Token
  = EnableAction
  | DisableAction
  | MulAction MulOper
  deriving (Show, Eq)

doOper :: Parser Token
doOper = EnableAction <$ string "do()"

dontOper :: Parser Token
dontOper = DisableAction <$ string "don't()"

wrapMulOper :: Parser Token
wrapMulOper = MulAction <$> mulOper

manyMaybeToken :: Parser [Maybe Token]
manyMaybeToken = many $ Just <$> (doOper <|> dontOper <|> wrapMulOper) <|> Nothing <$ satisfy (const True)

runCommands :: [Token] -> Int
runCommands tokens = runCommands' tokens True 0
  where
    runCommands' :: [Token] -> Bool -> Int -> Int
    runCommands' [] _ s = s
    runCommands' ((MulAction _):cs) False s = s + runCommands' cs False s
    runCommands' ((MulAction (MulOper a b)):cs) True s = s + a * b + runCommands' cs True s
    runCommands' (EnableAction:cs) _ s = s + runCommands' cs True s
    runCommands' (DisableAction:cs) _ s = s + runCommands' cs False s

extractTokens :: Maybe ([Maybe Token], String) -> [Token]
extractTokens = catMaybes . fst . fromJust

part2 :: String -> IO()
part2 s = print $ runCommands $ extractTokens $ parse manyMaybeToken s

-- | Main
main :: IO ()
main = readFile "resources/day3.txt" >>= part2