(*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

let out_of_time ~options ~start_time =
  Unix.gettimeofday () -. start_time > Options.max_seconds_for_check_per_worker options

(** Check as many files as it can before it hits the timeout. The timeout is soft,
  so the file which exceeds the timeout won't be canceled. We expect most buckets
  to not hit the timeout *)
let rec job_helper ~check ~options ~start_time acc = function
  | [] -> (acc, [])
  | unfinished_files when out_of_time ~options ~start_time ->
    Hh_logger.debug
      "Bucket timed out! %d files finished, %d files unfinished"
      (List.length acc)
      (List.length unfinished_files);
    (acc, unfinished_files)
  | file :: rest ->
    let result =
      match check file with
      | Ok (Some (_, acc)) -> Ok (Some acc)
      | (Ok None | Error _) as result -> result
    in
    job_helper ~check ~options ~start_time ((file, result) :: acc) rest

let mk_job ~mk_check ~options () files =
  let start_time = Unix.gettimeofday () in
  let check = mk_check () in
  job_helper ~check ~options ~start_time [] files

(** A stateful (next, merge) pair. This lets us re-queue unfinished files which are returned
  when a bucket times out *)
let mk_next ~intermediate_result_callback ~max_size ~workers ~files =
  let total_count = List.length files in
  let todo = ref (files, total_count) in
  let finished_count = ref 0 in
  let num_workers = max 1 (Base.Option.value_map workers ~default:1 ~f:List.length) in
  let status_update () =
    MonitorRPC.status_update
      ~event:
        ServerStatus.(Checking_progress { total = Some total_count; finished = !finished_count })
  in
  let next () =
    let (remaining_files, remaining_count) = !todo in
    (* When we get near the end of the file list, start using smaller buckets in order
     * to spread the work across the available workers *)
    let bucket_size =
      if remaining_count >= max_size * num_workers then
        max_size
      else
        (remaining_count / num_workers) + 1
    in
    let (bucket, remaining_files) = Base.List.split_n remaining_files bucket_size in
    let bucket_size = List.length bucket in
    todo := (remaining_files, remaining_count - bucket_size);
    if bucket_size = 0 then
      Bucket.Done
    else
      Bucket.Job bucket
  in
  let merge (finished_file_accs, unfinished_files) acc =
    intermediate_result_callback finished_file_accs;
    let (remaining_files, remaining_count) = !todo in
    todo :=
      ( List.rev_append unfinished_files remaining_files,
        remaining_count + List.length unfinished_files
      );
    finished_count := !finished_count + List.length finished_file_accs;
    status_update ();
    List.rev_append finished_file_accs acc
  in
  (next, merge)
