module Day19 where

import AoC

import Control.Applicative
import Data.Char
import Data.Foldable
import Data.Map (Map, (!))
import qualified Data.Map as M
import Data.Maybe
import Data.Tuple
import Text.ParserCombinators.ReadP
import Text.Read.Lex

data RuleElement = Literal String | Reference Int

parseElement :: ReadP RuleElement
parseElement = parseRef <|> parseLit
    where parseRef = Reference <$> readDecP
          parseLit = Literal <$> between (char '"') (char '"') (many1 $ satisfy isAlphaNum)

parseRule :: ReadP (Int, [[RuleElement]])
parseRule = do ruleNum <- readDecP
               string ": "
               alternatives <- parseAlternatives
               return (ruleNum, alternatives)
    where parseAlternatives = parseSequence `sepBy` string " | "
          parseSequence = parseElement `sepBy` char ' '

buildParser :: Map Int [[RuleElement]] -> Int -> ReadP String
buildParser m i = buildAlternatives $ m ! i
    where buildAlternatives = asum . map buildSequence
          buildSequence = foldl1 (liftA2 (++)) . map buildElement
          buildElement (Literal s) = string s
          buildElement (Reference j) = buildParser m j

countRootMatches :: [String] -> Map Int [[RuleElement]] -> Int
countRootMatches input rules = length . filter (isJust . oneCompleteResult rootParser) $ input
    where rootParser = buildParser rules 0

main = runAoC splitInput run (run . withExtraRules)
    where splitInput = swap . fmap tail . break (=="") . lines
          withExtraRules = fmap (++ ["8: 42 | 42 8", "11: 42 31 | 42 11 31"])
          run = uncurry countRootMatches . fmap parseRules
          parseRules = M.fromList . map (fromJust . oneCompleteResult parseRule)