{-# LANGUAGE TemplateHaskell #-}
import Control.Lens
import Data.List.Split
import Data.Maybe
import Data.List

data Monkey =
  Monkey
    { mID :: Int
    , _items :: [Int]
    , operation :: Int -> Int
    , test :: Int -> Bool
    , nextMonkey :: (Int, Int)
    , _inspectCounter :: Int
    }
makeLenses ''Monkey

instance Show Monkey where
  show (Monkey a b c d (e, f) g) =
    "Monkey " ++
    show a ++
    ":\n  Items: " ++
    show b ++
    "\n  Operation: e.g. 1 -> " ++
    show (c 1) ++
    ", 2 -> " ++
    show (c 2) ++
    "\n  Test: <<function>>" ++
    "\n   If true: throw to monkey " ++
    show e ++
    "\n    If false: throw to monkey " ++
    show f ++ "\nThis monkey has inspected " ++ show g ++ " items.\n"

main :: IO ()
main = do
  input <- map (map (splitOn ":") . lines) . splitOn "\n\n" <$> readFile "input"
  let monkes = mapMaybe parseMonkey input
  print $ product $ take 2 $ reverse $ sort $ map _inspectCounter $ last $ take 21 $ iterate turns monkes

parseMonkey :: [[String]] -> Maybe Monkey
parseMonkey [[a, _], [_, b], [_, c], [_, d], [_, e], [_, f]] =
  Just $
  Monkey
    (f' a)
    (read ('[' : b ++ "]"))
    (parseOP $ words c)
    (mtest $ read $ last $ words d)
    (f' e, f' f)
    0
  where
    f' = read . last . words
    mtest x y = y `mod` x == 0
    parseOP [_, _, _, x, y]
      | y == "old" = (^ 2)
      | x == "*" = (*) (read y)
      | otherwise = (+) (read y)
parseMonkey _ = Nothing

turns :: [Monkey] -> [Monkey]
turns = turns' 0
    where turns' n ms
            | length ms == n = ms
            | otherwise = turns' (succ n) monkeNew
                where itemsNew = [ flip div 3 $ operation (ms !! n) x  | x<-_items (ms !! n)]
                      monkeEmptied = ms & ix n .~ ((ms !! n) & items .~ [] & inspectCounter +~ length itemsNew)
                      monkeNew = monkeThrow itemsNew n monkeEmptied

monkeThrow :: [Int] -> Int -> [Monkey] -> [Monkey]
monkeThrow (x:xs) n ms = monkeThrow xs n (ms & ix next .~ ((ms !! next) & items .~ (_items (ms !! next) ++ [x]) ))
    where next = if test (ms !! n) x then fst targets else snd targets
          targets = nextMonkey (ms !! n)
monkeThrow _ _ ms = ms