open StdLabels

let rec pp :
    ?topLevel:bool -> ?path:string list -> Format.formatter -> Otoml.t -> unit =
 fun ?(topLevel = true) ?(path = []) format -> function
  | Otoml.TomlString v -> begin
      match String.contains v '\n' with
      | false ->
          Format.pp_print_string format "\"";
          Format.pp_print_string format v;
          Format.pp_print_string format "\""
      | true ->
          Format.pp_print_string format {|"""|};
          Format.pp_print_text format v;
          Format.pp_print_string format {|"""|}
    end
  | Otoml.TomlInteger i -> Format.pp_print_int format i
  | Otoml.TomlFloat f -> Format.pp_print_float format f
  | Otoml.TomlBoolean b -> Format.pp_print_bool format b
  | Otoml.TomlArray l -> begin
      match (topLevel, l) with
      | _, [] -> Format.pp_print_string format "[]"
      | false, _ -> Format.pp_print_string format "..."
      | true, l ->
          Format.pp_print_string format "[";
          Format.pp_print_break format 0 4;
          Format.pp_open_vbox format 0;
          Format.pp_print_list
            ~pp_sep:(fun f () ->
              Format.pp_print_string f ",";
              Format.pp_print_cut f ())
            pp format l;
          Format.pp_close_box format ();
          Format.pp_print_cut format ();
          Format.pp_print_string format "]"
    end
  | Otoml.TomlTable elements | Otoml.TomlInlineTable elements ->
      pp_table ~path format elements
  | Otoml.TomlTableArray t -> Format.pp_print_list pp format t
  | Otoml.TomlOffsetDateTime _
  | Otoml.TomlLocalDate _
  | Otoml.TomlLocalDateTime _
  | Otoml.TomlLocalTime _ -> ()

and pp_key_values : Format.formatter -> string * Otoml.t -> unit =
 fun format (name, value) ->
  Format.fprintf format "%s = %a" name (pp ~topLevel:false ~path:[]) value

and pp_table :
    ?path:string list -> Format.formatter -> (string * Otoml.t) list -> unit =
 fun ?(path = []) format elements ->
  (* Create two lists, one for the subtables, and one for the values. 

    As a table is valid until the next table, we then start to print the
    preoperties before printintg the subtables inside the element.
   *)
  let subtables, properties =
    List.partition elements ~f:(fun (_, v) ->
        match v with
        | Otoml.TomlTable _ -> true
        | _ -> false)
  in

  let () =
    match properties with
    | [] -> ()
    | _ ->
        let isTopLevel =
          match path with
          | [] -> true
          | _ -> false
        in
        if not isTopLevel then begin
          let path = List.rev path in
          Format.pp_print_cut format ();
          Format.fprintf format "[%a]"
            (Format.pp_print_list
               ~pp_sep:(fun f () -> Format.pp_print_string f ".")
               Format.pp_print_string)
            path;
          Format.pp_print_break format 0 4;
          Format.pp_open_vbox format 0
        end;
        if isTopLevel then begin
          Format.pp_print_list ~pp_sep:Format.pp_print_cut
            (fun format v ->
              match v with
              | key, Otoml.TomlTable elements ->
                  pp_table ~path:(key :: path) format elements
              | other -> pp_key_values format other)
            format properties
        end
        else begin
          Format.pp_print_string format "..."
        end;

        if not isTopLevel then begin
          Format.pp_close_box format ()
        end;
        Format.pp_print_cut format ()
  in
  (* Then go deeper inside each subtable *)
  List.iter subtables ~f:(function
    | name, Otoml.TomlTable v -> pp_table ~path:(name :: path) format v
    | _ -> (* Because of the partition, this should not happen *) ())

module Decode = struct
  module S = struct
    type value = Otoml.t

    let pp : Format.formatter -> value -> unit =
     fun formatter v ->
      Format.pp_open_vbox formatter 0;
      pp formatter v;
      Format.pp_close_box formatter ()

    let of_string : string -> (value, string) result =
      Otoml.Parser.from_string_result

    let of_file : string -> (value, string) result =
      Otoml.Parser.from_file_result

    let get_string : value -> string option = Otoml.get_opt Otoml.get_string
    let get_int : value -> int option = Otoml.get_opt Otoml.get_integer
    let get_float : value -> float option = Otoml.get_opt Otoml.get_float
    let get_bool : value -> bool option = Otoml.get_opt Otoml.get_boolean
    let get_null : value -> unit option = fun _ -> None

    let get_list : value -> value list option =
      Otoml.get_opt @@ Otoml.get_array Fun.id

    let get_key_value_pairs : value -> (value * value) list option =
      Otoml.get_opt (fun key ->
          Otoml.get_table key |> List.map ~f:(fun (k, v) -> (Otoml.string k, v)))

    let to_list : value list -> value = Otoml.array
  end

  include Decoders.Decode.Make (S)
end