(*
 * 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.
 *)

open Ppxlib

let flowlib_dir_ref = ref ""

let prelude_dir_ref = ref ""

let tslib_dir_ref = ref ""

let flowlib_cache = ref None

let prelude_cache = ref None

let tslib_cache = ref None

(* Read in all the flowlib files *)
let is_lib_file file =
  Filename.check_suffix file ".js"
  || Filename.check_suffix file ".d.ts"
  || Filename.check_suffix file ".d.mts"
  || Filename.check_suffix file ".d.cts"

let get_libs dir =
  Sys.readdir dir
  |> Array.fold_left
       (fun acc file ->
         if is_lib_file file then
           let contents = Script_utils.string_of_file (Filename.concat dir file) in
           (file, contents) :: acc
         else
           acc)
       []

(* Turn the (name, contents) list into a PPX ast (string * string) array
 * expression *)
let compute ~loc lib_dir =
  let open Ast_builder.Default in
  let libs = get_libs lib_dir in
  let hash =
    let state = Xx.init 0L in
    List.iter
      (fun (file, contents) ->
        Xx.update state file;
        Xx.update state contents)
      libs;
    state |> Xx.digest |> Xx.to_string |> estring ~loc
  in
  let contents =
    let elements =
      List.map
        (fun (name, contents) -> pexp_tuple ~loc [estring ~loc name; estring ~loc contents])
        libs
    in
    { pexp_desc = Pexp_array elements; pexp_loc = loc; pexp_loc_stack = []; pexp_attributes = [] }
  in
  (hash, contents)

let memo ~loc path cache_ref =
  match !cache_ref with
  | Some result -> result
  | None ->
    let result = compute ~loc path in
    cache_ref := Some result;
    result

let flowlib_hash ~loc ~path:_ =
  let (hash, _contents) = memo ~loc !flowlib_dir_ref flowlib_cache in
  hash

(* Whenever we see [%flowlib_contents], replace it wil the flowlib contents *)
let flowlib_contents ~loc ~path:_ =
  let (_hash, contents) = memo ~loc !flowlib_dir_ref flowlib_cache in
  contents

let prelude_hash ~loc ~path:_ =
  let (hash, _contents) = memo ~loc !prelude_dir_ref prelude_cache in
  hash

let prelude_contents ~loc ~path:_ =
  let (_hash, contents) = memo ~loc !prelude_dir_ref prelude_cache in
  contents

let tslib_hash ~loc ~path:_ =
  let (hash, _contents) = memo ~loc !tslib_dir_ref tslib_cache in
  hash

let tslib_contents ~loc ~path:_ =
  let (_hash, contents) = memo ~loc !tslib_dir_ref tslib_cache in
  contents

let make_rule name f =
  Context_free.Rule.extension
    (Extension.declare name Extension.Context.expression Ast_pattern.(pstr nil) f)

let rules =
  [
    make_rule "flowlib_hash" flowlib_hash;
    make_rule "flowlib_contents" flowlib_contents;
    make_rule "prelude_hash" prelude_hash;
    make_rule "prelude_contents" prelude_contents;
    make_rule "tslib_hash" tslib_hash;
    make_rule "tslib_contents" tslib_contents;
  ]

let () =
  Driver.add_arg "-flowlib" (Arg.Set_string flowlib_dir_ref) ~doc:"Path to flowlib directory";
  Driver.add_arg "-prelude" (Arg.Set_string prelude_dir_ref) ~doc:"Path to prelude directory";
  Driver.add_arg "-tslib" (Arg.Set_string tslib_dir_ref) ~doc:"Path to tslib directory";
  Driver.register_transformation ~rules "ppx_gen_flowlibs"
