import Data.List.Split (splitOn)
import qualified Control.Monad.Trans.State as St
import qualified Data.Map as M
import Control.Monad (replicateM_)
import Data.List (sort)

data Expr = Old | Lit Int | Add Expr Expr | Mul Expr Expr
data Test = T Int Int Int
data Part = Part1 | Part2
data Monkey = M [Int] Expr Test

getFactor :: Monkey -> Int 
getFactor (M _ _ (T x _ _)) = x 

parseElem :: String -> Expr
parseElem "old" = Old
parseElem i = Lit (read i)

parseExpr :: [String] -> Expr
parseExpr s
    | s !! 1 == "*" = Mul (parseElem $ s !! 0) (parseElem $ s !! 2)
    | s !! 1 == "+" = Add (parseElem $ s !! 0) (parseElem $ s !! 2)
    | otherwise = error $ "unknown expression: " ++ concat s

parseMonkey :: [String] -> Monkey
parseMonkey s = M startingItems expr test where
    startingItems = read $ "[" ++ (splitOn ": " (s !! 1) !! 1) ++ "]"
    expr = parseExpr $ last $ splitOn ["="] (words (s !! 2))
    testFactor = read $ last $ words (s !! 3)
    testTargetT = read $ last $ words (s !! 4)
    testTargetF = read $ last $ words (s !! 5)
    test = T testFactor testTargetT testTargetF

evalExpr :: Expr -> Int -> Int
evalExpr Old x = x
evalExpr (Lit i) x = i
evalExpr (Add a b) x = evalExpr a x + evalExpr b x
evalExpr (Mul a b) x = evalExpr a x * evalExpr b x

type Monkeys = M.Map Int (Monkey, Int)

steps :: Part -> Int -> Monkey -> St.State Monkeys (Int, Monkey)
steps part factor monkey@(M [] _ _) = do
    pure (0, monkey)
steps part factor (M (item:items) op test@(T fac tr fa)) = do
    let newItem' = evalExpr op item --`div` 3
    let newItem = case part of
                                Part1 -> evalExpr op item `div` 3
                                Part2 -> evalExpr op item `mod` factor
    let targetMonkey = if newItem `mod` fac == 0 then tr else fa
    monkeys <- St.get
    let (M items' op' test', count') = monkeys M.! targetMonkey
    let updated = M (newItem:items') op' test'
    St.put $ M.insert targetMonkey (updated, count') monkeys
    (count, final) <- steps part factor $ M items op test
    pure (count+1, final)

simulate :: Part -> Int -> Int -> St.State Monkeys ()
simulate part factor i = do
    monkeys <- St.get
    let (monkey, count) = monkeys M.! i
    (iters, finished) <- steps part factor monkey
    monkeys <- St.get
    St.put $ M.insert i (finished, count + iters) monkeys

round :: Part -> Int -> St.State Monkeys ()
round part factor = do
    monkeys <- St.get
    mapM_ (simulate part factor . fst) (M.toAscList monkeys)

main :: IO ()
main = do
    monkeys <- map parseMonkey . splitOn [""] . lines <$> getContents
    let factor = foldr1 lcm $ map getFactor monkeys
    let initialState = M.fromAscList $ zip [0..] (map (\x -> (x,0)) monkeys)
    
    putStr "part 1: "
    let action = replicateM_ 20 (Main.round Part1 factor)
    let results = St.execState action initialState
    let counts = map (snd . snd) (M.toAscList results)
    print $ counts
    print $ product $ take 2 $ reverse $ sort counts
    
    putStr "part 2: "
    let action = replicateM_ 10000 (Main.round Part2 factor)
    let results = St.execState action initialState
    let counts = map (snd . snd) (M.toAscList results)
    print $ counts
    print $ product $ take 2 $ reverse $ sort counts