(*
  This file is part of scilla.

  Copyright (c) 2018 - present Zilliqa Research Pvt. Ltd.

  scilla is free software: you can redistribute it and/or modify it under the
  terms of the GNU General Public License as published by the Free Software
  Foundation, either version 3 of the License, or (at your option) any later
  version.

  scilla is distributed in the hope that it will be useful, but WITHOUT ANY
  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  A PARTICULAR PURPOSE.  See the GNU General Public License for more details.

  You should have received a copy of the GNU General Public License along with
  scilla.  If not, see <http://www.gnu.org/licenses/>.
*)

open Core
open OUnit2
open Scilla_base.ScillaUtil.FilePathInfix

(* Parses and pretty-prints JSON to ignore whitespace changes
   in the future versions of JSON pretty-printers *)
let normalize_json s =
  s |> Yojson.Basic.from_string |> Yojson.Basic.pretty_to_string

(* Helper function inspired by one from Batteries library  *)
let stream_to_string fl =
  let buf = Buffer.create 4096 in
  (* OUnit uses potentially infinite streams and
     if a stream is finite, it may throw an exception *)
  (try Seq.iter (Buffer.add_char buf) fl with _exn -> ());
  Buffer.contents buf

type tsuite_env = {
  tests_dir : test_ctxt -> string;
  stdlib_dir : test_ctxt -> string;
  print_cli : test_ctxt -> bool;
  update_gold : test_ctxt -> bool;
  print_diff : test_ctxt -> bool;
  ext_ipc_server : test_ctxt -> string;
  server : test_ctxt -> bool;
}

let run_tests tests =
  let tests_dir_default = Sys_unix.getcwd () ^/ "tests" in
  let stdlib_dir_default = Sys_unix.getcwd () ^/ "src" ^/ "stdlib" in
  let ext_ipc_server_default = "" in
  let tests_dir =
    Conf.make_string "tests_dir" tests_dir_default "Directory containing tests"
  in
  let stdlib_dir =
    Conf.make_string "stdlib_dir" stdlib_dir_default
      "Directory containing stdlib"
  in
  let print_cli =
    Conf.make_bool "print_cli" false
      "Print command line arguments used for test(s)"
  in
  let update_gold =
    Conf.make_bool "update_gold" false
      "Ignore compare mismatch and update gold file(s)"
  in
  let print_diff =
    Conf.make_bool "print_diff" false
      "Print the diff between gold file and actual output"
  in
  let server = Conf.make_bool "server" false "Run tests in server-mode" in
  let ext_ipc_server =
    Conf.make_string "ext_ipc_server" ext_ipc_server_default
      "Address of external IPC server for IPC tests. Ensure that \"-runner \
       sequential\" is set"
  in

  let env : tsuite_env =
    {
      tests_dir;
      stdlib_dir;
      print_cli;
      update_gold;
      print_diff;
      server;
      ext_ipc_server;
    }
  in
  run_test_tt_main ("tests" >::: List.map ~f:(( |> ) env) tests)

let output_verifier gold_output msg print_diff output =
  let pp_diff fmt =
    let open Patdiff_kernel in
    let gold = Diff_input.{ name = "expected output"; text = gold_output } in
    let out = Diff_input.{ name = "actual output"; text = output } in
    let open Patdiff.Compare_core in
    match diff_strings Configuration.default ~prev:gold ~next:out with
    | `Same -> ()
    | `Different s ->
        (* s contains ANSI color codes *)
        Caml.Format.pp_force_newline fmt ();
        Caml.Format.pp_print_string fmt s
  in
  if print_diff then
    assert_equal
      ~cmp:(fun e o -> String.(strip e = strip o))
      ~pp_diff:(fun fmt _ -> pp_diff fmt)
      gold_output output ~msg
  else
    assert_equal
      ~cmp:(fun e o -> String.(strip e = strip o))
      ~printer:(fun s -> s)
      gold_output output ~msg

(* XXX: in case of json_errors we should probably not do string conversions
        and operate directly on the JSON's AST *)
let output_updater goldoutput_file test_name data ~json_errors =
  let update_goldoutput_file () =
    Out_channel.write_all goldoutput_file ~data;
    Printf.printf "Updated gold output for test %s\n" test_name
  in
  if json_errors && Sys_unix.file_exists_exn goldoutput_file then (
    let normalized_gold =
      normalize_json @@ In_channel.read_all goldoutput_file
    in
    let normalized_actual = normalize_json @@ data in
    if String.( <> ) normalized_gold normalized_actual then
      update_goldoutput_file ())
  else update_goldoutput_file ()

let prepare_cli_usage bin args = bin ^ " " ^ String.concat ~sep:" " args

let cli_usage_on_err bin args =
  "Command " ^ prepare_cli_usage bin args ^ " failed.\n"

let print_cli_usage flag bin args =
  if flag then Printf.printf "\nUsing CLI: %s\n" (prepare_cli_usage bin args)

module type TestSuiteInput = sig
  val tests : string list
  val gold_path : string -> string -> string list
  val test_path : string -> string list
  val runner : string
  val ignore_predef_args : bool
  val json_errors : bool
  val exit_code : UnixLabels.process_status
  val additional_libdirs : string list list
  val gas_limit : Stdint.uint64
  val custom_args : string list
  val provide_init_arg : bool
  val diff_filter : string -> string
end

module DiffBasedTests (Input : TestSuiteInput) = struct
  open Input

  let build_exp_tests env =
    List.map ~f:(fun fname ->
        fname >:: fun test_ctxt ->
        let open FilePath in
        let dir = env.tests_dir test_ctxt in
        let input_file = make_filename (test_path fname) in
        let init_file =
          make_filename (test_path (chop_extension fname ^ ".json"))
        in
        (* Verify standard output of execution with gold file *)
        let goldoutput_file = make_filename (gold_path dir fname) in
        let additional_dirs = List.map ~f:make_filename additional_libdirs in
        let stdlib = make_relative dir (env.stdlib_dir test_ctxt) in
        let path = string_of_path @@ (stdlib :: additional_dirs) in
        let args' =
          if ignore_predef_args then custom_args @ [ input_file ]
          else
            custom_args
            @ [
                "-libdir";
                path;
                input_file;
                "-gaslimit";
                Stdint.Uint64.to_string gas_limit;
              ]
            @ if json_errors then [ "-jsonerrors" ] else []
        in
        let args =
          if provide_init_arg then args' @ [ "-init"; init_file ] else args'
        in
        let msg = cli_usage_on_err runner args in
        print_cli_usage (env.print_cli test_ctxt) runner args;
        assert_command
          ~foutput:(fun s ->
            let actual_output = stream_to_string s in
            if env.update_gold test_ctxt then
              output_updater goldoutput_file input_file actual_output
                ~json_errors
            else if Sys_unix.file_exists_exn goldoutput_file then
              let gold_output =
                (* load all data from file *)
                let non_normalized_gold_output =
                  In_channel.read_all goldoutput_file
                in
                if json_errors then normalize_json non_normalized_gold_output
                else non_normalized_gold_output
              in
              output_verifier (diff_filter gold_output) msg
                (env.print_diff test_ctxt)
                (diff_filter actual_output)
            else
              assert_failure
                ("The gold file " ^ goldoutput_file ^ " does not exist"))
          ~exit_code ~use_stderr:true ~chdir:dir ~ctxt:test_ctxt runner args)

  let tests env = "exptests" >::: build_exp_tests env tests
end

module type TestSuiteMultiInput = sig
  val tests : string list list
  val gold_path : string -> string -> string list
  val test_path : string -> string list
  val runner : string
  val ignore_predef_args : bool
  val exit_code : UnixLabels.process_status
  val additional_libdirs : string list list
  val gas_limit : Stdint.uint64
  val custom_args : string list
  val provide_init_arg : bool
  val diff_filter : string -> string
end

module DiffBasedMultiTests (Input : TestSuiteMultiInput) = struct
  open Input

  let rec longest_common_string lhs rhs =
    let cmp xs ys = if List.length xs > List.length ys then xs else ys in
    match (lhs, rhs) with
    | [], _ | _, [] -> []
    | l :: ls, r :: rs ->
        if Char.equal l r then l :: longest_common_string ls rs
        else cmp (longest_common_string lhs rs) (longest_common_string ls rhs)

  let get_basename fnames =
    List.fold_left fnames
      ~init:(List.hd_exn fnames |> String.to_list)
      ~f:(fun acc fname -> longest_common_string acc (String.to_list fname))
    |> String.of_char_list

  let build_exp_tests env =
    let f fnames =
      let open FilePath in
      let basename = get_basename fnames in
      basename >:: fun test_ctxt ->
      let dir = env.tests_dir test_ctxt in
      let input_files =
        List.map fnames ~f:(fun fname -> make_filename (test_path fname))
      in
      let init_file = make_filename (test_path (basename ^ ".json")) in
      (* Verify standard output of execution with gold file *)
      let goldoutput_file = make_filename (gold_path dir basename) in
      let additional_dirs = List.map ~f:make_filename additional_libdirs in
      let stdlib = make_relative dir (env.stdlib_dir test_ctxt) in
      let path = string_of_path @@ (stdlib :: additional_dirs) in
      let args' =
        if ignore_predef_args then custom_args @ input_files
        else
          custom_args
          @ [ "-libdir"; path; "-gaslimit"; Stdint.Uint64.to_string gas_limit ]
          @ input_files
      in
      let args =
        if provide_init_arg then args' @ [ "-init"; init_file ] else args'
      in
      let msg = cli_usage_on_err runner args in
      print_cli_usage (env.print_cli test_ctxt) runner args;
      assert_command
        ~foutput:(fun s ->
          let actual_output = stream_to_string s in
          if env.update_gold test_ctxt then
            output_updater goldoutput_file basename actual_output
              ~json_errors:false
          else if Sys_unix.file_exists_exn goldoutput_file then
            let gold_output =
              (* load all data from file *)
              let non_normalized_gold_output =
                In_channel.read_all goldoutput_file
              in
              non_normalized_gold_output
            in
            output_verifier (diff_filter gold_output) msg
              (env.print_diff test_ctxt)
              (diff_filter actual_output)
          else
            assert_failure
              ("The gold file " ^ goldoutput_file ^ " does not exist"))
        ~exit_code ~use_stderr:true ~chdir:dir ~ctxt:test_ctxt runner args
    in
    List.map ~f

  let tests env = "exptests" >::: build_exp_tests env tests
end

type json_test = {
  json_test_name : string;
  json_tests : (string * string) list;
}

module type TestSuiteJSONInput = sig
  val tests : json_test list
  val gold_path : string -> string -> string list
  val test_path : string -> string list
  val runner : string
  val ignore_predef_args : bool
  val exit_code : UnixLabels.process_status
  val additional_libdirs : string list list
  val gas_limit : Stdint.uint64
  val custom_args : string list
  val diff_filter : string -> string
end

module DiffBasedMergeJSONTests (Input : TestSuiteJSONInput) = struct
  open Input
  open OUnit2

  let rec longest_common_string lhs rhs =
    let cmp xs ys = if List.length xs > List.length ys then xs else ys in
    match (lhs, rhs) with
    | [], _ | _, [] -> []
    | l :: ls, r :: rs ->
        if Char.equal l r then l :: longest_common_string ls rs
        else cmp (longest_common_string lhs rs) (longest_common_string ls rhs)

  let build_exp_tests env =
    let f json_test =
      let open FilePath in
      json_test.json_test_name >:: fun test_ctxt ->
      let dir = env.tests_dir test_ctxt in
      let input_json =
        List.fold_left json_test.json_tests ~init:[]
          ~f:(fun acc (fname, test) ->
            acc @ [ "--json"; make_filename (test_path fname); test ])
      in
      (* Verify standard output of execution with gold file *)
      let goldoutput_file =
        make_filename (gold_path dir json_test.json_test_name)
      in
      let additional_dirs = List.map ~f:make_filename additional_libdirs in
      let stdlib = make_relative dir (env.stdlib_dir test_ctxt) in
      let path = string_of_path @@ (stdlib :: additional_dirs) in
      let args =
        if ignore_predef_args then custom_args @ input_json
        else
          custom_args
          @ [ "-libdir"; path; "-gaslimit"; Stdint.Uint64.to_string gas_limit ]
          @ input_json
      in
      let msg = cli_usage_on_err runner args in
      print_cli_usage (env.print_cli test_ctxt) runner args;
      assert_command
        ~foutput:(fun s ->
          let actual_output = stream_to_string s in
          if env.update_gold test_ctxt then
            output_updater goldoutput_file json_test.json_test_name
              actual_output ~json_errors:false
          else if Sys_unix.file_exists_exn goldoutput_file then
            let gold_output =
              (* load all data from file *)
              let non_normalized_gold_output =
                In_channel.read_all goldoutput_file
              in
              non_normalized_gold_output
            in
            output_verifier (diff_filter gold_output) msg
              (env.print_diff test_ctxt)
              (diff_filter actual_output)
          else
            assert_failure
              ("The gold file " ^ goldoutput_file ^ " does not exist"))
        ~exit_code ~use_stderr:true ~chdir:dir ~ctxt:test_ctxt runner args
    in
    List.map ~f

  let tests env = "exptests" >::: build_exp_tests env tests
end
