
(* ----------- aux -------------------------------------------------------- *)
let int_of_string s =
  try int_of_string s
  with Failure _ -> Printf.eprintf "ERROR: not a number \n"; exit 1

let unix_read ic len =
  let all = Bytes.create len in
  let rec loop cur len =
    if cur >= len then ()
    else
      let n = Unix.read ic all cur (len-cur) in
      loop (cur+n) len in
  loop 0 len;
  Bytes.to_string all

let unix_open_file file =
  let ic = Unix.openfile file [Unix.O_RDONLY] 0 in
  let len = Unix.lseek ic 0 Unix.SEEK_END in
  let _ = Unix.lseek ic 0 Unix.SEEK_SET in
  ic, len

let unix_read_whole_file file =
  let ic, len = unix_open_file file in
  let data = unix_read ic len in
  Unix.close ic;
  data

(* ----------- patches ----------------------------------------------------- *)

type patch = {
  start : int;
  stop : int;
  text : string;
}

let parse_one_patch text =
  let r = Str.regexp ".*characters \\([0-9]+\\)-\\([0-9]+\\).*" in
  let beginning = String.index text '\n' in
  let header = String.sub text 0 beginning in
  let start = int_of_string @@ Str.replace_first r "\\1" header in
  let stop = int_of_string @@ Str.replace_first r "\\2" header in
  let text = String.sub text (beginning+1) (String.length text - beginning -1) in
  { start; stop; text }

let delim_start = "(* ------------ generator ----------------"
let delim_mid   = "------------ generated ---------------- *)"
let delim_end   = "(* ----------------- end ----------------- *)"

let parse_patch_file file =
  let text = unix_read_whole_file file in
  let patches = Str.split (Str.regexp_string "\nHIERARCHY BUILDER PATCH v1\n") text in
  let patches = List.map parse_one_patch patches in
  patches

let patch_file file patches =
  let ic, len = unix_open_file file in

  let out = Buffer.create 1024 in

  (* HACK since the locs don't include #[attrs], to be removed in Coq 8.14 https://github.com/coq/coq/pull/13844 *)
  let looks_like_the_beginning_of_a_command before i =
    Str.string_match Str.(regexp " *\\(Time\\|Fail\\|#\\[\\)") before (i+1) in

  let apply_one_hunk cur { start; stop; text } =
    if cur > start then
      failwith "Outdated patch?";
    let before = unix_read ic (start - cur) in
    let before, extra =
      (* HACK since the locs don't include #[attrs], to be removed in Coq 8.14 https://github.com/coq/coq/pull/13844 *)
      match String.rindex_opt before '\n' with
      | Some i when looks_like_the_beginning_of_a_command before i ->
        String.sub before 0 (i+1),
        String.sub before (i+1) (String.length before - i-1)
      | _ -> before, ""
      in
    let current = unix_read ic (stop - start) in
    Printf.bprintf out "%s\n%s\n%s%s\n%s\n%s\n%s\n"
      before
      delim_start
      extra
      current
      delim_mid
      text
      delim_end;
    stop
  in

  Printf.printf "Patching %s ..." file;
  let cur = List.fold_left apply_one_hunk 0 patches in
  let rest = unix_read ic (len - cur) in
  Printf.bprintf out "%s" rest;
  Unix.close ic;
  let oc = open_out file in
  Buffer.output_buffer oc out;
  close_out oc;
  Printf.printf "done.\n"
;;

let reset_file file =
  let ic, len = unix_open_file file in
  let text = unix_read ic len in
  let lines = String.split_on_char '\n' text in
  let rec copy = function
    | x :: xs when x = delim_start -> `Glue :: copy xs
    | x :: xs when x = delim_mid -> erase xs
    | x :: xs -> `Keep x :: copy xs
    | [] -> []
  and erase = function
    | x :: xs when x = delim_end -> `Glue :: copy xs
    | _ :: xs -> erase xs
    | [] -> assert false
  in
  let new_lines = copy lines in
  let rec glue = function
    | [] -> []
    | `Keep x :: `Glue :: `Keep y :: rest -> glue (`Keep(x^y) :: rest)
    | `Keep x :: xs -> x :: glue xs
    | `Glue :: xs -> glue xs
  in
  Unix.close ic;
  let new_text = String.concat "\n" (glue new_lines) in
  if new_text = text then
    Printf.printf "Skip %s since it contains no generated code\n" file
  else
    let oc = open_out file in
    Printf.fprintf oc "%s" new_text;
    close_out oc;
    try
      Sys.rename (file ^ ".hb") (file ^ ".hb.old");
      Printf.printf "Reset %s (patch file renamed to %s.hb.old)\n" file file
    with Sys_error _ -> Printf.printf "Reset %s\n" file


(* ----------- CLI ----------------------------------------------------- *)

let usage () =
  Printf.printf {|
Command line utility to apply patches generated by HB.

After building your project with logging enabled, eg

  COQ_ELPI_ATTRIBUTES='hb(log(raw))' make

each file.v file contanining calls to HB is paired with a file.v.hb patch
file which can be applied by this utility. You can also use
COQ_ELPI_ATTRIBUTES='hb(log)' to get terms printed in a nicer, more compact,
way but without all details they may not be parsable without manual editing.

usage:
  hb patch [-f] <file.v>..  Applies the patches in <file.v.hb> to <file.v>.
                            -f forces patch application even if the source file
                            is newer than the patch.
  hb reset <file.v>..       Erases all generated code from source file. It does
                            nothing if the file is not patched. If a patch file
                            <file.v.hb> exists it is renamed to <file.v.hb.old>.

  hb show <file.v.hb>..     Lists all the patches contained in <file.v.hb> (debugging).

|};
  exit 1

let apply_patches ~force file =
  let pfile = file ^ ".hb" in
  if Sys.file_exists pfile then
    let patches = parse_patch_file pfile in
    let { Unix.st_mtime = time_patch } = Unix.stat pfile in
    let { Unix.st_mtime = time_orig } = Unix.stat file in
    if time_orig > time_patch && not force then
      Printf.eprintf "Skip %s since it is more recent than %s.\n" file pfile
    else
      patch_file (Filename.chop_extension pfile) patches
  else
     Printf.printf "Skip %s since it has no patch\n" file
 

let list_patches pfile =
  let patches = parse_patch_file pfile in
  List.iter (fun { start; stop; text } ->
    Printf.printf "Replace %d-%d with\n%s\n\n" start stop text)
    patches

let reset_patches file =
  reset_file file

let () =
  let args = Array.to_list Sys.argv in
  match args with
  | _ :: "patch" :: "-f" :: (_ :: _ as args) -> List.iter (apply_patches ~force:true) args
  | _ :: "patch" :: (_ :: _ as args) -> List.iter (apply_patches ~force:false) args
  | _ :: "reset" :: (_ :: _ as args) -> List.iter reset_patches args
  | _ :: "show"  :: (_ :: _ as args) -> List.iter list_patches args
  | _ -> usage ()

