open Containers

type monkey =
  { mutable items : int list
  ; mutable inspected : int
  ; operation : int -> int
  ; test : int -> int
  ; test_num : int
  }

let parse_monkeys () =
  let lines = Utils.lines_of_input 11 in
  let rec parse lines monkeys =
    let lines = List.drop_while (String.equal "") lines in
    let lines, rest = List.take_drop_while Fun.(String.equal "" %> not) lines in
    let lines = List.map Fun.(String.trim %> String.split ~by:" ") lines in
    match lines with
    | [ [ "Monkey"; _ ]
      ; "Starting" :: "items:" :: items
      ; [ "Operation:"; "new"; "="; "old"; op; op_num ]
      ; [ "Test:"; "divisible"; "by"; test_num ]
      ; [ "If"; "true:"; "throw"; "to"; "monkey"; if_true ]
      ; [ "If"; "false:"; "throw"; "to"; "monkey"; if_false ]
      ] ->
      let items =
        List.map Fun.(String.replace ~which:`All ~sub:"," ~by:"" %> int_of_string) items
      in
      let op = if String.equal op "+" then ( + ) else ( * ) in
      let operation =
        match op_num with
        | "old" -> fun old -> op old old
        | num -> fun old -> op old @@ int_of_string num
      in
      let test_num = int_of_string test_num in
      let if_true = int_of_string if_true in
      let if_false = int_of_string if_false in
      let test score = if Int.(score mod test_num) = 0 then if_true else if_false in
      let monkey = { items; operation; test; test_num; inspected = 0 } in
      parse rest (monkey :: monkeys)
    | _ -> monkeys
  in
  Array.of_list @@ List.rev @@ parse lines []
;;

let throw item monkey = monkey.items <- List.append [ item ] monkey.items

let round monkies ~manageable =
  let divisor = Array.fold_left (fun acc monkey -> acc * monkey.test_num) 1 monkies in
  let destress score = if manageable then score / 3 else Int.(score mod divisor) in
  monkies
  |> Array.iter (fun monkey ->
    monkey.items
    |> List.iter (fun item ->
      let item = destress @@ monkey.operation item in
      let throw_to = monkey.test item in
      throw item monkies.(throw_to));
    monkey.inspected <- monkey.inspected + List.length monkey.items;
    monkey.items <- [])
;;

let solve rounds ~manageable =
  let monkies = parse_monkeys () in
  for _ = 1 to rounds do
    round monkies ~manageable
  done;
  monkies
  |> Array.to_list
  |> List.map (fun monkey -> monkey.inspected)
  |> List.sort Int.compare
  |> List.rev
  |> List.take 2
  |> List.reduce ( * )
  |> Option.get_exn_or "aoc goofed up the input"
;;

let part1 () = solve 20 ~manageable:true
let part2 () = solve 10_000 ~manageable:false

let%expect_test "Day 11.1" =
  Printf.printf "%d" @@ part1 ();
  [%expect {| 56350 |}]
;;

let%expect_test "Day 11.2" =
  Printf.printf "%d" @@ part2 ();
  [%expect {| 13954061248 |}]
;;