{-# LANGUAGE TemplateHaskell #-}

import Control.Lens
import Data.List
import Data.List.Split
import Data.Maybe

data Monkey =
  Monkey
    { mID :: Int
    , _items :: [Int]
    , operation :: Int -> Int
    , test :: Int
    , 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: divisible by " ++
    show d ++
    "\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 $ day11 20 (`div` 3) monkes
  print $ day11 10000 id monkes

day11 :: Int -> (Int -> Int) -> [Monkey] -> Int
day11 x f m = g $ take (succ x) $ iterate (turns f) m
  where
    g = product . take 2 . reverse . sort . map _inspectCounter . last

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

turns :: (Int -> Int) -> [Monkey] -> [Monkey]
turns f = turns' 0
  where
    turns' n ms
      | length ms == n = ms
      | otherwise = turns' (succ n) monkeNew
      where
        itemsNew =
          [f $ flip mod modNum $ operation (ms !! n) x | x <- _items (ms !! n)]
        monkeEmptied =
          ms &
          ix n .~ ((ms !! n) & items .~ [] & inspectCounter +~ length itemsNew)
        monkeNew = monkeThrow itemsNew n monkeEmptied
        modNum = foldl1 lcm $ map test ms

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 mod x (test (ms !! n)) == 0
        then fst targets
        else snd targets
    targets = nextMonkey (ms !! n)
monkeThrow _ _ ms = ms