(* See buffer.mli for the contract and the alignment invariant. *)

type bytes_ba =
  (int, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t

type bigarray = bytes_ba

(* [data] is a view into a larger allocation, positioned so its first byte is
   64-byte aligned (for buffers from [create]). [read_only] is set on [sub]
   views, which share storage and must not be mutated through this handle. *)
type t = { data : bytes_ba; len : int; read_only : bool }

let alignment = 64

(* The raw address of a Bigarray's first byte; see buffer_stubs.c. *)
external address_of : bytes_ba -> nativeint = "arrow_buffer_address"

(* Padding (0 .. alignment-1) to add to [addr] to reach the next multiple of
   [alignment]. *)
let padding_to_align addr =
  let m = Nativeint.to_int (Nativeint.rem addr (Nativeint.of_int alignment)) in
  if m = 0 then 0 else alignment - m

let create ~len =
  if len < 0 then invalid_arg "Arrow.Buffer.create: negative length";
  (* Over-allocate by a full alignment so any starting address can be shifted
     up to a 64-byte boundary with [len] bytes still in range. *)
  let raw =
    Bigarray.Array1.create Bigarray.int8_unsigned Bigarray.c_layout
      (len + alignment)
  in
  let pad = padding_to_align (address_of raw) in
  let data = Bigarray.Array1.sub raw pad len in
  Bigarray.Array1.fill data 0;
  { data; len; read_only = false }

let of_bigarray ?(read_only = false) data =
  { data; len = Bigarray.Array1.dim data; read_only }

let length t = t.len
let base_address t = address_of t.data

let is_aligned t =
  Nativeint.equal
    (Nativeint.rem (base_address t) (Nativeint.of_int alignment))
    0n

let get t i = Bigarray.Array1.get t.data i

let set t i v =
  if t.read_only then invalid_arg "Arrow.Buffer.set: read-only view";
  Bigarray.Array1.set t.data i v

(* Bounds check shared by every blit; raising up front avoids a partial copy. *)
let check_blit ~name ~src_len ~src_pos ~dst_len ~dst_pos ~len =
  if
    len < 0 || src_pos < 0 || dst_pos < 0
    || src_pos + len > src_len
    || dst_pos + len > dst_len
  then invalid_arg name

let blit_from_bytes ~src ~src_pos t ~dst_pos ~len =
  if t.read_only then invalid_arg "Arrow.Buffer.blit_from_bytes: read-only view";
  check_blit ~name:"Arrow.Buffer.blit_from_bytes" ~src_len:(Bytes.length src)
    ~src_pos ~dst_len:t.len ~dst_pos ~len;
  for i = 0 to len - 1 do
    Bigarray.Array1.unsafe_set t.data (dst_pos + i)
      (Char.code (Bytes.unsafe_get src (src_pos + i)))
  done

let blit_from_string ~src ~src_pos t ~dst_pos ~len =
  if t.read_only then
    invalid_arg "Arrow.Buffer.blit_from_string: read-only view";
  check_blit ~name:"Arrow.Buffer.blit_from_string" ~src_len:(String.length src)
    ~src_pos ~dst_len:t.len ~dst_pos ~len;
  for i = 0 to len - 1 do
    Bigarray.Array1.unsafe_set t.data (dst_pos + i)
      (Char.code (String.unsafe_get src (src_pos + i)))
  done

let blit_to_bytes t ~src_pos ~dst ~dst_pos ~len =
  check_blit ~name:"Arrow.Buffer.blit_to_bytes" ~src_len:t.len ~src_pos
    ~dst_len:(Bytes.length dst) ~dst_pos ~len;
  for i = 0 to len - 1 do
    Bytes.unsafe_set dst (dst_pos + i)
      (Char.unsafe_chr (Bigarray.Array1.unsafe_get t.data (src_pos + i)))
  done

let blit ~src ~src_pos ~dst ~dst_pos ~len =
  if dst.read_only then invalid_arg "Arrow.Buffer.blit: read-only view";
  check_blit ~name:"Arrow.Buffer.blit" ~src_len:src.len ~src_pos
    ~dst_len:dst.len ~dst_pos ~len;
  (* One memcpy (memmove semantics when the ranges overlap). *)
  Bigarray.Array1.blit
    (Bigarray.Array1.sub src.data src_pos len)
    (Bigarray.Array1.sub dst.data dst_pos len)

let sub t ~pos ~len =
  if pos < 0 || len < 0 || pos + len > t.len then invalid_arg "Arrow.Buffer.sub";
  { data = Bigarray.Array1.sub t.data pos len; len; read_only = true }