module List = ListLabels
module String = StringLabels

module Json = struct
  type t =
    [ `Int of int
    | `Float of float
    | `String of string
    | `List of t list
    | `Bool of bool
    | `Assoc of (string * t) list
    | `Null
    ]
end

module Timestamp : sig
  type t

  val to_json : t -> Json.t
  val of_float_seconds : float -> t
  val to_float_seconds : t -> float
end = struct
  type t = float

  let of_float_seconds x = x
  let to_float_seconds x = x

  let to_json f =
    let n = int_of_float @@ (f *. 1_000_000.) in
    `Int n
  ;;
end

module Id = struct
  type t =
    [ `Int of int
    | `String of string
    ]

  let create x = x

  let to_string = function
    | `String s -> s
    | `Int i -> string_of_int i
  ;;

  let to_json (t : t) = (t :> Json.t)
  let field id = "id", to_json id
end

module Stack_frame = struct
  module Raw = struct
    type t = string list

    let create t = t
    let to_json t = `List (List.map t ~f:(fun s -> `String s))
  end

  type t =
    { parent : Id.t option
    ; name : string
    ; category : string
    }

  let create ?parent ~name ~category () = { parent; name; category }

  let to_json { parent; name; category } : Json.t =
    let json = [ "name", `String name; "category", `String category ] in
    let json =
      match parent with
      | None -> json
      | Some id -> ("parent", Id.to_json id) :: json
    in
    `Assoc json
  ;;
end

module Event = struct
  [@@@ocaml.warning "-37"]

  module Timestamp = Timestamp

  type common_fields =
    { name : string
    ; cat : string list
    ; ts : Timestamp.t
    ; tts : Timestamp.t option
    ; pid : int
    ; tid : int
    ; cname : string option
    ; stackframe : [ `Id of Id.t | `Raw of Stack_frame.Raw.t ] option
    }

  let common_fields ?tts ?cname ?(cat = []) ?(pid = 0) ?(tid = 0) ?stackframe ~ts ~name ()
    =
    { tts; cname; cat; ts; pid; tid; name; stackframe }
  ;;

  let set_ts t ts = { t with ts }
  let ts t = t.ts

  type scope =
    | Global
    | Process
    | Thread

  type async =
    | Start
    | Instant
    | End

  type args = (string * Json.t) list

  type object_kind =
    | New
    | Snapshot of
        { cat : string list option
        ; args : args
        }
    | Destroy

  type metadata =
    | Process_name of
        { pid : int
        ; name : string
        }
    | Process_labels of
        { pid : int
        ; labels : string
        }
    | Thread_name of
        { tid : int
        ; pid : int
        ; name : string
        }
    | Process_sort_index of
        { pid : int
        ; sort_index : int
        }
    | Thread_sort_index of
        { pid : int
        ; tid : int
        ; sort_index : int
        }

  (* TODO support flow, samples, references, memory dumps *)
  type t =
    | Counter of common_fields * args * Id.t option
    | Duration_start of common_fields * args * Id.t option
    | Duration_end of
        { pid : int
        ; tid : int
        ; ts : float
        ; args : args option
        }
    | Complete of
        { common : common_fields
        ; args : args option
        ; dur : Timestamp.t
        ; tdur : Timestamp.t option
        }
    | Instant of common_fields * scope option * args option
    | Async of
        { common : common_fields
        ; async : async
        ; scope : string option
        ; id : Id.t
        ; args : args option
        }
    | Object of
        { common : common_fields
        ; object_kind : object_kind
        ; id : Id.t
        ; scope : string option
        }
    | Metadata of metadata

  let phase s = "ph", `String s

  let add_field_opt to_field field fields =
    match field with
    | None -> fields
    | Some f -> to_field f :: fields
  ;;

  let json_fields_of_common_fields { name; cat; ts; tts; pid; tid; cname; stackframe } =
    let fields =
      [ "name", `String name
      ; "cat", `String (String.concat ~sep:"," cat)
      ; "ts", Timestamp.to_json ts
      ; "pid", `Int pid
      ; "tid", `Int tid
      ]
    in
    let fields = add_field_opt (fun cname -> "cname", `String cname) cname fields in
    let fields = add_field_opt (fun tts -> "tts", Timestamp.to_json tts) tts fields in
    add_field_opt
      (fun stackframe ->
         match stackframe with
         | `Id id -> "sf", Id.to_json id
         | `Raw r -> "stack", Stack_frame.Raw.to_json r)
      stackframe
      fields
  ;;

  let json_of_scope = function
    | Global -> `String "g"
    | Process -> `String "p"
    | Thread -> `String "t"
  ;;

  let args_field fields = "args", `Assoc fields

  let json_fields_of_metadata m =
    let fields =
      let common pid name = [ "name", `String name; "pid", `Int pid ] in
      match m with
      | Process_name { pid; name } ->
        args_field [ "name", `String name ] :: common pid "thread_name"
      | Process_labels { pid; labels } ->
        args_field [ "labels", `String labels ] :: common pid "process_labels"
      | Thread_name { tid; pid; name } ->
        ("tid", `Int tid)
        :: args_field [ "name", `String name ]
        :: common pid "process_name"
      | Process_sort_index { pid; sort_index } ->
        args_field [ "sort_index", `Int sort_index ] :: common pid "process_sort_index"
      | Thread_sort_index { pid; sort_index; tid } ->
        ("tid", `Int tid)
        :: args_field [ "sort_index", `Int sort_index ]
        :: common pid "thread_sort_index"
    in
    phase "M" :: fields
  ;;

  let to_json_fields : t -> (string * Json.t) list = function
    | Counter (common, args, id) ->
      let fields = json_fields_of_common_fields common in
      let fields = phase "C" :: args_field args :: fields in
      add_field_opt Id.field id fields
    | Duration_start (common, args, id) ->
      let fields = json_fields_of_common_fields common in
      let fields = phase "B" :: args_field args :: fields in
      add_field_opt Id.field id fields
    | Duration_end { pid; tid; ts; args } ->
      let fields = [ "tid", `Int tid; "pid", `Int pid; "ts", `Float ts; phase "E" ] in
      add_field_opt args_field args fields
    | Complete { common; dur; args; tdur } ->
      let fields = json_fields_of_common_fields common in
      let fields = phase "X" :: ("dur", Timestamp.to_json dur) :: fields in
      let fields =
        add_field_opt (fun tdur -> "tdur", Timestamp.to_json tdur) tdur fields
      in
      add_field_opt args_field args fields
    | Instant (common, scope, args) ->
      let fields = json_fields_of_common_fields common in
      let fields = phase "i" :: fields in
      let fields = add_field_opt (fun s -> "s", json_of_scope s) scope fields in
      add_field_opt args_field args fields
    | Async { common; async; scope; id; args } ->
      let fields = json_fields_of_common_fields common in
      let fields = Id.field id :: fields in
      let fields =
        let ph =
          let s =
            match async with
            | Start -> "b"
            | Instant -> "n"
            | End -> "e"
          in
          phase s
        in
        ph :: fields
      in
      let fields = add_field_opt (fun s -> "scope", `String s) scope fields in
      add_field_opt args_field args fields
    | Object { common; object_kind; id; scope } ->
      let fields = json_fields_of_common_fields common in
      let fields = Id.field id :: fields in
      let fields =
        let ph, args =
          match object_kind with
          | New -> "N", None
          | Destroy -> "D", None
          | Snapshot { cat; args } ->
            let snapshot =
              add_field_opt
                (fun cat -> "cat", `String (String.concat ~sep:"," cat))
                cat
                args
            in
            "O", Some [ "snapshot", `Assoc snapshot ]
        in
        let fields = phase ph :: fields in
        add_field_opt args_field args fields
      in
      add_field_opt (fun s -> "scope", `String s) scope fields
    | Metadata m -> json_fields_of_metadata m
  ;;

  let to_json t = `Assoc (to_json_fields t)
  let counter ?id common args = Counter (common, args, id)
  let complete ?tdur ?args ~dur common = Complete { common; tdur; dur; args }
  let async ?scope ?args id async common = Async { common; args; scope; id; async }
  let instant ?args ?scope common = Instant (common, scope, args)
end

module Output_object = struct
  type t =
    { displayTimeUnit : [ `Ms | `Ns ] option
    ; traceEvents : Event.t list
    ; stackFrames : (Id.t * Stack_frame.t) list option
    ; extra_fields : (string * Json.t) list option
    }

  let to_json { displayTimeUnit; traceEvents; extra_fields; stackFrames } =
    let json = [ "traceEvents", `List (List.map traceEvents ~f:Event.to_json) ] in
    let json =
      match displayTimeUnit with
      | None -> json
      | Some u ->
        ( "displayTimeUnit"
        , `String
            (match u with
             | `Ms -> "ms"
             | `Ns -> "ns") )
        :: json
    in
    let json : (string * Json.t) list =
      match stackFrames with
      | None -> json
      | Some frames ->
        let frames =
          List.map frames ~f:(fun (id, frame) ->
            let id = Id.to_string id in
            id, Stack_frame.to_json frame)
        in
        ("stackFrames", `Assoc frames) :: json
    in
    let json =
      match extra_fields with
      | None -> json
      | Some extra_fields -> json @ extra_fields
    in
    `Assoc json
  ;;

  let create ?displayTimeUnit ?extra_fields ?stackFrames ~traceEvents () =
    { displayTimeUnit; extra_fields; traceEvents; stackFrames }
  ;;
end
