type error = { path : string list; message : string }

exception Decode_error of error

let error_to_string e =
  let path =
    match e.path with [] -> "" | p -> String.concat "." (List.rev p) ^ ": "
  in
  path ^ e.message

let decode_error path msg = raise (Decode_error { path; message = msg })

type 'a decoder = string list -> Simdjsont_raw.element -> 'a
type encoder = Buffer.t -> unit

type 'a t = {
  decode : 'a decoder;
  encode : 'a -> encoder;
  to_json : 'a -> Simdjsont_json.t;
}

let decode_element codec elt =
  try Ok (codec.decode [] elt)
  with Decode_error e -> Error (error_to_string e)

let decode_string codec s =
  let parser = Simdjsont_raw.create_parser () in
  match Simdjsont_raw.parse_string parser s with
  | Ok elt ->
      let result = decode_element codec elt in
      let _ = Sys.opaque_identity parser in
      result
  | Error e -> Error e.Simdjsont_raw.message

let decode_string_exn codec s =
  match decode_string codec s with Ok v -> v | Error msg -> failwith msg

let encode_to_buffer codec v =
  let buf = Buffer.create 256 in
  codec.encode v buf;
  buf

let encode_string codec v = Buffer.contents (encode_to_buffer codec v)
let to_json codec v = codec.to_json v

let write_escaped_string buf s =
  Buffer.add_char buf '"';
  for i = 0 to String.length s - 1 do
    match String.unsafe_get s i with
    | '"' -> Buffer.add_string buf "\\\""
    | '\\' -> Buffer.add_string buf "\\\\"
    | '\b' -> Buffer.add_string buf "\\b"
    | '\012' -> Buffer.add_string buf "\\f"
    | '\n' -> Buffer.add_string buf "\\n"
    | '\r' -> Buffer.add_string buf "\\r"
    | '\t' -> Buffer.add_string buf "\\t"
    | c when Char.code c < 0x20 ->
        Buffer.add_string buf (Printf.sprintf "\\u%04x" (Char.code c))
    | c -> Buffer.add_char buf c
  done;
  Buffer.add_char buf '"'

let null : unit t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.element_type elt with
        | Simdjsont_raw.Null -> ()
        | _ -> decode_error path "expected null");
    encode = (fun () buf -> Buffer.add_string buf "null");
    to_json = (fun () -> Simdjsont_json.Null);
  }

let bool : bool t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.get_bool elt with
        | Ok b -> b
        | Error _ -> decode_error path "expected bool");
    encode =
      (fun b buf -> Buffer.add_string buf (if b then "true" else "false"));
    to_json = (fun b -> Simdjsont_json.Bool b);
  }

let int : int t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.get_int64 elt with
        | Ok i -> Int64.to_int i
        | Error _ -> decode_error path "expected int");
    encode = (fun i buf -> Buffer.add_string buf (string_of_int i));
    to_json = (fun i -> Simdjsont_json.Int (Int64.of_int i));
  }

let int64 : int64 t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.get_int64 elt with
        | Ok i -> i
        | Error _ -> decode_error path "expected int64");
    encode = (fun i buf -> Buffer.add_string buf (Int64.to_string i));
    to_json = (fun i -> Simdjsont_json.Int i);
  }

let float : float t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.element_type elt with
        | Simdjsont_raw.Double -> (
            match Simdjsont_raw.get_double elt with
            | Ok f -> f
            | Error _ -> Float.nan)
        | Simdjsont_raw.Int64 -> (
            match Simdjsont_raw.get_int64 elt with
            | Ok i -> Int64.to_float i
            | Error _ -> Float.nan)
        | Simdjsont_raw.Uint64 -> (
            match Simdjsont_raw.get_uint64 elt with
            | Ok u -> Int64.to_float (Unsigned.UInt64.to_int64 u)
            | Error _ -> Float.nan)
        | Simdjsont_raw.Null -> Float.nan
        | _ -> decode_error path "expected number");
    encode =
      (fun f buf ->
        if Float.is_finite f then
          Buffer.add_string buf (Printf.sprintf "%.17g" f)
        else Buffer.add_string buf "null");
    to_json = (fun f -> Simdjsont_json.Float f);
  }

let string : string t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.get_string elt with
        | Ok s -> s
        | Error _ -> decode_error path "expected string");
    encode = (fun s buf -> write_escaped_string buf s);
    to_json = (fun s -> Simdjsont_json.String s);
  }

let list (item : 'a t) : 'a list t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.get_array elt with
        | Ok arr ->
            let result = ref [] in
            let idx = ref 0 in
            Simdjsont_raw.array_iter
              (fun e ->
                let item_path = string_of_int !idx :: path in
                result := item.decode item_path e :: !result;
                incr idx)
              arr;
            List.rev !result
        | Error _ -> decode_error path "expected array");
    encode =
      (fun items buf ->
        Buffer.add_char buf '[';
        (match items with
        | [] -> ()
        | x :: xs ->
            item.encode x buf;
            List.iter
              (fun v ->
                Buffer.add_char buf ',';
                item.encode v buf)
              xs);
        Buffer.add_char buf ']');
    to_json = (fun items -> Simdjsont_json.Array (List.map item.to_json items));
  }

let array (item : 'a t) : 'a array t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.get_array elt with
        | Ok arr ->
            let len = Simdjsont_raw.array_length arr in
            if len = 0 then [||]
            else begin
              let result = ref [||] in
              let idx = ref 0 in
              Simdjsont_raw.array_iter
                (fun e ->
                  let item_path = string_of_int !idx :: path in
                  let v = item.decode item_path e in
                  if !idx = 0 then result := Array.make len v
                  else Array.unsafe_set !result !idx v;
                  incr idx)
                arr;
              !result
            end
        | Error _ -> decode_error path "expected array");
    encode =
      (fun items buf ->
        Buffer.add_char buf '[';
        let len = Array.length items in
        if len > 0 then begin
          item.encode (Array.unsafe_get items 0) buf;
          for i = 1 to len - 1 do
            Buffer.add_char buf ',';
            item.encode (Array.unsafe_get items i) buf
          done
        end;
        Buffer.add_char buf ']');
    to_json =
      (fun items ->
        Simdjsont_json.Array (List.map item.to_json (Array.to_list items)));
  }

let optional (inner : 'a t) : 'a option t =
  {
    decode =
      (fun path elt ->
        match Simdjsont_raw.element_type elt with
        | Simdjsont_raw.Null -> None
        | _ -> Some (inner.decode path elt));
    encode =
      (fun opt buf ->
        match opt with
        | None -> Buffer.add_string buf "null"
        | Some v -> inner.encode v buf);
    to_json =
      (function None -> Simdjsont_json.Null | Some v -> inner.to_json v);
  }

let map (f : 'a -> 'b) (g : 'b -> 'a) (codec : 'a t) : 'b t =
  {
    decode = (fun path elt -> f (codec.decode path elt));
    encode = (fun v buf -> codec.encode (g v) buf);
    to_json = (fun v -> codec.to_json (g v));
  }

let decode_obj_field path (obj : Simdjsont_raw.object_) name (dec : 'a decoder)
    : 'a =
  let field_path = name :: path in
  match Simdjsont_raw.object_find obj name with
  | Ok elt -> dec field_path elt
  | Error _ -> decode_error path ("missing field: " ^ name)

let decode_obj_opt_field path (obj : Simdjsont_raw.object_) name
    (dec : 'a decoder) : 'a option =
  let field_path = name :: path in
  match Simdjsont_raw.object_find obj name with
  | Ok elt -> (
      match Simdjsont_raw.element_type elt with
      | Simdjsont_raw.Null -> None
      | _ -> Some (dec field_path elt))
  | Error _ -> None

module Obj = struct
  type ('o, 'dec) builder = {
    dec : string list -> Simdjsont_raw.object_ -> 'dec;
    enc : 'o -> Buffer.t -> bool -> bool;
    ast : 'o -> (string * Simdjsont_json.t) list;
  }

  let field constructor =
    {
      dec = (fun _path _obj -> constructor);
      enc = (fun _v _buf first -> first);
      ast = (fun _v -> []);
    }

  let mem name (codec : 'a t) ~enc:(get : 'o -> 'a)
      (builder : ('o, 'a -> 'b) builder) : ('o, 'b) builder =
    {
      dec =
        (fun path obj ->
          let value = decode_obj_field path obj name codec.decode in
          builder.dec path obj value);
      enc =
        (fun v buf first ->
          let first = builder.enc v buf first in
          if not first then Buffer.add_char buf ',';
          write_escaped_string buf name;
          Buffer.add_char buf ':';
          codec.encode (get v) buf;
          false);
      ast = (fun v -> (name, codec.to_json (get v)) :: builder.ast v);
    }

  let opt_mem name (codec : 'a t) ~enc:(get : 'o -> 'a option)
      (builder : ('o, 'a option -> 'b) builder) : ('o, 'b) builder =
    {
      dec =
        (fun path obj ->
          let value = decode_obj_opt_field path obj name codec.decode in
          builder.dec path obj value);
      enc =
        (fun v buf first ->
          let first = builder.enc v buf first in
          match get v with
          | None -> first
          | Some inner_v ->
              if not first then Buffer.add_char buf ',';
              write_escaped_string buf name;
              Buffer.add_char buf ':';
              codec.encode inner_v buf;
              false);
      ast =
        (fun v ->
          match get v with
          | None -> builder.ast v
          | Some inner_v -> (name, codec.to_json inner_v) :: builder.ast v);
    }

  let finish (builder : ('o, 'o) builder) : 'o t =
    {
      decode =
        (fun path elt ->
          match Simdjsont_raw.get_object elt with
          | Ok obj -> builder.dec path obj
          | Error _ -> decode_error path "expected object");
      encode =
        (fun v buf ->
          Buffer.add_char buf '{';
          let _ = builder.enc v buf true in
          Buffer.add_char buf '}');
      to_json = (fun v -> Simdjsont_json.Object (builder.ast v));
    }
end

let value : Simdjsont_json.t t =
  let rec decode_value path elt =
    match Simdjsont_raw.element_type elt with
    | Simdjsont_raw.Null -> Simdjsont_json.Null
    | Simdjsont_raw.Bool ->
        Simdjsont_json.Bool
          (match Simdjsont_raw.get_bool elt with Ok b -> b | _ -> false)
    | Simdjsont_raw.Int64 ->
        Simdjsont_json.Int
          (match Simdjsont_raw.get_int64 elt with Ok i -> i | _ -> 0L)
    | Simdjsont_raw.Uint64 ->
        let i =
          match Simdjsont_raw.get_uint64 elt with
          | Ok u -> Unsigned.UInt64.to_int64 u
          | _ -> 0L
        in
        Simdjsont_json.Int i
    | Simdjsont_raw.Double ->
        Simdjsont_json.Float
          (match Simdjsont_raw.get_double elt with Ok f -> f | _ -> 0.0)
    | Simdjsont_raw.String ->
        Simdjsont_json.String
          (match Simdjsont_raw.get_string elt with Ok s -> s | _ -> "")
    | Simdjsont_raw.Array -> (
        match Simdjsont_raw.get_array elt with
        | Ok arr ->
            let items = ref [] in
            let idx = ref 0 in
            Simdjsont_raw.array_iter
              (fun e ->
                items := decode_value (string_of_int !idx :: path) e :: !items;
                incr idx)
              arr;
            Simdjsont_json.Array (List.rev !items)
        | Error _ -> Simdjsont_json.Array [])
    | Simdjsont_raw.Object -> (
        match Simdjsont_raw.get_object elt with
        | Ok obj ->
            let members = ref [] in
            Simdjsont_raw.object_iter
              (fun k v ->
                members := (k, decode_value (k :: path) v) :: !members)
              obj;
            Simdjsont_json.Object (List.rev !members)
        | Error _ -> Simdjsont_json.Object [])
  in
  {
    decode = decode_value;
    encode = (fun v buf -> Simdjsont_json.write buf v);
    to_json = (fun v -> v);
  }