module IntSet = struct
  include Set.Make(Int)

  let exists_opt (type a) (f: int -> a option) (set: t): a option =
    let exception Found of a in
    let f' x = Option.iter (fun x -> raise (Found x)) (f x) in
    match iter f' set with
    | exception Found x -> Some x
    | _                 -> None
end

let read_ints =
  Bracket.infile_lines ~line:int_of_string ~of_seq:IntSet.of_seq


let rec find ?(target=2020) ~n set =
  if n < 1 then
    raise (Invalid_argument "n must be at least 1")
  else if n = 1 then
    if IntSet.mem target set then Some [target] else None
  else
    let joined x =
      Option.map (List.cons x) (find ~target:(target - x) ~n:(n - 1) set) in
    IntSet.exists_opt joined set

let print_result xs =
  let res = List.fold_left ( * ) 1 xs in
  let rec print_prod fmt = function
    | []    -> Format.pp_print_text fmt "{}"
    | [x]   -> Format.fprintf fmt "%d" x
    | x::xs -> Format.fprintf fmt "%d * %a" x print_prod xs in
  Format.printf "%a = %d\n" print_prod xs res


let main' n input =
  match find ~n (read_ints input) with
  | Some xs -> print_result xs
  | None    -> raise Not_found

let mains = [|main' 2; main' 3|]

let main = function
  | [part; input] -> mains.(int_of_string part - 1) input
  | _             -> Usage.exit "1 <part> <infile>"


let%expect_test "part 1" =
  mains.(0) "../../data/day1"; (* is this always the right path??? *)
  [%expect{| 534 * 1486 = 793524 |}]

let%expect_test "part 2" =
  mains.(1) "../../data/day1";
  [%expect{| 71 * 686 * 1263 = 61515678 |}]