open Cmdliner
open Syntax

let parse filename =
  match Parsing.Driver.parse_file filename with
  | Some program -> program
  | None -> failwith "failed to parse program, program is empty"

(* create file if it doesn't exist overwrite otherwise *)
let write_to_file filename f =
  let oc = open_out filename in
  let fmt = Format.formatter_of_out_channel oc in
  f fmt;
  Format.pp_print_newline fmt ();
  close_out oc

let process_file ?(include_extremes = true) program output_dir num_samples
    num_bins =
  let samples =
    Dynamizer.sample_from_bins ~include_extremes program num_samples num_bins
  in
  Format.printf "replacement_strings:\n%!";
  List.iteri
    (fun i sample_list ->
      List.iteri
        (fun j (sample, replacement_list) ->
          let replacement_string = Dynamizer.to_hex_string replacement_list in
          Format.printf "%s\n%!" replacement_string;
          let filename =
            Filename.concat output_dir
              (Printf.sprintf "sample_%d_%d_%s.txt" i j replacement_string)
          in
          write_to_file filename (fun fmt ->
              (Parse_tree.Program.pp ~pp_meta:(fun _ _ -> ())) fmt sample))
        sample_list)
    samples

let clear_directory dir_path =
  let files = Sys.readdir dir_path in
  Array.iter
    (fun file ->
      let file_path = Filename.concat dir_path file in
      if not (Sys.is_directory file_path) then Sys.remove file_path)
    files

let main input_file num_samples num_bins include_extremes verbose =
  (* Set verbosity *)
  Dynamizer.verbose := verbose;
  Format.printf "Processing file: %s\n" input_file;
  (* drop extension for input file and add _samples to the name
     e.g. input_file = "foo.ml" -> output_dir = "foo_samples"
  *)
  let output_dir =
    Format.sprintf "%s_samples" (Filename.chop_extension input_file)
  in
  Format.printf "Output directory: '%s'\n" output_dir;
  if Sys.file_exists output_dir && (not @@ Sys.is_directory output_dir) then
    failwith
      (Format.sprintf
         "A file with the name '%s' already exists, but it is not a directory"
         output_dir)
  else if Sys.file_exists output_dir then (
    Format.printf "'%s' already exists, clearing it\n" output_dir;
    clear_directory output_dir)
  else Sys.mkdir output_dir 0o755;
  let program = parse input_file in
  let lattice_size, static_type_count = Dynamizer.Count.count_program program in
  Format.printf "Number of all possible configurations (lattice size): %d\n"
    lattice_size;
  Format.printf "Number of all type nodes: %d\n" static_type_count;
  let num_samples =
    if num_samples = -1 then
      let scaling_factor = 10 in
      lattice_size * scaling_factor / num_bins
    else num_samples
  in
  Format.printf "Number of requested configurations (samples): %d\n" num_samples;
  Format.printf "Number of requested bins: %d\n" num_bins;
  process_file ~include_extremes program output_dir num_samples num_bins;
  Format.printf "Done\n%!"

let input_file_arg =
  let doc = "The input file to be processed" in
  let docv = "INPUT_FILE" in
  Arg.(required & pos 0 (some string) None & info [] ~docv ~doc)

let num_samples_arg =
  let doc = "The number of samples to generate" in
  let docv = "NUM_SAMPLES" in
  Arg.(value & opt int 10 & info [ "n"; "num-samples" ] ~docv ~doc)

let num_bins_arg =
  let doc = "The number of bins to use for sampling" in
  let docv = "NUM_BINS" in
  Arg.(value & opt int 5 & info [ "b"; "num-bins" ] ~docv ~doc)

let include_extremes_arg =
  let doc =
    "Whether to include the fully untyped and fully typed configurations"
  in
  let docv = "INCLUDE_EXTREMES" in
  Arg.(value & opt bool true & info [ "e"; "include-extremes" ] ~docv ~doc)

let verbose_arg =
  let doc = "Enable verbose logging output" in
  Arg.(value & flag & info [ "v"; "verbose" ] ~doc)

let cmd =
  let doc = "Generate multiple versions of a program using Dynamizer" in
  let info = Cmd.info "dynamizer" ~doc in
  Cmd.v info
    Term.(
      const main
      $ input_file_arg
      $ num_samples_arg
      $ num_bins_arg
      $ include_extremes_arg
      $ verbose_arg)

let () = exit (Cmd.eval cmd)
