open Why3
open Why3.Json_base

type prover_stat =
  {
    mutable count     : int;
    mutable max_time  : float;
    mutable max_steps : int;
  }

type stats = prover_stat Whyconf.Hprover.t

type result_info =
  | Proved of stats * int
  | Not_Proved of
       Gnat_expl.extra_info *
       (Model_parser.model * Check_ce.rac_result option) list *
       (string * string) option *
        Gnat_expl.unproved_status

type msg =
  { result           : bool;
    stats            : stats option;
    stats_checker    : int;
    check_tree       : Json_base.json;
    extra_info       : Gnat_expl.extra_info;
    cntexmp_models   : (Model_parser.model * Check_ce.rac_result option) list;
    manual_proof     : (string * string) option;
    unproved_status : Gnat_expl.unproved_status
  }

let msg_set : msg Gnat_expl.HCheck.t = Gnat_expl.HCheck.create 17

let warning_list : string list ref = ref []

let add_warning ?loc s =
  let s =
    match loc with
    | Some loc ->
        let file, line1, col1, line2, col2 = Loc.get loc in
        Format.sprintf "%s:(%d:%d)-(%d:%d): %s" file line1 col1 line2 col2 s
    | None -> s
  in
  warning_list := s :: !warning_list

let () = Loc.set_warning_hook add_warning

let adapt_stats statsopt =
  match statsopt with
  | None -> None
  | Some stats ->
      let newstats = Whyconf.Hprover.create 3 in
      Whyconf.Hprover.iter (fun k v ->
        let name = k.Whyconf.prover_name in
        Whyconf.Hprover.add newstats k {v with max_steps =
          Gnat_config.back_convert_steps ~prover:name v.max_steps}) stats;
      Some newstats

let register check check_tree result =
  let valid, extra_info, stats, models, manual, unproved_status =
    match result with
    | Proved (stats, stats_checker) ->
        true,
        {Gnat_expl.pretty_node = None; inlined = None},
        Some (stats, stats_checker),
        [],
        None,
        Gnat_expl.Unknown
    | Not_Proved (extra_info, models, manual, unproved_status) ->
        false, extra_info, None, models, manual, unproved_status
  in
  if (Gnat_expl.HCheck.mem msg_set check) then assert false
  else begin
    let stats_g = Option.map (fun x -> let (stats, _) = x in stats) stats in
    let stats_checker = Option.map (fun x -> let (_, sc) = x in sc) stats in
    let msg =
    { result         = valid;
      extra_info     = extra_info;
      stats          = adapt_stats stats_g;
      stats_checker  = Option.value ~default:0 stats_checker;
      check_tree     = check_tree;
      cntexmp_models = models;
      manual_proof   = manual;
      unproved_status = unproved_status } in
    Gnat_expl.HCheck.add msg_set check msg
  end

let get_cntexmp_models =
  List.map (function
  (m, r) ->
      let json_result_state (state, _log) =
        let more = match state with
          | Check_ce.Res_fail (ctx, t) ->
              let t =
                Ident.Sattr.fold Term.t_attr_add ctx.Pinterp_core.attrs t in
              let check = match Gnat_expl.search_labels t with
                | Some Gnat_expl.{check={id; check_kind; sloc}} ->
                    let file, line, col =
                      Gnat_loc.explode (Gnat_loc.orig_loc sloc) in
                    Record [
                      "id", Int id;
                      "vc_kind", String (Gnat_expl.check_kind_to_ada check_kind);
                      "sloc", Record [
                        "file", String file;
                        "line", Int line;
                        "col", Int col;
                      ]
                    ]
                | None -> Null in
              ["check", check]
          | Check_ce.Res_stuck reason | Check_ce.Res_incomplete reason ->
              ["reason", String reason]
          | _ -> [] in
        let state =
          ["state", String (Check_ce.string_of_rac_result_state state)] in
        Record (state @ more) in
      let rcd_only_cntexmp_model = [ "cntexmp", (Model_parser.json_model m) ] in
      match r with
      | None -> Record rcd_only_cntexmp_model
      | Some (Check_ce.RAC_not_done _) -> Record rcd_only_cntexmp_model
      | Some (Check_ce.RAC_done (state, _log)) ->
          Record (List.append
            rcd_only_cntexmp_model
            [ "giant_step_rac_result", (json_result_state (state, _log)) ])
    )

let get_manual_proof_info info =
  match info with
  | None -> []
  | Some (fname, cmd) ->
      [ "vc_file", String (Sys.getcwd () ^ Filename.dir_sep ^ fname);
        "editor_cmd", String cmd
      ]

let get_prover_stats stat =
  Record [
    "count", Int stat.count;
    "max_steps", Int stat.max_steps;
    "max_time", StandardFloat stat.max_time
  ]

let get_stats (stats, stat_checker) =
  (* Compute the stats for goal solved by transformations as a fake prover. *)
  let fake_prover_ch =
    Whyconf.{prover_name = "Checker"; prover_version = ""; prover_altern = ""}
  in
  match stats with
  | None -> []
  | Some s ->
    if stat_checker <> 0 then
      Whyconf.Hprover.add s fake_prover_ch
        {count = stat_checker; max_steps = 0; max_time = 0.0};
    let kv_list = Whyconf.Hprover.fold (fun k v acc -> (k,v)::acc) s [] in
    let get_name pr = pr.Whyconf.prover_name in
    [ "stats",
      Record (List.map (fun (k,v) -> get_name k, get_prover_stats v) kv_list) ]

let opt_int i = Option.value ~default:0 i

let get_extra_info i =
  match i with
  | { Gnat_expl.pretty_node = i; inlined = inline } ->
      Record [
        "node", Int (opt_int i);
        "inline", Int (opt_int inline)
      ]

let get_unproved_status status =
  let status, time, steps, memory =
    match status with
    | Gnat_expl.Unknown ->
        "unknown", false, false, false
    | Gnat_expl.Gave_up ->
        "gave_up", false, false, false
    | Gnat_expl.Limit { timeout; step; memory } ->
        "limit", timeout, step, memory
  in
  Record [
    "status", String status;
    "time", Bool time;
    "steps", Bool steps;
    "memory", Bool memory
  ]

let get_msg (check, m) =
  Record (
    List.concat [
      [ "id", Int check.Gnat_expl.id ;
        "check_kind", String (Gnat_expl.check_kind_to_ada check.Gnat_expl.check_kind) ;
        "result", Bool m.result ;
        "check_tree", m.check_tree ;
        "extra_info", get_extra_info m.extra_info ;
        "unproved_status", get_unproved_status m.unproved_status
      ] ;
      get_stats (m.stats, m.stats_checker);
      ["cntexmps", List (get_cntexmp_models m.cntexmp_models)];
      get_manual_proof_info m.manual_proof
    ]
  )

(* In case several models are returned by Why3, we sort them by id. *)
let sort_messages (l : (Gnat_expl.check * msg) list) =
  List.sort (fun x y -> compare (fst x).Gnat_expl.id (fst y).Gnat_expl.id) l

let get_results () =
  let msg_set =
    Gnat_expl.HCheck.fold (fun k v acc -> (k,v) :: acc) msg_set []
  in
  "results", List (List.map get_msg (sort_messages msg_set))

let get_warnings () =
  match !warning_list with
  | [] -> []
  | l -> [("warnings", List (List.map (fun elem -> String elem) l))]

let get_timings () =
  let l = Hashtbl.fold (fun k v acc -> (k,v)::acc) (Debug.Stats.get_timings ()) [] in
  let get_name s = s in
  "timings", Record (List.map (fun (k,(v,_n)) -> get_name k, StandardFloat v) l)

let get_entity () =
  "entity", Int Gnat_config.entity

(* The main function, where we build the final Record to be printed using
Why3.Json_base.print_json function. *)
let print_messages () =
  print_json_single_line Format.std_formatter (
    Record (
      List.concat [
        [get_results (); get_timings (); get_entity ()];
        get_warnings ()
      ]
    )
  )
