open Ppxlib

let expand_test_pred ~loc:_ ~path:_ typ =
  let loc = { typ.ptyp_loc with loc_ghost = true } in
  [%expr
    fun ?(here = []) ?message predicate t ->
      let pos = [%e Ppx_here_expander.lift_position_as_string ~loc] in
      let sexpifier = [%e Ppx_sexp_conv_expander.Sexp_of.core_type typ] in
      Ppx_assert_lib.Runtime.test_pred ~pos ~sexpifier ~here ?message predicate t]
;;

let expand_test_eq ~loc:_ ~path:_ typ =
  let loc = { typ.ptyp_loc with loc_ghost = true } in
  [%expr
    fun ?(here = []) ?message ?equal t1 t2 ->
      let pos = [%e Ppx_here_expander.lift_position_as_string ~loc] in
      let sexpifier = [%e Ppx_sexp_conv_expander.Sexp_of.core_type typ] in
      let comparator =
        [%e
          Merlin_helpers.hide_expression
            (Ppx_compare_expander.Compare.core_type ~with_local:false typ)]
      in
      Ppx_assert_lib.Runtime.test_eq
        ~pos
        ~sexpifier
        ~comparator
        ~here
        ?message
        ?equal
        t1
        t2]
;;

let expand_test_result ~loc:_ ~path:_ typ =
  let loc = { typ.ptyp_loc with loc_ghost = true } in
  [%expr
    fun ?(here = []) ?message ?equal ~expect got ->
      let pos = [%e Ppx_here_expander.lift_position_as_string ~loc] in
      let sexpifier = [%e Ppx_sexp_conv_expander.Sexp_of.core_type typ] in
      let comparator =
        [%e
          Merlin_helpers.hide_expression
            (Ppx_compare_expander.Compare.core_type ~with_local:false typ)]
      in
      Ppx_assert_lib.Runtime.test_result
        ~pos
        ~sexpifier
        ~comparator
        ~here
        ?message
        ?equal
        ~expect
        ~got]
;;

let extensions =
  let declare name expand =
    [ Extension.declare name Extension.Context.expression Ast_pattern.(ptyp __) expand
    ; Extension.declare
        name
        Extension.Context.core_type
        Ast_pattern.(ptyp __)
        (fun ~loc ~path:_ ty ->
          let loc = { loc with loc_ghost = true } in
          let open Ast_builder.Default in
          let ident = Located.lident ~loc ("Ppx_assert_lib.Runtime." ^ name) in
          ptyp_constr ~loc ident [ ty ])
    ]
  in
  List.concat
    [ declare "test_pred" expand_test_pred
    ; declare "test_eq" expand_test_eq
    ; declare "test_result" expand_test_result
    ]
;;

let () = Driver.register_transformation "assert" ~extensions
