include Types
module A = Atomic_
module Collector = Collector
module Meta_map = Meta_map
module Level = Level
module Core_ext = Core_ext

type collector = Collector.t

(* ## globals ## *)

(** Global collector. *)
let collector : collector A.t = A.make Collector.C_none

(* default level for spans without a level *)
let default_level_ = A.make Level.Trace
let current_level_ = A.make Level.Trace

(* ## implementation ## *)

let data_empty_build_ () = []
let[@inline] enabled () = Collector.is_some (A.get collector)
let[@inline] get_default_level () = A.get default_level_
let[@inline] set_default_level l = A.set default_level_ l
let[@inline] set_current_level l = A.set current_level_ l
let[@inline] get_current_level () = A.get current_level_

let[@inline] check_level ?(level = A.get default_level_) () : bool =
  Level.leq level (A.get current_level_)

let parent_of_span_opt_opt = function
  | None -> P_unknown
  | Some None -> P_none
  | Some (Some p) -> P_some p

let enter_span_st st (cbs : _ Collector.Callbacks.t) ?__FUNCTION__ ~__FILE__
    ~__LINE__ ?parent ?(params = []) ?(data = data_empty_build_) name : span =
  let parent = parent_of_span_opt_opt parent in
  let data = data () in
  cbs.enter_span st ~__FUNCTION__ ~__FILE__ ~__LINE__ ~parent ~params ~data name

let with_span_collector_ st (cbs : _ Collector.Callbacks.t) ?__FUNCTION__
    ~__FILE__ ~__LINE__ ?parent ?params ?data name f =
  let sp : span =
    enter_span_st st cbs ?__FUNCTION__ ~__FILE__ ~__LINE__ ?parent ?params ?data
      name
  in
  match f sp with
  | res ->
    cbs.exit_span st sp;
    res
  | exception exn ->
    let bt = Printexc.get_raw_backtrace () in
    cbs.exit_span st sp;
    Printexc.raise_with_backtrace exn bt

let[@inline] with_span ?level ?__FUNCTION__ ~__FILE__ ~__LINE__ ?parent ?params
    ?data name f =
  match A.get collector with
  | C_some (st, cbs) when check_level ?level () ->
    with_span_collector_ st cbs ?__FUNCTION__ ~__FILE__ ~__LINE__ ?parent
      ?params ?data name f
  | _ ->
    (* fast path: no collector, no span *)
    f Collector.dummy_span

let[@inline] enter_span ?level ?__FUNCTION__ ~__FILE__ ~__LINE__ ?flavor ?parent
    ?(params = []) ?data name : span =
  match A.get collector with
  | C_some (st, cbs) when check_level ?level () ->
    let params =
      match flavor with
      | None -> params
      | Some f -> Core_ext.Extension_span_flavor f :: params
    in
    (enter_span_st [@inlined never]) st cbs ?__FUNCTION__ ~__FILE__ ~__LINE__
      ?parent ~params ?data name
  | _ -> Collector.dummy_span

let[@inline] exit_span sp : unit =
  match A.get collector with
  | C_none -> ()
  | C_some (st, cbs) -> cbs.exit_span st sp

let[@inline] add_data_to_span sp data : unit =
  if sp != Collector.dummy_span && data <> [] then (
    match A.get collector with
    | C_none -> ()
    | C_some (st, cbs) -> cbs.add_data_to_span st sp data
  )

let message_collector_ st (cbs : _ Collector.Callbacks.t) ?span ?(params = [])
    ?(data = data_empty_build_) msg : unit =
  let data = data () in
  cbs.message st ~span ~params ~data msg

let[@inline] message ?level ?span ?params ?data msg : unit =
  match A.get collector with
  | C_some (st, cbs) when check_level ?level () ->
    (message_collector_ [@inlined never]) st cbs ?span ?params ?data msg
  | _ -> ()

let messagef ?level ?span ?params ?data k =
  match A.get collector with
  | C_some (st, cbs) when check_level ?level () ->
    k (fun fmt ->
        Format.kasprintf
          (fun str -> message_collector_ st cbs ?span ?params ?data str)
          fmt)
  | _ -> ()

let counter_int ?level ?(params = []) ?(data = data_empty_build_) name n : unit
    =
  match A.get collector with
  | C_some (st, cbs) when check_level ?level () ->
    let data = data () in
    cbs.counter_int st ~params ~data name n
  | _ -> ()

let counter_float ?level ?(params = []) ?(data = data_empty_build_) name f :
    unit =
  match A.get collector with
  | C_some (st, cbs) when check_level ?level () ->
    let data = data () in
    cbs.counter_float st ~params ~data name f
  | _ -> ()

let setup_collector c : unit =
  while
    let cur = A.get collector in
    match cur with
    | C_some _ -> invalid_arg "trace: collector already present"
    | C_none -> not (A.compare_and_set collector cur c)
  do
    ()
  done;

  (* initialize collector *)
  match c with
  | C_none -> ()
  | C_some (st, cb) -> cb.init st

let shutdown () =
  match A.exchange collector C_none with
  | C_none -> ()
  | C_some (st, cbs) -> cbs.shutdown st

let with_setup_collector c f =
  setup_collector c;
  Fun.protect ~finally:shutdown f

type extension_event = Types.extension_event = ..

let[@inline] extension_event ev : unit =
  match A.get collector with
  | C_none -> ()
  | C_some (st, cbs) -> cbs.extension st ev

let set_thread_name name : unit =
  extension_event @@ Core_ext.Extension_set_thread_name name

let set_process_name name : unit =
  extension_event @@ Core_ext.Extension_set_process_name name

module Internal_ = struct
  module Atomic_ = Atomic_
end

(* ### deprecated *)

[@@@ocaml.alert "-deprecated"]

let enter_manual_span ~parent ?flavor ?level ?__FUNCTION__ ~__FILE__ ~__LINE__
    ?data name : explicit_span =
  let params =
    match flavor with
    | None -> []
    | Some f -> [ Core_ext.Extension_span_flavor f ]
  in
  enter_span ~parent ~params ?level ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name

let enter_manual_toplevel_span ?flavor ?level ?__FUNCTION__ ~__FILE__ ~__LINE__
    ?data name : explicit_span =
  enter_manual_span ~parent:None ?flavor ?level ?__FUNCTION__ ~__FILE__
    ~__LINE__ ?data name

let enter_manual_sub_span ~parent ?flavor ?level ?__FUNCTION__ ~__FILE__
    ~__LINE__ ?data name : explicit_span =
  enter_manual_span ~parent:(Some parent) ?flavor ?level ?__FUNCTION__ ~__FILE__
    ~__LINE__ ?data name

let exit_manual_span = exit_span
let add_data_to_manual_span = add_data_to_span

[@@@ocaml.alert "+deprecated"]