(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2018 Dynamic Ledger Solutions, Inc. <contact@tezos.com>     *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

(* error encoding *)
let write_error_encoding =
  let open Binary_error in
  let open Encoding in
  union
    [
      case
        (Tag 0)
        ~title:"Size limit exceeded"
        empty
        (function Size_limit_exceeded -> Some () | _ -> None)
        (fun () -> Size_limit_exceeded);
      case
        (Tag 1)
        ~title:"No case matched"
        empty
        (function No_case_matched -> Some () | _ -> None)
        (fun () -> No_case_matched);
      case
        (Tag 2)
        ~title:"Invalid int"
        (obj3 (req "min" int31) (req "v" int31) (req "max" int31))
        (function
          | Invalid_int {min; v; max} -> Some (min, v, max) | _ -> None)
        (fun (min, v, max) -> Invalid_int {min; v; max});
      case
        (Tag 3)
        ~title:"Invalid float"
        (obj3 (req "min" float) (req "v" float) (req "max" float))
        (function
          | Invalid_float {min; v; max} -> Some (min, v, max) | _ -> None)
        (fun (min, v, max) -> Invalid_float {min; v; max});
      case
        (Tag 4)
        ~title:"Invalid bytes length"
        (obj2 (req "expected" int31) (req "found" int31))
        (function
          | Invalid_bytes_length {expected; found} ->
              Some (expected, found)
          | _ ->
              None)
        (fun (expected, found) -> Invalid_bytes_length {expected; found});
      case
        (Tag 5)
        ~title:"Invalid string length"
        (obj2 (req "expected" int31) (req "found" int31))
        (function
          | Invalid_string_length {expected; found} ->
              Some (expected, found)
          | _ ->
              None)
        (fun (expected, found) -> Invalid_bytes_length {expected; found});
      case
        (Tag 6)
        ~title:"Invalid natural"
        empty
        (function Invalid_natural -> Some () | _ -> None)
        (fun () -> Invalid_natural);
      case
        (Tag 7)
        ~title:"List too long"
        empty
        (function List_too_long -> Some () | _ -> None)
        (fun () -> List_too_long);
      case
        (Tag 8)
        ~title:"Array too long"
        empty
        (function Array_too_long -> Some () | _ -> None)
        (fun () -> Array_too_long);
      case
        (Tag 9)
        ~title:"Error in custom encoder"
        string
        (function Custom_write_error s -> Some s | _ -> None)
        (fun s -> Custom_write_error s) ]

let raise (error : Binary_error.write_error) =
  raise (Binary_error.Write_error error)

(** One specific instance of the generic writer *)

type state = {
  mutable buffer : Bytes.t;  (** The buffer where to write. *)
  mutable offset : int;
      (** The offset of the next byte to be written in [buffer]. *)
  mutable allowed_bytes : int option;
      (** Maximum number of bytes that are allowed to be write in [buffer]
      (after [offset]) before to fail (None = unlimited). *)
}

let check_allowed_bytes state size =
  match state.allowed_bytes with
  | Some len when len < size ->
      raise Size_limit_exceeded
  | Some len ->
      state.allowed_bytes <- Some (len - size)
  | None ->
      ()

(** [may_resize state size] will first ensure there is enough
    space in [state.buffer] for writing [size] bytes (starting at
    [state.offset]).

    When the buffer does not have enough space for writing [size] bytes,
    but still has enough [allowed_bytes], it will replace the buffer
    with a buffer large enough.

    @raise [Binary_error.Write_error Size_limit_exceeded] when there is
           not enough allowed bytes to write [size] bytes. *)
let may_resize state size =
  check_allowed_bytes state size ;
  let buffer_len = Bytes.length state.buffer in
  if buffer_len - state.offset < size then (
    let new_buffer = Bytes.create (max (2 * buffer_len) (buffer_len + size)) in
    Bytes.blit state.buffer 0 new_buffer 0 state.offset ;
    state.buffer <- new_buffer ) ;
  state.offset <- state.offset + size


module Atom = Binary_generic_writer_backend.FullOfMinimal (struct
  type nonrec state = state

  let get_size_limit { allowed_bytes ; _ } = allowed_bytes
  let set_size_limit state limit = state.allowed_bytes <- limit

  let set_int kind buffer ofs v =
    match kind with
    | `Int31 | `Uint30 ->
        TzEndian.set_int32 buffer ofs (Int32.of_int v)
    | `Int16 | `Uint16 ->
        TzEndian.set_int16 buffer ofs v
    | `Int8 | `Uint8 ->
        TzEndian.set_int8 buffer ofs v

  let int kind state v =
    let ofs = state.offset in
    may_resize state (Binary_size.integer_to_size kind) ;
    set_int kind state.buffer ofs v

  let int8 = int `Int8

  let uint8 = int `Uint8

  let int16 = int `Int16

  let uint16 = int `Uint16

  let uint30 = int `Uint30

  let int31 = int `Int31

  let int32 state v =
    let ofs = state.offset in
    may_resize state Binary_size.int32 ;
    TzEndian.set_int32 state.buffer ofs v

  let int64 state v =
    let ofs = state.offset in
    may_resize state Binary_size.int64 ;
    TzEndian.set_int64 state.buffer ofs v

  let float state v =
    let ofs = state.offset in
    may_resize state Binary_size.float ;
    TzEndian.set_double state.buffer ofs v

  let fixed_kind_bytes length state s =
    if Bytes.length s <> length then
      raise (Invalid_bytes_length {expected = length; found = Bytes.length s}) ;
    let ofs = state.offset in
    may_resize state length ;
    Bytes.blit s 0 state.buffer ofs length

  let fixed_kind_string length state s =
    if String.length s <> length then
      raise
        (Invalid_string_length {expected = length; found = String.length s}) ;
    let ofs = state.offset in
    may_resize state length ;
    Bytes.blit_string s 0 state.buffer ofs length
end)

(** ******************** *)

(** Various entry points *)

let write_exn e v buffer offset len =
  (* By hardcoding [allowed_bytes] with the buffer length,
       we ensure that [write] will never reallocate the buffer. *)
  let state = {buffer; offset; allowed_bytes = Some len} in
  Binary_generic_writer.write_rec e (module Atom) state v ; state.offset

let write e v buffer offset len =
  try Ok (write_exn e v buffer offset len) with Binary_error.Write_error err -> Error err

let write_opt e v buffer offset len =
  try Some (write_exn e v buffer offset len) with Binary_error.Write_error _ -> None

let to_bytes_exn ?(buffer_size = 128) e v =
  match Encoding.classify e with
  | `Fixed n ->
      (* Preallocate the complete buffer *)
      let state =
        {buffer = Bytes.create n; offset = 0; allowed_bytes = Some n}
      in
      Binary_generic_writer.write_rec e (module Atom) state v ; state.buffer
  | `Dynamic | `Variable ->
      (* Preallocate a minimal buffer and let's not hardcode a
         limit to its extension. *)
      let state =
        {buffer = Bytes.create buffer_size; offset = 0; allowed_bytes = None}
      in
      Binary_generic_writer.write_rec e (module Atom) state v ;
      Bytes.sub state.buffer 0 state.offset

let to_bytes_opt ?buffer_size e v =
  try Some (to_bytes_exn ?buffer_size e v) with Binary_error.Write_error _ -> None

let to_bytes ?buffer_size e v =
  try Ok (to_bytes_exn ?buffer_size e v) with Binary_error.Write_error err -> Error err