(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2022                                                    *)
(*    CEA (Commissariat à l'énergie atomique et aux énergies              *)
(*         alternatives)                                                  *)
(*                                                                        *)
(*  You can redistribute it and/or modify it under the terms of the GNU   *)
(*  Lesser General Public License as published by the Free Software       *)
(*  Foundation, version 2.1.                                              *)
(*                                                                        *)
(*  It 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 Lesser General Public License for more details.                   *)
(*                                                                        *)
(*  See the GNU Lesser General Public License version 2.1                 *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).            *)
(*                                                                        *)
(**************************************************************************)

open Base
open Cmdliner

let caisar = "caisar"

let () =
  Nn2smt.init ();
  Vars_on_lhs.init ()

let () =
  Pyrat.init ();
  Marabou.init ();
  Vnnlib.init ()

(* -- Logs. *)

let pp_header =
  let x =
    match Array.length Caml.Sys.argv with
    | 0 -> Caml.(Filename.basename Sys.executable_name)
    | _ -> Caml.(Filename.basename Sys.argv.(0))
  in
  let pp_h ppf style h = Fmt.pf ppf "[%s][%a] " x Fmt.(styled style string) h in
  fun ppf (l, h) ->
    let open Logs_fmt in
    match l with
    | Logs.App -> Fmt.pf ppf "[%a] " Fmt.(styled app_style string) x
    | Logs.Error ->
      pp_h ppf err_style (match h with None -> "ERROR" | Some h -> h)
    | Logs.Warning ->
      pp_h ppf warn_style (match h with None -> "WARNING" | Some h -> h)
    | Logs.Info ->
      pp_h ppf info_style (match h with None -> "INFO" | Some h -> h)
    | Logs.Debug ->
      pp_h ppf debug_style (match h with None -> "DEBUG" | Some h -> h)

let setup_logs =
  let setup_log level =
    Fmt_tty.setup_std_outputs ~style_renderer:`Ansi_tty ();
    Logs.set_level level;
    Logs.set_reporter (Logs_fmt.reporter ~pp_header ())
  in
  Term.(const setup_log $ Logs_cli.level ())

let log_level_is_debug () =
  match Logs.level () with Some Debug -> true | _ -> false

(* -- Commands. *)

let config detect () =
  if detect
  then (
    Logs.debug (fun m -> m "Automatic detection");
    let config =
      let debug = log_level_is_debug () in
      Autodetect.autodetection ~debug ()
    in
    let open Why3 in
    let provers = Whyconf.get_provers config in
    if not (Whyconf.Mprover.is_empty provers)
    then
      Logs.app (fun m ->
        m "@[<v>%a@]"
          (Pp.print_iter2 Whyconf.Mprover.iter Pp.newline Pp.nothing
             Whyconf.print_prover Pp.nothing)
          provers))

let memlimit_of_string s =
  let s = String.strip s in
  let value =
    let re = Re__Pcre.regexp "[0-9]+" in
    Re__Core.matches re s
  in
  let unit =
    let re = Re__Pcre.regexp "MB?|GB?|TB?" in
    Re__Core.matches re s
  in
  match (value, unit) with
  | [ v ], ([] | [ "M" ] | [ "MB" ]) -> Int.of_string v
  | [ v ], ([ "G" ] | [ "GB" ]) -> Int.of_string v * 1000
  | [ v ], ([ "T" ] | [ "TB" ]) -> Int.of_string v * 1000000
  | _ -> invalid_arg "Unrecognized memory limit"

let timelimit_of_string s =
  let s = String.strip s in
  let value =
    let re = Re__Pcre.regexp "[0-9]+" in
    Re__Core.matches re s
  in
  let unit =
    let re = Re__Pcre.regexp "[a-z]" in
    Re__Core.matches re s
  in
  match (value, unit) with
  | [ v ], ([] | [ "s" ]) -> Int.of_string v
  | [ v ], [ "m" ] -> Int.of_string v * 60
  | [ v ], [ "h" ] -> Int.of_string v * 3600
  | _ -> invalid_arg "Unrecognized time limit"

let verify ?format ~loadpath ?memlimit ?timelimit prover ?dataset_csv files =
  let debug = log_level_is_debug () in
  let memlimit = Option.map memlimit ~f:memlimit_of_string in
  let timelimit = Option.map timelimit ~f:timelimit_of_string in
  List.iter
    ~f:
      (Verification.verify ~debug ?format ~loadpath ?memlimit ?timelimit prover
         ?dataset_csv)
    files

let verify_json ?memlimit ?timelimit file =
  let jc = Result.ok_exn (Verification.JSON.of_string file) in
  let prover = jc.Verification.JSON.prover in
  let dataset_csv = jc.Verification.JSON.property.dataset in
  let file = Result.ok_or_failwith (Verification.File.of_json_config jc) in
  verify ~loadpath:[] ?memlimit ?timelimit prover ~dataset_csv [ file ]

let exec_cmd cmdname cmd =
  Logs.debug (fun m -> m "Command `%s'" cmdname);
  cmd ()

(* -- Command line. *)

let config_cmd =
  let cmdname = "config" in
  let info =
    let doc = Fmt.str "%s configuration." caisar in
    let exits = Cmd.Exit.defaults in
    let man =
      [
        `S Manpage.s_description;
        `P (Fmt.str "Handle the configuration of %s." caisar);
      ]
    in
    Cmd.info cmdname ~sdocs:Manpage.s_common_options ~exits ~doc ~man
  in
  let term =
    let detect =
      let doc = "Detect solvers in \\$PATH." in
      Arg.(value & flag & info [ "d"; "detect" ] ~doc)
    in
    Term.(
      ret
        (const (fun detect _ ->
           if not detect
           then `Help (`Pager, Some cmdname)
           else
             (* TODO: Do not only check for [detect], and enable it by default,
                as soon as other options are available. *)
             `Ok (exec_cmd cmdname (fun () -> config true ())))
        $ detect $ setup_logs))
  in
  Cmd.v info term

let memlimit =
  let doc =
    "Memory limit. $(docv) is intended in megabytes (MB) by default, but \
     gigabytes (GB) and terabytes (TB) may also be specified instead."
  in
  Arg.(
    value & opt (some string) None & info [ "m"; "memlimit" ] ~doc ~docv:"MEM")

let timelimit =
  let doc =
    "Time limit. $(docv) is intended in seconds (s) by default, but minutes \
     (m) and hours (h) may also be specified instead."
  in
  Arg.(
    value & opt (some string) None & info [ "t"; "timelimit" ] ~doc ~docv:"TIME")

let verify_cmd =
  let cmdname = "verify" in
  let info =
    let doc = "Property verification using external provers." in
    Cmd.info cmdname ~sdocs:Manpage.s_common_options ~exits:Cmd.Exit.defaults
      ~doc
      ~man:[ `S Manpage.s_description; `P doc ]
  in
  let loadpath =
    let doc = "Additional loadpath." in
    Arg.(value & opt_all string [ "." ] & info [ "L"; "loadpath" ] ~doc)
  in
  let format =
    let doc = "File format." in
    Arg.(value & opt (some string) None & info [ "format" ] ~doc)
  in
  let term =
    let files =
      let doc = "File to verify." in
      let file_or_stdin = Arg.conv' Verification.File.(of_string, pretty) in
      Arg.(non_empty & pos_all file_or_stdin [] & info [] ~doc ~docv:"FILE")
    in
    let prover =
      let all_provers = Prover.list_available () in
      let doc =
        Fmt.str
          "Prover to use. Support is provided for the following provers: %a."
          (Fmt.list ~sep:Fmt.comma Fmt.string)
          (List.map ~f:Prover.to_string all_provers)
      in
      let provers =
        Arg.enum (List.map ~f:(fun p -> (Prover.to_string p, p)) all_provers)
      in
      Arg.(required & opt (some provers) None & info [ "p"; "prover" ] ~doc)
    in
    let dataset_csv =
      let doc = "Dataset in CSV format." in
      Arg.(value & opt (some file) None & info [ "dataset-csv" ] ~doc)
    in
    Term.(
      const
        (fun format loadpath memlimit timelimit prover dataset_csv files _ ->
        exec_cmd cmdname (fun () ->
          verify ?format ~loadpath ?memlimit ?timelimit prover ?dataset_csv
            files))
      $ format $ loadpath $ memlimit $ timelimit $ prover $ dataset_csv $ files
      $ setup_logs)
  in
  Cmd.v info term

let verify_json_cmd =
  let cmdname = "verify-json" in
  let info =
    let doc =
      "Property verification using external provers via json configuration \
       file."
    in
    Cmd.info cmdname ~sdocs:Manpage.s_common_options ~exits:Cmd.Exit.defaults
      ~doc
      ~man:[ `S Manpage.s_description; `P doc ]
  in
  let term =
    let json =
      let doc = "JSON file." in
      Arg.(
        required & pos ~rev:true 0 (some file) None & info [] ~doc ~docv:"FILE")
    in
    Term.(
      const (fun memlimit timelimit json _ ->
        exec_cmd cmdname (fun () -> verify_json ?memlimit ?timelimit json))
      $ memlimit $ timelimit $ json $ setup_logs)
  in
  Cmd.v info term

let default_info =
  let doc =
    "A platform for characterizing the safety and robustness of artificial \
     intelligence based software."
  in
  let sdocs = Manpage.s_common_options in
  let man =
    [
      `S Manpage.s_common_options;
      `P "These options are common to all commands.";
      `S "MORE HELP";
      `P "Use `$(mname) $(i,COMMAND) --help' for help on a single command.";
      `S Manpage.s_bugs;
      `P "Email bug reports to <...>";
    ]
  in
  let version = "0.1" in
  let exits = Cmd.Exit.defaults in
  Cmd.info caisar ~version ~doc ~sdocs ~exits ~man

let default_cmd = Term.(ret (const (fun _ -> `Help (`Pager, None)) $ const ()))

let () =
  (* We register custom printers for a selection of exceptions otherwise Why3
     will print the related messages as "anomaly: <exception with message>". *)
  Why3.Exn_printer.register (fun fmt exn ->
    match exn with
    | Invalid_argument msg -> Fmt.pf fmt "Invalid argument: %s" msg
    | Failure msg -> Fmt.pf fmt "Failure: %s" msg
    | _ -> raise exn)

let () =
  try
    Cmd.group ~default:default_cmd default_info
      [ config_cmd; verify_cmd; verify_json_cmd ]
    |> Cmd.eval ~catch:false |> Caml.exit
  with exn when not (log_level_is_debug ()) ->
    Logs.err (fun m -> m "@[%a@]" Why3.Exn_printer.exn_printer exn)