(** [matched_group_opt n s] returns [Some (Str.matched_group n s)] if the [n]th
    group matched, or [None] if the group did not match. *)
let matched_group_opt n s = try Some (Str.matched_group n s) with _ -> None

(** [find_nums s] returns the list of instructions given in [s]. *)
let instrs_of_string s =
  let r =
    Str.regexp
      {|do()\|don't()\|mul(\([0-9][0-9]?[0-9]?\),\([0-9][0-9]?[0-9]?\))|}
  in
  let rec impl acc pos =
    try
      let _ = Str.search_forward r s pos in
      let p0 = Str.matched_group 0 s in
      let p1 = matched_group_opt 1 s in
      let p2 = matched_group_opt 2 s in
      impl
        ((p0, Option.map int_of_string p1, Option.map int_of_string p2) :: acc)
        (Str.match_end ())
    with Not_found -> acc
  in
  List.rev (impl [] 0)

let instrs_of_file fname =
  Aoc.strings_of_file fname |> List.map instrs_of_string |> List.concat

(** [mac_opt acc a b] returns [acc + a' * b'] if [a = Some a'] and 
    [b = Some b'].  If either [a] or [b] are [None] then the result is [acc]. *)
let mac_opt acc a b =
  acc + (Option.value a ~default:0 * Option.value b ~default:0)

let day2403a = List.fold_left (fun acc (_, a, b) -> mac_opt acc a b) 0

let day2403b lst =
  let rec impl acc enabled = function
    | [] -> acc
    | ("do()", _, _) :: t -> impl acc true t
    | ("don't()", _, _) :: t -> impl acc false t
    | (_, a, b) :: t ->
        if enabled then impl (mac_opt acc a b) true t else impl acc false t
  in
  impl 0 true lst

let _ =
  Aoc.main instrs_of_file
    [ (string_of_int, day2403a); (string_of_int, day2403b) ]