open Vanity

type profile = {
  name : string;
  flags : Pipeline.optimisation_flags;
  relation : Relation.relation;
}

type output_mode =
  | Text
  | Json

let safe_flags =
  { Pipeline.default_flags with unsafe_repr_eq = false; unsafe_strict_unroll = false; unsafe_effect_drop = false }

let profiles =
  [
    { name = "safe"; flags = safe_flags; relation = Relation.Boxed_unboxed };
    { name = "specialise-only"; flags = { safe_flags with inline = false; repr_lower = false }; relation = Relation.Baseline };
    { name = "inline-no-repr-leak"; flags = { safe_flags with inline = true }; relation = Relation.Boxed_unboxed };
    { name = "unsafe-inline"; flags = { safe_flags with unsafe_repr_eq = true }; relation = Relation.Boxed_unboxed };
    { name = "unsafe-unbox"; flags = Pipeline.default_flags; relation = Relation.Boxed_unboxed };
    { name = "unsafe-strictness"; flags = { safe_flags with inline = false; unsafe_strict_unroll = true }; relation = Relation.Boxed_unboxed };
    { name = "unsafe-effects"; flags = { safe_flags with unsafe_effect_drop = true }; relation = Relation.Boxed_unboxed };
  ]

let json_escape s =
  let b = Buffer.create (String.length s + 16) in
  String.iter
    (function
      | '"' -> Buffer.add_string b "\\\""
      | '\\' -> Buffer.add_string b "\\\\"
      | '\n' -> Buffer.add_string b "\\n"
      | '\r' -> Buffer.add_string b "\\r"
      | '\t' -> Buffer.add_string b "\\t"
      | c -> Buffer.add_char b c)
    s;
  Buffer.contents b

let json_string s = "\"" ^ json_escape s ^ "\""
let json_bool b = if b then "true" else "false"
let json_list xs = "[" ^ String.concat "," xs ^ "]"

let json_field name value =
  json_string name ^ ":" ^ value

let json_object fields =
  "{" ^ String.concat "," fields ^ "}"

let option_default all selected =
  match selected with
  | [] -> all
  | _ -> List.rev selected

let find_named kind name xs get_name =
  match List.find_opt (fun x -> String.equal (get_name x) name) xs with
  | Some x -> x
  | None ->
      prerr_endline ("unknown " ^ kind ^ " " ^ name);
      exit 2

let selected_cases names =
  option_default Corpus.all (List.map (fun name -> find_named "case" name Corpus.all (fun c -> c.Corpus.name)) names)

let selected_profiles names =
  option_default profiles (List.map (fun name -> find_named "profile" name profiles (fun p -> p.name)) names)

let failed_invariants audit =
  audit.Audit.invariants
  |> List.filter (fun result -> not result.Audit.passed)
  |> List.map (fun result -> Audit.invariant_kind_to_string result.Audit.kind)

let emit_plain_obligations audit =
  if audit.Audit.obligations = [] then "obligations none"
  else
    "obligations\n" ^
    (audit.Audit.obligations
     |> List.map (fun result ->
            let o = result.Audit.obligation in
            o.Pipeline.pass ^ "  " ^
            Reporting.obligation_kind_to_string o.Pipeline.kind ^ "  " ^
            Audit.obligation_status_to_string result.Audit.status ^ "  " ^
            result.Audit.note)
     |> String.concat "\n")

let emit_plain_invariants audit =
  "invariants\n" ^
  (audit.Audit.invariants
   |> List.map (fun result ->
          Audit.invariant_kind_to_string result.Audit.kind ^ "  " ^
          (if result.Audit.passed then "ok" else "failed") ^ "  " ^
          result.Audit.detail)
   |> String.concat "\n")

let emit_plain_boundaries audit =
  "pass observations\n" ^
  (audit.Audit.pass_boundaries
   |> List.map (fun (boundary : Audit.pass_boundary) -> boundary.Audit.pass ^ "  " ^ boundary.Audit.verdict)
   |> String.concat "\n")

let run_case_text ~details case =
  let audit = Audit.audit_case case in
  let failed = failed_invariants audit in
  Printf.printf
    "case %s\nclassification %s\nverdict %s\nsource %s\ntarget %s\nfailed invariants %s\n"
    case.Corpus.name
    (Audit.failure_mode_to_string audit.Audit.failure_mode)
    (Reporting.verdict_to_string audit.Audit.comparison.Relation.verdict)
    (Reporting.string_of_source_trace audit.Audit.source_trace)
    (Reporting.string_of_target_trace audit.Audit.target_trace)
    (if failed = [] then "none" else String.concat ", " failed);
  if details || audit.Audit.failure_mode <> Audit.Preserved then begin
    Printf.printf "\n%s\n\n" (Reporting.emit_pipeline_visualisation audit.Audit.compiled);
    Printf.printf "%s\n\n" (emit_plain_boundaries audit);
    Printf.printf "%s\n\n" (emit_plain_invariants audit);
    Printf.printf "%s\n" (emit_plain_obligations audit)
  end;
  Printf.printf "\n"

let outcome_json_source outcome =
  json_string (Reporting.string_of_source_outcome outcome)

let outcome_json_target outcome =
  json_string (Reporting.string_of_target_outcome outcome)

let obligation_json result =
  let o = result.Audit.obligation in
  json_object
    [
      json_field "pass" (json_string o.Pipeline.pass);
      json_field "kind" (json_string (Reporting.obligation_kind_to_string o.Pipeline.kind));
      json_field "status" (json_string (Audit.obligation_status_to_string result.Audit.status));
      json_field "subject" (json_string o.Pipeline.subject);
      json_field "detail" (json_string o.Pipeline.detail);
      json_field "note" (json_string result.Audit.note);
    ]

let invariant_json result =
  json_object
    [
      json_field "kind" (json_string (Audit.invariant_kind_to_string result.Audit.kind));
      json_field "passed" (json_bool result.Audit.passed);
      json_field "detail" (json_string result.Audit.detail);
    ]

let boundary_json (boundary : Audit.pass_boundary) =
  json_object
    [
      json_field "pass" (json_string boundary.Audit.pass);
      json_field "observation" (json_string boundary.Audit.verdict);
      json_field "effects" (string_of_int boundary.Audit.effects);
    ]

let audit_json audit =
  json_object
    [
      json_field "case" (json_string audit.Audit.case.Corpus.name);
      json_field "summary" (json_string audit.Audit.case.Corpus.summary);
      json_field "classification" (json_string (Audit.failure_mode_to_string audit.Audit.failure_mode));
      json_field "verdict" (json_string (Reporting.verdict_to_string audit.Audit.comparison.Relation.verdict));
      json_field "source" (json_string (Reporting.string_of_source_trace audit.Audit.source_trace));
      json_field "target" (json_string (Reporting.string_of_target_trace audit.Audit.target_trace));
      json_field "pipeline"
        (json_object
           [
             json_field "source" (json_string (Source.string_of_term audit.Audit.compiled.Pipeline.source_program));
             json_field "specialise" (json_string (Source.string_of_term audit.Audit.compiled.Pipeline.specialised));
             json_field "inline" (json_string (Source.string_of_term audit.Audit.compiled.Pipeline.inlined));
             json_field "effect" (json_string (Source.string_of_term audit.Audit.compiled.Pipeline.effect_rewritten));
             json_field "repr_lower" (json_string (Target.string_of_term audit.Audit.compiled.Pipeline.target_program));
           ]);
      json_field "pass_boundaries" (json_list (List.map boundary_json audit.Audit.pass_boundaries));
      json_field "invariants" (json_list (List.map invariant_json audit.Audit.invariants));
      json_field "obligations" (json_list (List.map obligation_json audit.Audit.obligations));
    ]

let run_case_json case =
  audit_json (Audit.audit_case case)

let failure_json failure =
  json_object
    [
      json_field "original_type" (json_string (Types.string_of_typ failure.Gen.original.Gen.ty));
      json_field "original_term" (json_string (Source.string_of_term failure.Gen.original.Gen.term));
      json_field "minimal_type" (json_string (Types.string_of_typ failure.Gen.minimal.Gen.ty));
      json_field "minimal_term" (json_string (Source.string_of_term failure.Gen.minimal.Gen.term));
      json_field "verdict" (json_string (Reporting.verdict_to_string failure.Gen.comparison.Relation.verdict));
      json_field "source_effects" (string_of_int failure.Gen.comparison.Relation.source_effects);
      json_field "target_effects" (string_of_int failure.Gen.comparison.Relation.target_effects);
    ]

let run_profile_text ~details ~count ~max_depth profile =
  let result = Gen.run_campaign profile.flags profile.relation ~count ~max_depth () in
  Printf.printf
    "profile %s\nrelation %s\nrelated %d/%d\nviolations %d\n"
    profile.name
    (Reporting.string_of_relation profile.relation)
    result.Gen.related
    result.Gen.total
    (result.Gen.total - result.Gen.related);
  if details && result.Gen.failures <> [] then begin
    Printf.printf "minimal failures\n";
    result.Gen.failures
    |> List.to_seq
    |> Seq.take 5
    |> Seq.iter (fun failure ->
           Printf.printf
             "type %s\nterm %s\nverdict %s\n\n"
             (Types.string_of_typ failure.Gen.minimal.Gen.ty)
             (Source.string_of_term failure.Gen.minimal.Gen.term)
             (Reporting.verdict_to_string failure.Gen.comparison.Relation.verdict))
  end;
  Printf.printf "\n"

let run_profile_json ~count ~max_depth profile =
  let result = Gen.run_campaign profile.flags profile.relation ~count ~max_depth () in
  json_object
    [
      json_field "profile" (json_string profile.name);
      json_field "relation" (json_string (Reporting.string_of_relation profile.relation));
      json_field "total" (string_of_int result.Gen.total);
      json_field "related" (string_of_int result.Gen.related);
      json_field "violations" (string_of_int (result.Gen.total - result.Gen.related));
      json_field "failures" (json_list (List.map failure_json result.Gen.failures));
    ]

let list_available () =
  print_endline "cases";
  List.iter (fun case -> print_endline ("  " ^ case.Corpus.name)) Corpus.all;
  print_endline "";
  print_endline "profiles";
  List.iter (fun profile -> print_endline ("  " ^ profile.name)) profiles

let () =
  let mode = ref Text in
  let details = ref false in
  let list_only = ref false in
  let run_corpus = ref true in
  let run_profiles = ref true in
  let case_names = ref [] in
  let profile_names = ref [] in
  let count = ref 160 in
  let max_depth = ref 4 in
  let add_case name = case_names := name :: !case_names in
  let add_profile name = profile_names := name :: !profile_names in
  let specs =
    [
      ("--case", Arg.String add_case, "run one corpus case by name; repeatable");
      ("--profile", Arg.String add_profile, "run one campaign profile by name; repeatable");
      ("--cases", Arg.Clear run_profiles, "run selected corpus cases without campaign profiles");
      ("--profiles", Arg.Clear run_corpus, "run selected campaign profiles without corpus cases");
      ("--count", Arg.Set_int count, "set generated campaign specimen count");
      ("--max-depth", Arg.Set_int max_depth, "set generated campaign term/type depth");
      ("--details", Arg.Set details, "print pass observations, invariants, obligations, and shrunk failures");
      ("--json", Arg.Unit (fun () -> mode := Json), "emit machine-readable JSON");
      ("--list", Arg.Set list_only, "list available cases and profiles");
    ]
  in
  Arg.parse specs (fun name -> add_case name) "vanity [--case name] [--profile name] [--details] [--json]";
  if !list_only then list_available ()
  else
    let cases = if !run_corpus then selected_cases !case_names else [] in
    let selected_profiles = if !run_profiles then selected_profiles !profile_names else [] in
    match !mode with
    | Text ->
        if !details then Printf.printf "%s\n\n" (Audit.emit_relation_ledger ());
        if cases <> [] then begin
          Printf.printf "corpus\n\n";
          List.iter (run_case_text ~details:!details) cases
        end;
        if selected_profiles <> [] then begin
          Printf.printf "profiles\n\n";
          List.iter (run_profile_text ~details:!details ~count:!count ~max_depth:!max_depth) selected_profiles
        end
    | Json ->
        let case_json = List.map run_case_json cases in
        let profile_json = List.map (run_profile_json ~count:!count ~max_depth:!max_depth) selected_profiles in
        print_endline
          (json_object
             [
               json_field "version" (json_string Project.version);
               json_field "relation_ledger"
                 (json_list
                    (List.map
                       (fun entry ->
                          json_object
                            [
                              json_field "pass" (json_string entry.Audit.pass);
                              json_field "relation" (json_string entry.Audit.relation);
                              json_field "condition" (json_string entry.Audit.condition);
                            ])
                       Audit.relation_ledger));
               json_field "cases" (json_list case_json);
               json_field "profiles" (json_list profile_json);
             ])