(**************************************************************************)
(*                                                                        *)
(*    Copyright (c) 2014 - 2018.                                          *)
(*    Dynamic Ledger Solutions, Inc. <contact@tezos.com>                  *)
(*                                                                        *)
(*    All rights reserved. No warranty, explicit or implicit, provided.   *)
(*                                                                        *)
(**************************************************************************)

let warnings = "+a-4-6-7-9-29-40..42-44-45-48"
let warn_error = "-a+8"

let () =
  Clflags.unsafe_string := false

(** Override the default 'Env.Persistent_signature.load'
    with a lookup in locally defined hashtable.
*)

let preloaded_cmis : (string, Env.Persistent_signature.t) Hashtbl.t =
  Hashtbl.create ~random:true 42

(* Set hook *)
let () =
  Env.Persistent_signature.load :=
    (fun ~unit_name ->
       try Some (Hashtbl.find preloaded_cmis (String.capitalize_ascii unit_name))
       with Not_found -> None)

let load_cmi_from_file file =
  Hashtbl.add preloaded_cmis
    (String.capitalize_ascii Filename.(basename (chop_extension file)))
    { filename = file ;
      cmi = Cmi_format.read_cmi file ;
    }

let load_embeded_cmi (unit_name, content) =
  let content = Bytes.of_string content in
  (* Read cmi magic *)
  let magic_len = String.length Config.cmi_magic_number in
  let magic = Bytes.sub content 0 magic_len in
  assert (magic = Bytes.of_string Config.cmi_magic_number) ;
  (* Read cmi_name and cmi_sign *)
  let pos = magic_len in
  let (cmi_name, cmi_sign) = Marshal.from_bytes content pos in
  let pos = pos + Marshal.total_size content pos in
  (* Read cmi_crcs *)
  let cmi_crcs = Marshal.from_bytes content pos in
  let pos = pos + Marshal.total_size content pos  in
  (* Read cmi_flags *)
  let cmi_flags = Marshal.from_bytes content pos in
  (* TODO check crcrs... *)
  Hashtbl.add
    preloaded_cmis
    (String.capitalize_ascii unit_name)
    { filename  =  unit_name ^ ".cmi" ;
      cmi = { cmi_name; cmi_sign; cmi_crcs; cmi_flags } ;
    }

let load_embeded_cmis cmis = List.iter load_embeded_cmi cmis

(** Compilation environment.

    [tezos_protocol_env] defines the list of [cmi] available while compiling
    the protocol version. The [cmi] are packed into the [tezos-node]
    binary by using [ocp-ocamlres], see the Makefile.

    [register_env] defines a complementary list of [cmi] available
    while compiling the generated [register.ml] file (that register
    the protocol first-class module into the [Updater.versions]
    hashtable).

*)


let tezos_protocol_env =
  let open Embedded_cmis in
  [
    "CamlinternalFormatBasics", camlinternalFormatBasics_cmi ;
    "Tezos_protocol_environment_sigs", tezos_protocol_environment_sigs_cmi ;
    "Tezos_protocol_environment_sigs__V1", tezos_protocol_environment_sigs__V1_cmi ;
  ]

let register_env =
  let open Embedded_cmis in
  [
    "tezos_protocol_registerer__Registerer", tezos_protocol_registerer__Registerer_cmi ;
  ]


(** Helpers *)

let (//) = Filename.concat

let create_file ?(perm = 0o644) name content =
  let open Unix in
  let fd = openfile name [O_TRUNC; O_CREAT; O_WRONLY] perm in
  ignore(write_substring fd content 0 (String.length content));
  close fd

let safe_unlink file =
  try Unix.unlink file
  with Unix.Unix_error(Unix.ENOENT, _, _) -> ()

let unlink_cmi dir (file, _) =
  safe_unlink (dir // file ^ ".cmi")

let unlink_object obj =
  safe_unlink obj;
  safe_unlink (Filename.chop_suffix obj ".cmx" ^ ".cmi");
  safe_unlink (Filename.chop_suffix obj ".cmx" ^ ".o")

let debug_flag = ref false

let debug fmt =
  if !debug_flag then Format.eprintf fmt
  else Format.ifprintf Format.err_formatter fmt

let hash_file file =
  let open Blake2 in
  let buflen = 8092 in
  let buf = BytesLabels.create buflen in
  let fd = Unix.openfile file [Unix.O_RDONLY] 0o600 in
  let state = Blake2b.init 32 in
  let loop () =
    match Unix.read fd buf 0 buflen with
    | 0 -> ()
    | nb_read ->
        Blake2b.update state
          (Cstruct.of_bytes
             (if nb_read = buflen then buf else BytesLabels.sub buf ~pos:0 ~len:nb_read))
  in
  loop () ;
  Unix.close fd ;
  let Blake2b.Hash h = Blake2b.final state in
  Cstruct.to_string h

let mktemp_dir () =
  Filename.get_temp_dir_name () //
  Printf.sprintf "tezos-protocol-build-%06X" (Random.int 0xFFFFFF)

(** Main *)

type driver = {
  compile_ml: ?for_pack:string -> string -> string ;
  pack_objects: string -> string list -> string ;
  link_shared: string -> string list -> unit ;
}

let main { compile_ml ; pack_objects ; link_shared } =
  Random.self_init () ;
  let anonymous = ref []
  and static = ref false
  and register = ref false
  and build_dir = ref None in
  let args_spec = [
    "-static", Arg.Set static, " Only build the static library (no .cmxs)";
    "-register", Arg.Set register, " Generete the `Registerer` module";
    "-bin-annot", Arg.Set Clflags.binary_annotations, " (see ocamlopt)" ;
    "-g", Arg.Set Clflags.debug, " (see ocamlopt)" ;
    "-build-dir", Arg.String (fun s -> build_dir := Some s),
    "use custom build directory and preserve build artifacts"
  ] in
  let usage_msg =
    Printf.sprintf
      "Usage: %s [options] <out> <srcdir>\nOptions are:"
      Sys.argv.(0) in
  Arg.parse args_spec (fun s -> anonymous := s :: !anonymous) usage_msg ;
  let (output, source_dir) =
    match List.rev !anonymous with
    | [ output ; protocol_dir ] -> output, protocol_dir
    | _ -> Arg.usage args_spec usage_msg ; Pervasives.exit 1 in
  let build_dir =
    match !build_dir with
    | None ->
        let dir = mktemp_dir () in
        at_exit (fun () -> Lwt_main.run (Lwt_utils_unix.remove_dir dir)) ;
        dir
    | Some dir -> dir in
  let hash, protocol =
    match Lwt_main.run (Lwt_utils_unix.Protocol.read_dir source_dir) with
    | Ok v -> v
    | Error err ->
        Format.kasprintf Pervasives.failwith
          "Failed to read TEZOS_PROTOCOL: %a" pp_print_error err in
  Lwt_main.run (Lwt_utils_unix.create_dir ~perm:0o755 build_dir) ;
  Lwt_main.run (Lwt_utils_unix.create_dir ~perm:0o755 (Filename.dirname output)) ;
  (* Generate the 'functor' *)
  let functor_file = build_dir // "functor.ml" in
  let oc = open_out functor_file in
  Packer.dump oc
    (Array.map
       begin fun { Protocol.name }  ->
         let name_lowercase = String.uncapitalize_ascii name in
         source_dir // name_lowercase ^ ".ml"
       end
       (Array.of_list protocol.components)) ;
  close_out oc ;
  (* Compile the protocol *)
  let proto_cmi = Filename.chop_extension functor_file ^ ".cmi" in
  let functor_unit =
    String.capitalize_ascii
      Filename.(basename (chop_extension functor_file)) in
  let for_pack = String.capitalize_ascii (Filename.basename output) in
  (* Initialize the compilers *)
  Compenv.(readenv Format.err_formatter Before_args);
  Clflags.nopervasives := true;
  Clflags.no_std_include := true ;
  Clflags.include_dirs := [Filename.dirname functor_file] ;
  Warnings.parse_options false warnings ;
  Warnings.parse_options true warn_error ;

  load_embeded_cmis tezos_protocol_env ;
  let packed_protocol_object = compile_ml ~for_pack functor_file in

  let register_objects =
    if not !register then
      []
    else begin
      load_embeded_cmis register_env ;
      load_cmi_from_file proto_cmi ;
      (* Compiler the 'registering module' *)
      let register_file = Filename.dirname functor_file // "register.ml" in
      create_file register_file
        (Printf.sprintf
           "module Name = struct let name = %S end\n\
           \ let () = Tezos_protocol_compiler_native__Registerer.register Name.name (module %s.Make)"
           (Protocol_hash.to_b58check hash)
           functor_unit) ;
      let register_object = compile_ml ~for_pack register_file in
      [ register_object ]
    end
  in

  let resulting_object =
    pack_objects output (packed_protocol_object :: register_objects) in

  (* Create the final [cmxs] *)
  if not !static then begin
    Clflags.link_everything := true ;
    link_shared (output ^ ".cmxs") [resulting_object] ;
  end