#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [attoparsec multiset containers pretty-simple])"
#! nix-shell -i "ghcid -c 'ghci -Wall' -T main"

{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative
import Control.Monad.Zip
import Data.Attoparsec.Text (Parser, parseOnly)
import Data.Bool (bool)
import Data.Map (Map)
import Data.Monoid (Sum(..))
import Data.MultiSet (MultiSet)
import Data.Tree (Tree)
import Text.Pretty.Simple
-- import Debug.Trace (trace)
import qualified Data.Attoparsec.Text as A
import qualified Data.Map.Strict as M
import qualified Data.MultiSet as B
import qualified Data.Text as Text
import qualified Data.Tree as T

exampleData :: [String]
exampleData =
  [ "light red bags contain 1 bright white bag, 2 muted yellow bags."
  , "dark orange bags contain 3 bright white bags, 4 muted yellow bags."
  , "bright white bags contain 1 shiny gold bag."
  , "muted yellow bags contain 2 shiny gold bags, 9 faded blue bags."
  , "shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags."
  , "dark olive bags contain 3 faded blue bags, 4 dotted black bags."
  , "vibrant plum bags contain 5 faded blue bags, 6 dotted black bags."
  , "faded blue bags contain no other bags."
  , "dotted black bags contain no other bags."
  ]

exampleData2 :: [String]
exampleData2 =
  [ "shiny gold bags contain 2 dark red bags."
  , "dark red bags contain 2 dark orange bags."
  , "dark orange bags contain 2 dark yellow bags."
  , "dark yellow bags contain 2 dark green bags."
  , "dark green bags contain 2 dark blue bags."
  , "dark blue bags contain 2 dark violet bags."
  , "dark violet bags contain no other bags."
  ]


data Color = Color String String
  deriving (Show, Ord, Eq)

-- An efficient implementation of multisets, also sometimes called bags.
type Bag = MultiSet Color

parseFullColor :: Parser Color
parseFullColor = do
  adjective <- A.many' A.letter
  A.skipSpace
  colorName <- A.many' A.letter
  A.skipSpace
  _ <- ("bags" <|> "bag")
  pure $ Color adjective colorName

parseBag' :: Parser Bag
parseBag' = do
  i <- A.decimal
  A.skipSpace
  c <- parseFullColor
  pure $ B.insertMany c i B.empty

parseNoBag :: Parser Bag
parseNoBag = do
  _ <- "no other bags"
  pure $ B.empty

parseBag :: Parser Bag
parseBag = parseNoBag <|> parseBag'

type Rules = Map Color Bag

parseRule :: Parser Rules
parseRule = do
  k <- parseFullColor
  A.skipSpace
  _ <- "contain"
  A.skipSpace
  vx <- parseBag `A.sepBy` (", ")
  _ <- "."
  let v = B.unions vx
  pure $ M.insert k v M.empty


eitherToMaybe :: Either b a -> Maybe a
eitherToMaybe (Right v) = Just v
eitherToMaybe (Left _) = Nothing

parseStrs :: [String] -> Maybe Rules
parseStrs s = M.unions <$> parse s
 where
  parse = eitherToMaybe
        . sequence
        . map (parseOnly (parseRule <* A.endOfInput) . Text.pack)

mkTree :: (Color,Int) -> Rules -> Tree (Color, Int)
mkTree color rules = T.unfoldTree mkNode color
  where
    mkNode :: (Color, Int) -> ((Color, Int), [(Color, Int)])
    mkNode n@(c,_) = (n, concat $ B.toOccurList <$> M.lookup c rules)

mkSimpleTree :: Color -> Rules -> Tree Color
mkSimpleTree color rules = T.unfoldTree mkNode color
  where
    mkNode :: Color -> (Color, [Color])
    mkNode c = (c, fmap fst <$> concat $ B.toOccurList <$> M.lookup c rules)

solvePart1 :: Color -> [String] -> Maybe Int
solvePart1 color strs = do
  rules <- parseStrs strs
  let trees = map (\c -> mkSimpleTree c rules) (filter (/= color) $ M.keys rules)
  let containColor = map (elem color) trees
  pure $ getSum $ foldMap (Sum . bool 0 1) containColor
  -- ^^ complicated exp for: pure $ length $ filter (== True) containColor

updateTree :: Tree Int -> Tree Int
updateTree t@(T.Node _ []) = t
updateTree (T.Node v vx) = T.Node v (fmap ((fmap ((*) v)) . updateTree) vx)

solvePart2 :: Color -> [String] -> Maybe ( Sum Int )
solvePart2 color strs = do
  rules <- parseStrs strs
  let tree = mkTree (color,1) rules
  pure $ foldMap Sum $ (updateTree . snd . munzip) tree

main :: IO ()
main = do
  input <- lines <$> readFile "day7/input"
  putStrLn ":: Day 7 - Part 1"
  print $ exampleData
  putStrLn ":: Tests"
  print $ parseOnly (parseFullColor <* A.endOfInput) "light red bag"
  print $ parseOnly (parseBag       <* A.endOfInput) "2 muted blue bags"
  print $ parseOnly (parseBag       <* A.endOfInput) "3 muted red bages"
  print $ parseOnly (parseRule      <* A.endOfInput) "dark orange bags contain 3 bright white bags, 4 muted yellow bags."
  putStrLn ":: Parsing test data"
  print $ parseStrs exampleData
  putStrLn ":: Building tree"
  pPrint $ mkTree (Color "shiny" "gold", 1) <$> (parseStrs exampleData)
  putStrLn ":: Solving Part 1"
  print $ solvePart1 (Color "shiny" "gold") exampleData
  print $ solvePart1 (Color "shiny" "gold") input
  putStrLn ":: Solving Part 2"
  print $ solvePart2 (Color "shiny" "gold") exampleData
  print $ solvePart2 (Color "shiny" "gold") exampleData2
  print $ solvePart2 (Color "shiny" "gold") input