open Angstrom

type t =
  | Nil
  | Bool of bool
  | Int_fix_pos of int
  | Int_fix_neg of int
  | Int8 of int
  | Int16 of int
  | Int32 of int32
  | Int64 of int64
  | Uint8 of int
  | Uint16 of int
  | Uint32 of int32
  | Uint64 of int64
  | Float of float
  | Double of float
  | Str_fix of string
  | Str8 of string
  | Str16 of string
  | Str32 of string
  | Bin8 of string
  | Bin16 of string
  | Bin32 of string
  | Array_fix of t list
  | Array16 of t list
  | Array32 of t list
  | Map_fix of (t * t) list
  | Map16 of (t * t) list
  | Map32 of (t * t) list
  | Ext_fix_1 of int * string
  | Ext_fix_2 of int * string
  | Ext_fix_4 of int * string
  | Ext_fix_8 of int * string
  | Ext_fix_16 of int * string
  | Ext8 of int * string
  | Ext16 of int * string
  | Ext32 of int * string

module Be = struct
  include Be

  module Es = EndianString.BigEndian

  let get_float s = Es.get_float s 0
  let get_double s = Es.get_double s 0

  let get_int8 s = Es.get_int8 s 0
  let get_int16 s = Es.get_int16 s 0
  let get_int32 s = Es.get_int32 s 0
  let get_int64 s = Es.get_int64 s 0
  let get_uint32_as_int s =
    (* Trickery is afoot:
       The Int32 module doesn't provide a simple way to convert a Int32.t
       value to an int in an unsigned-friendly way.  So we're going to fake
       it here by padding the value we read into a (theoretically signed)
       Int64.t. *)
    let v = get_int64 ("\x00\x00\x00\x00" ^ s) in
    if Int64.compare v (Int64.of_int max_int) > 0 then (
      invalid_arg "msgpack field with length longer than max_int"
    );
    Int64.to_int v

  let get_uint8 s = Es.get_uint8 s 0
  let get_uint16 s = Es.get_uint16 s 0
  let get_uint32 s = Es.get_int32 s 0
  let get_uint64 s = Es.get_int64 s 0

  let make_es size set x =
    let buf = Bytes.create size in
    set buf 0 x;
    Bytes.to_string buf

  let of_float x =
    make_es 4 Es.set_float x
  let of_double x =
    make_es 8 Es.set_double x

  let of_int8 x =
    make_es 1 Es.set_int8 x
  let of_int16 x =
    make_es 2 Es.set_int16 x
  let of_int32 x =
    make_es 4 Es.set_int32 x
  let of_int64 x =
    make_es 8 Es.set_int64 x

  let of_uint8 x =
    make_es 1 Es.set_int8 x
  let of_uint16 x =
    make_es 2 Es.set_int16 x
  let of_uint32 x =
    make_es 4 Es.set_int32 x
  let of_uint64 x =
    make_es 8 Es.set_int64 x
  let of_uint32_as_int x =
    (* Trickery is afoot:
       The Int32 module doesn't provide a simple way to convert an int value
       to an Int32.t in an unsigned-friendly way.  So we're going to fake it
       here by padding the value we write into a (theoretically signed)
       Int64.t. *)
    let s = make_es 8 Es.set_int64 (Int64.of_int x) in
    String.sub s 4 4
end

let byte x = Angstrom.char @@ Char.chr x

(* nil *)
let nil_ =
  byte 0xc0 *> return Nil

let nil = nil_

(* bool *)
let false_ =
  byte 0xc2 *> return (Bool false)
let true_ =
  byte 0xc3 *> return (Bool true)

let bool =
  choice [false_; true_]

(* int *)
let int_fix_pos =
  satisfy (fun c -> Char.code c land 0x80 = 0) >>| fun c ->
  Int_fix_pos (Be.get_uint8 (String.make 1 c))
let int_fix_neg =
  satisfy (fun c -> Char.code c land 0xe0 = 0xe0) >>| fun c ->
  Int_fix_neg (Be.get_int8 (String.make 1 c))
let uint8 =
  byte 0xcc *> Be.uint8 >>| fun i -> Uint8 i
let uint16 =
  byte 0xcd *> Be.uint16 >>| fun i -> Uint16 i
let uint32 =
  byte 0xce *> Be.uint32 >>| fun i -> Uint32 i
let uint64 =
  byte 0xcf *> Be.uint64 >>| fun i -> Uint64 i
let int8 =
  byte 0xd0 *> Be.int8 >>| fun i -> Int8 i
let int16 =
  byte 0xd1 *> Be.int16 >>| fun i -> Int16 i
let int32 =
  byte 0xd2 *> Be.int32 >>| fun i -> Int32 i
let int64 =
  byte 0xd3 *> Be.int64 >>| fun i -> Int64 i

let int =
  choice [
    int_fix_pos;
    int_fix_neg;
    int8;
    int16;
    int32;
    int64;
    uint8;
    uint16;
    uint32;
    uint64;
  ]

(* float *)
let float =
  byte 0xca *> Be.float >>| fun f -> Float f
let double =
  byte 0xcb *> Be.double >>| fun d -> Double d

let raw8 tag =
  byte tag *> Be.uint8 >>= fun length ->
  take length
let raw16 tag =
  byte tag *> Be.uint16 >>= fun length ->
  take length
let raw32 tag =
  byte tag *> take 4 >>= fun s ->
  let length = Be.get_uint32_as_int s in
  take length

(* str *)
let str_fix =
  satisfy (fun c -> Char.code c land 0xe0 = 0xa0) >>= fun c ->
  let length = Be.get_uint8 (String.make 1 c) land 0x1f in
  take length >>| fun x -> Str_fix x
let str8 = raw8 0xd9 >>| fun x -> Str8 x
let str16 = raw16 0xda >>| fun x -> Str16 x
let str32 = raw32 0xdb >>| fun x -> Str32 x
let str =
  choice [
    str_fix;
    str8;
    str16;
    str32;
  ]

(* bin *)
let bin8 = raw8 0xc4 >>| fun x -> Bin8 x
let bin16 = raw16 0xc5 >>| fun x -> Bin16 x
let bin32 = raw32 0xc6 >>| fun x -> Bin32 x
let bin =
  choice [
    bin8;
    bin16;
    bin32;
  ]

(* array *)
let array_fix any =
  satisfy (fun c -> Char.code c land 0xf0 = 0x90) >>= fun c ->
  let length = Be.get_uint8 (String.make 1 c) land 0x0f in
  count length any >>| fun x -> Array_fix x
let array16 any =
  byte 0xdc *> Be.uint16 >>= fun length ->
  count length any >>| fun x -> Array16 x
let array32 any =
  byte 0xdd *> take 4 >>= fun s ->
  let length = Be.get_uint32_as_int s in
  count length any >>| fun x -> Array32 x
let array any =
  choice [
    array_fix any;
    array16 any;
    array32 any;
  ]

(* ext *)
let ext_fix_n tag n =
  byte tag *> Be.int8 >>= fun typ ->
  take n >>= fun content ->
  return (typ, content)

let ext_fix_1 = ext_fix_n 0xd4 1 >>| fun (typ, x) -> Ext_fix_1 (typ, x)
let ext_fix_2 = ext_fix_n 0xd5 2 >>| fun (typ, x) -> Ext_fix_2 (typ, x)
let ext_fix_4 = ext_fix_n 0xd6 4 >>| fun (typ, x) -> Ext_fix_4 (typ, x)
let ext_fix_8 = ext_fix_n 0xd7 8 >>| fun (typ, x) -> Ext_fix_8 (typ, x)
let ext_fix_16 = ext_fix_n 0xd8 16 >>| fun (typ, x) -> Ext_fix_16 (typ, x)

let ext8 =
  byte 0xc7 *> Be.uint8 >>= fun length ->
  Be.int8 >>= fun typ ->
  take length >>= fun content ->
  return @@ Ext8 (typ, content)
let ext16 =
  byte 0xc8 *> Be.uint16 >>= fun length ->
  Be.int8 >>= fun typ ->
  take length >>= fun content ->
  return @@ Ext16 (typ, content)
let ext32 =
  byte 0xc9 *> take 4 >>= fun s ->
  let length = Be.get_uint32_as_int s in
  Be.int8 >>= fun typ ->
  take length >>= fun content ->
  return @@ Ext32 (typ, content)

let ext =
  choice [
    ext_fix_1;
    ext_fix_2;
    ext_fix_4;
    ext_fix_8;
    ext_fix_16;
    ext8;
    ext16;
    ext32;
  ]

(* map *)
let tuple a b =
  a >>= fun a' ->
  b >>= fun b' ->
  return (a', b')

let map_fix ~key ~value =
  satisfy (fun c -> Char.code c land 0xf0 = 0x80) >>= fun c ->
  let length = Be.get_uint8 (String.make 1 c) land 0x0f in
  count length (tuple key value) >>| fun x -> Map_fix x
let map16 ~key ~value =
  byte 0xde *> Be.uint16 >>= fun length ->
  count length (tuple key value) >>| fun x -> Map16 x
let map32 ~key ~value =
  byte 0xdf *> take 4 >>= fun s ->
  let length = Be.get_uint32_as_int s in
  count length (tuple key value) >>| fun x -> Map32 x
let map ~key ~value =
  choice [
    map_fix ~key ~value;
    map16 ~key ~value;
    map32 ~key ~value;
  ]

(* One msgpack value *)
let msgpack =
  fix (
    fun v ->
      let map' = map ~key:v ~value:v in
      let array' = array v in
      choice [
        nil;
        bool;
        int;
        float;
        double;
        str;
        bin;
        array';
        map';
        ext;
      ]
  )

(* Many concatenated msgpack values *)
let msgpacks =
  many msgpack

let of_string s =
  parse_only msgpack (`String s)

let of_string_exn s =
  match of_string s with
  | Result.Ok msg -> msg
  | Result.Error e -> invalid_arg e

let msgs_of_string s =
  parse_only msgpacks (`String s)

let msgs_of_string_exn s =
  match msgs_of_string s with
  | Result.Ok msg -> msg
  | Result.Error e -> invalid_arg e

let to_buffer buf msg =
  let bc c = Buffer.add_char buf (Char.chr c) in
  let bs s = Buffer.add_string buf s in
  let rec add m =
    match m with
    | Nil -> bc 0xc0
    | Bool false -> bc 0xc2
    | Bool true -> bc 0xc3
    | Int_fix_pos i -> bs (Be.of_uint8 i)
    | Int_fix_neg i -> bs (Be.of_int8 i)
    | Int8 i -> bc 0xd0; bs (Be.of_int8 i)
    | Int16 i -> bc 0xd1; bs (Be.of_int16 i)
    | Int32 i -> bc 0xd2; bs (Be.of_int32 i)
    | Int64 i -> bc 0xd3; bs (Be.of_int64 i)
    | Uint8 i -> bc 0xcc; bs (Be.of_uint8 i)
    | Uint16 i -> bc 0xcd; bs (Be.of_uint16 i)
    | Uint32 i -> bc 0xce; bs (Be.of_uint32 i)
    | Uint64 i -> bc 0xcf; bs (Be.of_uint64 i)
    | Float f -> bc 0xca; bs (Be.of_float f)
    | Double d -> bc 0xcb; bs (Be.of_double d)
    | Str_fix s -> bc (0xa0 lor String.length s); bs s
    | Str8 s ->
      bc 0xd9;
      bs (Be.of_uint8 (String.length s));
      bs s
    | Str16 s ->
      bc 0xda;
      bs (Be.of_uint16 (String.length s));
      bs s
    | Str32 s ->
      bc 0xdb;
      bs (Be.of_uint32_as_int (String.length s));
      bs s
    | Bin8 b ->
      bc 0xc4;
      bs (Be.of_uint8 (String.length b));
      bs b
    | Bin16 b ->
      bc 0xc5;
      bs (Be.of_uint16 (String.length b));
      bs b
    | Bin32 b ->
      bc 0xc6;
      bs (Be.of_uint32_as_int (String.length b));
      bs b
    | Array_fix a -> bc (0x90 lor List.length a); List.iter add a
    | Array16 a ->
      bc 0xdc;
      bs (Be.of_uint16 (List.length a));
      List.iter add a
    | Array32 a ->
      bc 0xdd;
      bs (Be.of_uint32_as_int (List.length a));
      List.iter add a
    | Map_fix m ->
      bc (0x80 lor List.length m);
      List.iter (fun (k, v) -> add k; add v) m
    | Map16 m ->
      bc 0xde;
      bs (Be.of_uint16 (List.length m));
      List.iter (fun (k, v) -> add k; add v) m
    | Map32 m ->
      bc 0xdf;
      bs (Be.of_uint32_as_int (List.length m));
      List.iter (fun (k, v) -> add k; add v) m
    | Ext_fix_1 (typ, s) ->
      bc 0xd4;
      bs (Be.of_int8 typ);
      bs s
    | Ext_fix_2 (typ, s) ->
      bc 0xd5;
      bs (Be.of_int8 typ);
      bs s
    | Ext_fix_4 (typ, s) ->
      bc 0xd6;
      bs (Be.of_int8 typ);
      bs s
    | Ext_fix_8 (typ, s) ->
      bc 0xd7;
      bs (Be.of_int8 typ);
      bs s
    | Ext_fix_16 (typ, s) ->
      bc 0xd8;
      bs (Be.of_int8 typ);
      bs s
    | Ext8 (typ, s) ->
      bc 0xc7;
      bs (Be.of_uint8 (String.length s));
      bs (Be.of_int8 typ);
      bs s
    | Ext16 (typ, s) ->
      bc 0xc8;
      bs (Be.of_uint16 (String.length s));
      bs (Be.of_int8 typ);
      bs s
    | Ext32 (typ, s) ->
      bc 0xc9;
      bs (Be.of_uint32_as_int (String.length s));
      bs (Be.of_int8 typ);
      bs s
  in
  add msg

let to_string msg =
  let buf = Buffer.create 1_024 in
  to_buffer buf msg;
  Buffer.contents buf

let msgs_to_string msgs =
  let buf = Buffer.create 1_024 in
  List.iter (to_buffer buf) msgs;
  Buffer.contents buf