import Aoc
import Data.Char ( isDigit )
import Control.Arrow ( first, second )

data JSON a = Int a | Str String | Array [JSON a] | Object [(String, JSON a)]
    deriving ( Show, Eq )

parseElements :: (String -> (p, String)) -> Char -> String -> ([p], String)
parseElements parseElement closing cs
    | head cs == closing = ([], tail cs)
    | head cs == ',' = parseElements parseElement closing (tail cs)
    | otherwise =
        let (element, rest0) = parseElement cs in
        let (elements, rest1) = parseElements parseElement closing rest0 in
        (element:elements, rest1)

parseString :: String -> (String, String)
parseString rest0 =
    let (string, rest1) = span (/= '"') rest0 in -- no escaping
    let ('"':rest2) = rest1 in
    (string, rest2)

parseKeyValue :: Read a => [Char] -> ((String, JSON a), String)
parseKeyValue ('"':rest0) =
    let (string, rest1) = parseString rest0 in
    let (':':rest2) = rest1 in
    let (object, rest3) = parseJSON rest2 in
    ((string, object), rest3)

parseJSON :: Read a => String -> (JSON a, String)
parseJSON ('[':rest0) =
    let (elements, rest1) = parseElements parseJSON ']' rest0 in
    (Array elements, rest1)
parseJSON ('{':rest0) =
    let (elements, rest1) = parseElements parseKeyValue '}' rest0 in
    (Object elements, rest1)
parseJSON ('"':rest0) = first Str $ parseString rest0
parseJSON rest = if isNum (head rest) then first (Int . read) $ span isNum rest else error rest where isNum c = isDigit c || c == '-'

parseDone :: (p, String) -> p
parseDone (p, "") = p

instance Foldable JSON where
  foldMap :: Monoid m => (a -> m) -> JSON a -> m
  foldMap f (Int a) = f a
  foldMap f (Str _) = mempty
  foldMap f (Array xs) = foldMap (foldMap f) xs
  foldMap f (Object kvs) = foldMap (foldMap f . snd) kvs

hasRed :: Eq a => JSON a -> Bool
hasRed (Object kvs) = any ((==Str "red") . snd) kvs
hasRed _ = False

removeRed :: Eq a => JSON a -> JSON a
removeRed json = case json of
    Object kvs -> if hasRed json then Object [] else Object $ map (second removeRed) kvs
    Array xs -> Array $ map removeRed xs
    _ -> json

part1 :: JSON Int -> Int
part1 json = sum $ foldMap (:[]) json

part2 :: JSON Int -> Int
part2 = part1 . removeRed

main :: IO ()
main = aocMain (parseDone . parseJSON . single) part1 part2