type level =
  | TRACE (* 0 *)
  | DEBUG (* 1 *)
  | INFO  (* 2 *)
  | WARN  (* 3 *)
  | ERROR (* 4 *)

external int_of_level : level -> int = "%identity"

type writer =
  ts:float ->
  ns:string ->
  lvl:level ->
  string ->
  unit

let ( +++ ) w1 w2 ~ts ~ns ~lvl msg =
  w1 ~ts ~ns ~lvl msg;
  w2 ~ts ~ns ~lvl msg

type logger = {
  namespace : string;
  mutable child : logger;
  sibling : logger;
  mutable min_level : int;
  mutable writers : writer option;
}

let rec root_logger = {
  namespace = "";
  min_level = 0;
  child = root_logger;
  sibling = root_logger;
  writers = None;
}

let make_logger parent_logger ns =
  let namespace =
    if parent_logger == root_logger then ns
    else String.concat "." [parent_logger.namespace; ns];
  in
  let rec logger = {
    namespace;
    min_level = Int.max_int;
    child = logger;
    sibling = parent_logger.child;
    writers = None;
  } in
  parent_logger.child <- logger;
  logger

type 'a log_function =
  ((('a, Format.formatter, unit) format -> 'a) -> unit) -> unit

let write logger lvl msg =
  let ts = Unix.gettimeofday () in
  let ns = logger.namespace in
  Option.iter (fun w -> w ~ts ~ns ~lvl msg) logger.writers

let[@inline] logf logger lvl : _ log_function = fun logk ->
  if int_of_level lvl >= logger.min_level then
    logk (fun fmt ->
           Format.kasprintf (write logger lvl)
             ("@[<hov>" ^^ fmt))


module type Logs = sig
  val logger : logger
  val trace : _ log_function
  val debug : _ log_function
  val info  : _ log_function
  val warn  : _ log_function
  val error : _ log_function
  val log   : level -> _ log_function
end

let make_logs parent_logger ns =
  (module struct
    let logger = make_logger parent_logger ns
    let[@inline] trace k = logf logger TRACE k
    let[@inline] debug k = logf logger DEBUG k
    let[@inline] info  k = logf logger INFO  k
    let[@inline] warn  k = logf logger WARN  k
    let[@inline] error k = logf logger ERROR k
    let[@inline] log level k = logf logger level k
  end : Logs)

let logs ns = make_logs root_logger ns
let sublogs logger ns = make_logs logger ns

(* TODO: filters: namespace[*] => min_level *)

let rec add_writer_rec min_level writer parent_logger logger =
  if parent_logger != logger then
    begin
      logger.min_level <- min min_level logger.min_level;
      logger.writers <-
        (match logger.writers with
         | None -> Some writer
         | Some writer' -> Some (writer' +++ writer));
      add_writer_rec min_level writer logger logger.child;
      add_writer_rec min_level writer parent_logger logger.sibling
    end

let[@inline] add_writer ?(min_level = WARN) writer =
  add_writer_rec (int_of_level min_level) writer root_logger root_logger.child

module Pretty = struct
  type t = {
    mutex : Mutex.t;
    bp : Buffer.t;
    mutable align_to : int;
    out : out_channel;
  }

  (* TODO: config timestamp *)
  (* TODO: config color *)
  (* TODO: config namespace *)

  let header = function
    | TRACE -> "TRACE"
    | DEBUG -> "DEBUG"
    | INFO  -> "INFO"
    | WARN  -> "WARN"
    | ERROR -> "ERROR"

  let ansi_color = function
    | TRACE -> 34
    | DEBUG -> 36
    | INFO  -> 32
    | WARN  -> 33
    | ERROR -> 31

  let pr_timestamp bp ts =
    let ts_ms = int_of_float (ts *. 1000.0) mod 1000 in
    let tm = Unix.localtime ts in
    Printf.bprintf bp "%04d-%02d-%02d %02d:%02d:%02d.%03d"
      (tm.tm_year + 1900) (tm.tm_mon + 1) tm.tm_mday
      tm.tm_hour tm.tm_min tm.tm_sec ts_ms

  let pr_spaces bp n =
    for _ = 1 to n do
      Buffer.add_char bp ' '
    done

  let rec pr_lines_rec indent bp s i =
    match String.index_from s i '\n' with
    | exception Not_found ->
      Buffer.add_substring bp s i (String.length s - i)
    | k ->
      Buffer.add_substring bp s i (k + 1 - i);
      pr_spaces bp indent;
      pr_lines_rec indent bp s (k + 1)

  let pr_lines ~indent bp s =
    pr_lines_rec indent bp s 0

  let make (out : out_channel) = {
    mutex = Mutex.create ();
    bp = Buffer.create 512;
    align_to = 0;
    out;
  }

  let writer t ~ts ~ns ~lvl msg =
    begin
      Mutex.lock t.mutex;
      Printf.bprintf t.bp "\x1b[1m%a\x1b[22m %s "
        pr_timestamp ts ns;
      pr_spaces t.bp (t.align_to - Buffer.length t.bp);
      t.align_to <- Buffer.length t.bp;
      Printf.bprintf t.bp "\x1b[%dm%-5s\x1b[39;1m "
        (ansi_color lvl) (header lvl);
      pr_lines t.bp msg ~indent:(t.align_to - 3);
      Printf.bprintf t.bp "\x1b[0m\n";
      Buffer.output_buffer t.out t.bp;
      flush t.out;
      Buffer.clear t.bp;
      Mutex.unlock t.mutex
    end
end

let init_pretty_writer ?min_level out =
  Pretty.make out |>
  Pretty.writer |>
  add_writer ?min_level