(** Bitstream utilities for jam/cue serialization *)

(** A bitstream writer *)
type writer = {
  buf: bytes ref;          (** Buffer for bits *)
  mutable bit_pos: int;    (** Current bit position *)
}

(** A bitstream reader *)
type reader = {
  buf: bytes;              (** Buffer to read from *)
  mutable bit_pos: int;    (** Current bit position *)
  len: int;                (** Length in bits *)
}

(** Create a new bitstream writer *)
let writer_create () = {
  buf = ref (Bytes.create 1024);
  bit_pos = 0;
}

(** Grow the writer buffer if needed *)
let writer_ensure w bits_needed =
  let bytes_needed = (w.bit_pos + bits_needed + 7) / 8 in
  if bytes_needed > Bytes.length !(w.buf) then begin
    let old_buf = !(w.buf) in
    let new_size = max (bytes_needed * 2) (Bytes.length old_buf * 2) in
    let new_buf = Bytes.create new_size in
    Bytes.blit old_buf 0 new_buf 0 (Bytes.length old_buf);
    w.buf := new_buf
  end

(** Write a single bit *)
let write_bit w bit =
  writer_ensure w 1;
  let byte_pos = w.bit_pos / 8 in
  let bit_off = w.bit_pos mod 8 in
  if bit then begin
    let old_byte = Bytes.get_uint8 !(w.buf) byte_pos in
    Bytes.set_uint8 !(w.buf) byte_pos (old_byte lor (1 lsl bit_off))
  end;
  w.bit_pos <- w.bit_pos + 1

(** Write multiple bits from a Z.t value *)
let write_bits w value nbits =
  writer_ensure w nbits;
  for i = 0 to nbits - 1 do
    let bit = Z.testbit value i in
    write_bit w bit
  done

(** Get the final bytes from a writer *)
let writer_to_bytes w =
  let byte_len = (w.bit_pos + 7) / 8 in
  Bytes.sub !(w.buf) 0 byte_len

(** Create a bitstream reader *)
let reader_create buf =
  {
    buf;
    bit_pos = 0;
    len = Bytes.length buf * 8;
  }

(** Read a single bit *)
let read_bit r =
  if r.bit_pos >= r.len then
    raise (Invalid_argument "read_bit: end of stream");
  let byte_pos = r.bit_pos / 8 in
  let bit_off = r.bit_pos mod 8 in
  let byte_val = Bytes.get_uint8 r.buf byte_pos in
  r.bit_pos <- r.bit_pos + 1;
  (byte_val lsr bit_off) land 1 = 1

(** Read multiple bits as a Z.t *)
let read_bits r nbits =
  let result = ref Z.zero in
  for i = 0 to nbits - 1 do
    if read_bit r then
      result := Z.logor !result (Z.shift_left Z.one i)
  done;
  !result

(** Peek at a bit without advancing *)
let peek_bit r =
  if r.bit_pos >= r.len then
    raise (Invalid_argument "peek_bit: end of stream");
  let byte_pos = r.bit_pos / 8 in
  let bit_off = r.bit_pos mod 8 in
  let byte_val = Bytes.get_uint8 r.buf byte_pos in
  (byte_val lsr bit_off) land 1 = 1

(** Get current bit position *)
let reader_pos r = r.bit_pos

(** Check if at end of stream *)
let reader_at_end r = r.bit_pos >= r.len