open Vanity

let assert_true msg b =
  if not b then failwith msg

let assert_target_valid msg term =
  match Target.validate term with
  | Ok () -> ()
  | Error errors -> failwith (msg ^ ": " ^ String.concat "; " errors)

let assert_target_invalid msg term =
  match Target.validate term with
  | Ok () -> failwith msg
  | Error _ -> ()

let assert_invalid_arg msg f =
  try
    let _ = f () in
    failwith msg
  with
  | Invalid_argument _ -> ()

let find_case name =
  match List.find_opt (fun (case : Corpus.case) -> String.equal case.name name) Corpus.all with
  | Some case -> case
  | None -> failwith ("missing case " ^ name)

let has_obligation kind obligations =
  List.exists (fun (o : Pipeline.obligation) -> o.kind = kind) obligations

let () =
  List.iter
    (fun (case : Corpus.case) ->
      assert_true
        ("ill-typed corpus case " ^ case.name)
        (Typecheck.is_well_typed case.ty case.source);
      let compiled = Pipeline.compile case.flags case.ty case.source in
      assert_target_valid
        ("invalid target IR for corpus case " ^ case.name)
        compiled.target_program)
    Corpus.all;
  let repr = Audit.audit_case (find_case "free-theorem-fails-after-unsafe-inlining") in
  assert_true
    "expected representation exposure witness"
    (repr.failure_mode = Audit.Representation_exposure);
  assert_true
    "expected unsafe inlining to emit representation exposure obligation"
    (has_obligation Pipeline.Exposed_representation repr.compiled.obligations);
  let safe = Audit.audit_case (find_case "safe-polymorphic-instantiation") in
  assert_true
    "safe instantiation should not emit representation exposure obligation"
    (not (has_obligation Pipeline.Exposed_representation safe.compiled.obligations));
  assert_true
    "safe instantiation should emit preserve relation obligations at rewrite sites"
    (has_obligation Pipeline.Preserve_relation safe.compiled.obligations);
  let strict = Audit.audit_case (find_case "strictness-induced-termination-change") in
  assert_true
    "expected strictness shift witness"
    (strict.failure_mode = Audit.Strictness_shift);
  let generated = Gen.sample_terms ~count:80 ~max_depth:4 () in
  List.iter
    (fun specimen ->
      assert_true
        ("ill-typed generated specimen " ^ Source.string_of_term specimen.Gen.term)
        (Typecheck.is_well_typed specimen.Gen.ty specimen.Gen.term))
    generated;
  assert_target_invalid
    "expected tuple projection arity validation failure"
    (Target.Proj (2, Target.Tuple [Target.Int 1; Target.Bool true]));
  assert_target_invalid
    "expected worker-wrapper arity validation failure"
    (Target.WorkerWrapper
       {
         wrapper = "wrapper";
         worker = "worker";
         boxed_arg = Types.TPair (Types.TInt, Types.TBool);
         unboxed_args = [Target.RInt];
         result_repr = Target.RTuple [Target.RInt; Target.RBool];
         wrap_body = Target.Tuple [Target.Int 0; Target.Bool true];
         worker_body = Target.Tuple [Target.Var "u0"; Target.Bool true];
         in_term = Target.Var "wrapper";
       });
  assert_invalid_arg
    "expected unsupported worker-wrapper source shape to be rejected"
    (fun () ->
       Pipeline.compile
         { Pipeline.default_flags with unsafe_repr_eq = false; unsafe_strict_unroll = false }
         (Types.TArrow (Types.TPair (Types.TInt, Types.TBool), Types.TPair (Types.TInt, Types.TBool)))
         (Source.Lam
            ( "p",
              Types.TPair (Types.TInt, Types.TBool),
              Types.TPair (Types.TInt, Types.TBool),
              Source.Pair (Source.Int 0, Source.Bool true) )))