open Ppxlib

let log_expander ~ctxt (e : expression) (l : (arg_label * expression) list) : expression =
  let loc = Expansion_context.Extension.extension_point_loc ctxt in
  let open (val Ast_builder.make loc) in
  let (is_exn, err) = begin match e.pexp_desc with
  | Pexp_ident { txt = Lident "trace"; _ }
  | Pexp_ident { txt = Lident "debug"; _ }
  | Pexp_ident { txt = Lident "info"; _ }
  | Pexp_ident { txt = Lident "warn"; _ }
  | Pexp_ident { txt = Lident "error"; _ } -> (false, None)
  | Pexp_ident { txt = Lident "exn"; _ } -> (true, None)
  | _ ->
      (false,
       Some (pexp_extension @@ Location.error_extensionf ~loc
             "log type must be one of: trace, debug, info, warn, error, exn"))
  end in match err with
  | Some err -> err
  | None ->
    if is_exn then
      match l with
      | [] ->
        pexp_extension @@ Location.error_extensionf ~loc
          "must provide exn value"
      | (_, v)::l ->
          let application =
            eapply (pexp_ident { txt = Lident "m"; loc}) (List.map snd l) in
          [%expr [%e e] [%e v] (Printexc.get_raw_backtrace ()) ~__POS__ ~__FUNCTION__
                        (fun m -> [%e application])]
    else
      let application =
        eapply (pexp_ident { txt = Lident "m"; loc}) (List.map snd l) in
      [%expr [%e e] ~__POS__ ~__FUNCTION__ (fun m -> [%e application])]

let log_extension =
  let pattern =
    let open Ast_pattern in
    pstr (pstr_eval (pexp_apply __ __) nil ^:: nil)
  in
  Context_free.Rule.extension
    (Extension.V3.declare "xlog" Expression pattern log_expander)

let import_expander ~ctxt _e : structure_item =
  let loc = Expansion_context.Extension.extension_point_loc ctxt in
  [%stri include (val Xlog.logs __FUNCTION__)]

let import_extension =
  let pattern =
    let open Ast_pattern in
    __
  in
  Context_free.Rule.extension
    (Extension.V3.declare "xlog_import" Structure_item pattern import_expander)

let () =
  Driver.register_transformation
    ~rules:[ log_extension; import_extension ]
    "xlog.ppx"