{-# LANGUAGE OverloadedStrings #-}

import qualified Data.Text as T
import Data.Char (isSpace, isNumber)
import Data.List (find, isPrefixOf, stripPrefix, intercalate, sort)
import Data.Maybe (fromJust)
import Data.List.Split (startsWith, splitOn)
import GHC.Utils.Misc (nTimes)

{- cabal:
build-depends: base, split
-}

-- i've been told this needs bigints so i've switched all Ints storing worry values to Integers

data Operation = Add Integer | Multiply Integer | MultiplySelf
  deriving (Show, Eq)

execOperation :: Operation -> Integer -> Integer
execOperation (Add n) x = x + n
execOperation (Multiply n) x = x * n
execOperation MultiplySelf x = x * x

-- part 2 optimization checks mean that this can only be Divisible, sorry!
newtype Condition = Divisible Integer
  deriving (Show, Eq)

check :: Condition -> Integer -> Bool
check (Divisible n) x = x `mod` n == 0

newtype Test = Test (Condition, Int, Int)
  deriving (Show, Eq)

executeTest :: Test -> Integer -> Int
executeTest (Test (cond, a, b)) n = if check cond n then a else b

data Monkey = Monkey {
  monkeyItems :: [Integer],
  monkeyOperation :: Operation,
  monkeyTest :: Test
} deriving (Show, Eq)

trim :: String -> String
trim = reverse . dropWhile isSpace . reverse . dropWhile isSpace

replace :: Int -> a -> [a] -> [a]
replace pos newVal list = take pos list ++ newVal : drop (pos+1) list

findByPrefix :: String -> [String] -> Maybe String
findByPrefix prefix = find (prefix `isPrefixOf`)

findByPrefixAndStrip :: String -> [String] -> Maybe String
findByPrefixAndStrip prefix list = do
  el <- findByPrefix prefix list
  stripPrefix prefix el

giveMonkeyItem :: Monkey -> Integer -> Monkey
giveMonkeyItem monkey item = Monkey (monkeyItems monkey ++ [item]) (monkeyOperation monkey) (monkeyTest monkey)

parseMonkey :: String -> Monkey
parseMonkey m = monkey
  where
    -- let's just assume the dataset isn't malformed, okay................. just This Once, index promise
    monkey = Monkey (fromJust parseItems) (fromJust parseOperation) (fromJust parseTest)

    parseItems :: Maybe [Integer]
    parseItems = do
      itemsStr <- findByPrefixAndStrip "Starting items: " ln
      return $ map read $ splitOn "," itemsStr
    parseOperation :: Maybe Operation
    parseOperation = do
      operationStr <- findByPrefixAndStrip "Operation: new = " ln
      let [leftHand, op, rightHand] = words operationStr
      return $ case op of
        "+" -> Add $ read rightHand
        "*" -> if rightHand == "old" then MultiplySelf else Multiply $ read rightHand
        p -> error $ "Unknown operator " ++ p
    parseTest :: Maybe Test
    parseTest = do
      cond <- parseCondition
      outcomeTrue <- parseOutcome "true"
      outcomeFalse <- parseOutcome "false"
      return $ Test (cond, outcomeTrue, outcomeFalse)
    parseCondition :: Maybe Condition
    parseCondition = do
      conditionStr <- findByPrefixAndStrip "Test: " ln
      let w = words conditionStr
      nStr <- find (isNumber . head) w
      let n = read nStr
      return $ case head w of
        "divisible" -> Divisible n
        cond -> error $ "Unknown condition" ++ cond
    parseOutcome :: String -> Maybe Int
    parseOutcome s = do
      outcomeStr <- findByPrefixAndStrip ("If " ++ s ++ ": ") ln
      return $ read $ last $ words outcomeStr

    ln = map trim $ lines m

parse :: String -> [Monkey]
parse l = map (parseMonkey . unlines . drop 1 . lines . T.unpack) $ T.splitOn "\n\n" $ T.pack l

-- > The process of each monkey taking a single turn is called a round.
executeRound :: [Monkey] -> [(Monkey, Int)]
executeRound monkeys = foldl executeTurn (zip monkeys (repeat 0)) [0 .. length monkeys - 1]
  where
    executeTurn :: [(Monkey, Int)] -> Int -> [(Monkey, Int)]
    executeTurn monkeys idx = replace idx (newMonkeys !! idx, inspectCount + itemCount) $ zip newMonkeys (map snd monkeys)
      where
        newMonkeys = processItem (map fst monkeys) idx
        itemCount = length $ monkeyItems $ fst $ monkeys !! idx
        inspectCount = snd $ monkeys !! idx

    processItem :: [Monkey] -> Int -> [Monkey]
    processItem monkeys idx
      | null items = monkeys
      | otherwise = processItem'
      where
        items = monkeyItems monkey
        operation = monkeyOperation monkey
        test = monkeyTest monkey
        monkey = monkeys !! idx

        -- processItem' is called when we know the list is not empty. prevents Dumb Errors
        processItem' = processItem afterThrowItem idx
          where
            afterThrowItem = replace throwItemTo (giveMonkeyItem throwItemToMonkey newItem) afterRemoveItem
            afterRemoveItem = replace idx (Monkey newItems operation test) monkeys
            newItems = tail items
            throwItemToMonkey = monkeys !! throwItemTo
            throwItemTo = executeTest test newItem
            newItem = execOperation operation item `mod` worryLimit
            item = head items
            -- bit ugly but it's alright
            worryLimit = product $ map ((\(Test (Divisible n, _, _)) -> n) . monkeyTest) monkeys

executeRoundAccum :: [(Monkey, Int)] -> [(Monkey, Int)]
executeRoundAccum m = zipWith (\x (m, y) -> (m, x + y)) inspections (executeRound monkeys)
  where
    monkeys = map fst m
    inspections = map snd m

-- lol
getMonkeyBusiness :: [Int] -> Int
getMonkeyBusiness values = sorted !! 0 * sorted !! 1
  where sorted = reverse $ sort values

main = interact $
  show
  . getMonkeyBusiness . map snd
  . nTimes 10000 executeRoundAccum
  . (\monkeys -> zip monkeys (repeat 0))
  . parse