import Data.Vect
import Data.List1
import Data.SnocList
import Data.String
import Data.Fin
import System.File.ReadWrite

allFins' : (n : Nat) -> List (Fin n)
allFins' n = toList (allFins'' n)
  where
    allFins'' : (n : Nat) -> SnocList (Fin n)
    allFins'' 0 = Lin
    allFins'' 1 = Lin :< FZ
    allFins'' (S k) = map weaken (allFins'' k) :< last

%default total

data Error =
  StartingItems String
  | ParseTest String
  | ParseOperation String
  | ParseMonkey String
  | ParseMonkeys String
  | BadVect

Show Error where
  show (StartingItems str) = "StartingItems: " ++ str
  show (ParseTest str) = "ParseTest: " ++ str
  show (ParseOperation str) = "ParseOperation: " ++ str
  show (ParseMonkey str) = "ParseMonkey: " ++ str
  show (ParseMonkeys str) = "ParseMonkeys: " ++ str
  show BadVect = "Invalid vector length, somehow"

record Item where
  constructor MkItem
  worryLevel : Integer

Show Item where
  show (MkItem worryLevel) = "(worry: " ++ show worryLevel ++ ")"

-- parse the items line
parseItems : String -> Either Error (n : Nat ** Vect n Item)
parseItems str =
  let numeric = pack . dropWhile (not . isDigit) . unpack $ str
      numStrings = split (== ',') numeric
      numbers = traverse parsePositive (forget numStrings)
  in case numbers of
       Nothing => Left (StartingItems str)
       Just numbers =>
         let items = map MkItem numbers
             len = length items
         in case toVect len items of
              Nothing => Left BadVect
              Just items => Right (len ** items)

record Test (monkeys : Nat) where
  constructor MkTest
  divisible_by : Integer
  true_monkey : Fin monkeys
  false_monkey : Fin monkeys

Show (Test monkeys) where
  show (MkTest divisible_by true_monkey false_monkey) =
    "if divisible by " ++ show divisible_by ++ " then " ++ show true_monkey ++ " else " ++ show false_monkey

parseTest : {monkeys : Nat } -> Vect 3 String -> Either Error (Test monkeys)
parseTest {monkeys} input@[line_1, line_2, line_3] =
  do Just divisible_by <- pure . parsePositive . pack . dropWhile (not . isDigit) . unpack $ line_1
       | Nothing => Left error
     Just true_monkey <- pure . (lineToNat >=> finify) $ line_2
       | Nothing => Left error
     Just false_monkey <- pure . (lineToNat >=> finify) $ line_3
       | Nothing => Left error
     Right (MkTest divisible_by true_monkey false_monkey)
  where lineToNat : String -> Maybe Nat
        lineToNat = parsePositive . pack . dropWhile (not . isDigit) . unpack
        error : Error
        error = ParseTest (joinBy "\n" (toList input))
        finify : Nat -> Maybe (Fin monkeys)
        finify n = natToFin n monkeys

applyTest : Test n -> Item -> Fin n
applyTest (MkTest divisible_by true_monkey false_monkey) (MkItem worryLevel) =
  if worryLevel `mod` divisible_by == 0
    then true_monkey
    else false_monkey

data Operation =
  Multiply Integer
  | Add Integer
  | MultiplySelf

Show Operation where
  show (Multiply k) = "old * " ++ show k
  show (Add k) = "old + " ++ show k
  show MultiplySelf = "old * old"

parseOperation : String -> Either Error Operation
parseOperation str =
  let op = dropWhile (\x => not (x == '+' || x == '*')) . unpack $ str
  in case op of
       ('*' :: ' ' :: 'o' :: 'l' :: 'd' :: _) => Right MultiplySelf
       ('*' :: ' ' :: rest) =>
         case parsePositive (pack rest) of
           Nothing => Left (ParseOperation str)
           Just other => Right (Multiply other)
       ('+' :: ' ' :: rest) =>
         case parsePositive (pack rest) of
           Nothing => Left (ParseOperation str)
           Just other => Right (Add other)
       _ => Left (ParseOperation str)

applyOperation : Operation -> Item -> Item
applyOperation (Multiply k) (MkItem worryLevel) = MkItem (k * worryLevel)
applyOperation (Add k) (MkItem worryLevel) = MkItem (k + worryLevel)
applyOperation MultiplySelf (MkItem worryLevel) = MkItem (worryLevel * worryLevel)

calmItem : Item -> Item
calmItem (MkItem worryLevel) = MkItem (worryLevel `div` 3)

record Monkey (monkeys : Nat) where
  constructor MkMonkey
  id : Fin monkeys
  item_count : Nat
  items : Vect item_count Item
  operation : Operation
  test : Test monkeys
  inspected : Nat

Show (Monkey monkeys) where
  show (MkMonkey id item_count items operation test inspected) =
    "Monkey " ++ show id ++ ":\n" ++
    "Items(" ++ show item_count ++ "): " ++ show items ++ "\n" ++
    "Operation: " ++ show operation ++ "\n" ++
    "Test: " ++ show test ++ "\n" ++
    "Inspected: " ++ show inspected

parseMonkey : {monkeys : Nat} -> Vect 6 String -> Either Error (Monkey monkeys)
parseMonkey {monkeys} input@[id, starting_items, op, test_1, test_2, test_3] =
  let test : Vect 3 String
      test = [test_1, test_2, test_3]
  in do Just id <- pure . (parsePositive >=> (\x => natToFin x monkeys)) . pack . filter isDigit . unpack $ id
          | Nothing => Left error
        (item_count ** items) <- parseItems starting_items
        operation <- parseOperation op
        test <- parseTest test
        Right (MkMonkey id item_count items operation test 0)
  where error : Error
        error = ParseMonkey (joinBy "\n" (toList input))

record Monkeys (count : Nat) where
  constructor MkMonkeys
  monkeys : Vect count (Monkey count)

showMonkeys : Monkeys c -> String
showMonkeys (MkMonkeys monkeys) =
    show (length monkeys) ++ " Monkeys: \n" ++ joinBy "\n\n" (map show (toList monkeys))

Show (Monkeys c) where
  show monkeys = showMonkeys monkeys

parseMonkeys : String -> Either Error (n : Nat ** Monkeys n)
parseMonkeys str =
  let individuals = forget . split (== "") . map trim . lines $ str
      vects = traverse (toVect 6) individuals
  in do Just vects <- pure vects
          | Nothing => Left (ParseMonkeys str)
        let count = length vects
        Just vects <- pure (toVect count vects)
          | Nothing => Left (ParseMonkeys str)
        monkeys <- traverse (parseMonkey {monkeys = count}) vects
        Right (count ** MkMonkeys monkeys)

throw : Monkeys c -> Item -> Fin c -> Monkeys c
throw (MkMonkeys monkeys) item idx =
  let monkey@(MkMonkey id item_count items operation test inspected) = index idx monkeys
      newItems = snoc items item
      newMonkey = MkMonkey id (S item_count) newItems operation test inspected
      newMonkeys = replaceAt idx newMonkey monkeys
  in MkMonkeys newMonkeys

turn : (calmOp : Item -> Item) -> Monkeys c -> Fin c -> Monkeys c
turn calmOp inputMonkeys@(MkMonkeys monkeys) idx =
  let inputMonkey@(MkMonkey id item_count items operation test inspected) = index idx monkeys
      (outputMonkey, (MkMonkeys resMonkeys)) = itemsHelper inputMonkey inputMonkeys
      outputMonkeys = replaceAt idx outputMonkey resMonkeys
  in MkMonkeys outputMonkeys
  where itemsHelper : Monkey c -> Monkeys c -> (Monkey c, Monkeys c)
        itemsHelper monkey@(MkMonkey id 0 [] operation test inspected) y = (monkey, y)
        itemsHelper monkey@(MkMonkey id (S len) (x :: xs) operation test inspected) y =
          let operated = applyOperation operation x
              calmed = calmOp operated
              throwTo = applyTest test calmed
              monkeys = throw y calmed throwTo
          in itemsHelper (assert_smaller monkey (MkMonkey id len xs operation test (S inspected))) monkeys

round : {c : Nat} -> (calmOp : Item -> Item) -> Monkeys c -> Monkeys c
round {c} calmOp input =
  let indexes = allFins' c
  in foldl (turn calmOp) input indexes

rounds : {c: Nat} -> (d: Nat) -> (calmOp : Item -> Item) -> Monkeys c -> Vect d (Monkeys c)
rounds 0 calmOp x = []
rounds (S k) calmOp x =
  let newMonkeys = round calmOp x
  in newMonkeys :: rounds k calmOp newMonkeys

part1 : {c : Nat} -> Monkeys c -> Nat
part1 x =
  let rounds = rounds 20 calmItem x
      lastRound = last rounds
  in foldl (*) 1 . take 2 . reverse . sort . toList . map inspected . monkeys $ lastRound

part2 : {c : Nat} -> Monkeys c -> Nat
part2 x =
  let multiple = product . map divisible_by . map test . monkeys $ x
      rounds = rounds 10000 (\(MkItem x) => MkItem (x `mod` multiple)) x
      lastRound = last rounds
  in foldl (*) 1 . take 2 . reverse . sort . toList . map inspected . monkeys $ lastRound

partial main : IO ()
main =
  do Right file <- readFile "input"
       | Left err => printLn err
     Right (count ** monkeys) <- pure (parseMonkeys file)
       | Left err => printLn err
     putStr "Part 1: "
     printLn (part1 monkeys)
     putStr "Part 2: "
     printLn (part2 monkeys)