(** Bitstream reader for Zstandard decompression.

    This module wraps the Bitstream library, translating exceptions
    to Zstd_error for consistent error handling. *)

(** Helper to wrap Bitstream operations and translate exceptions *)
let[@inline] wrap_truncated f =
  try f ()
  with Bitstream.End_of_stream ->
    raise (Constants.Zstd_error Constants.Truncated_input)

let[@inline] wrap_all f =
  try f ()
  with
  | Bitstream.End_of_stream ->
    raise (Constants.Zstd_error Constants.Truncated_input)
  | Bitstream.Invalid_state _ ->
    raise (Constants.Zstd_error Constants.Corruption)
  | Bitstream.Corrupted_stream _ ->
    raise (Constants.Zstd_error Constants.Corruption)

(** Forward bitstream reader - reads from start to end *)
module Forward = struct
  type t = Bitstream.Forward_reader.t

  let create src ~pos ~len =
    Bitstream.Forward_reader.create src ~pos ~len

  let of_bytes src =
    Bitstream.Forward_reader.of_bytes src

  let[@inline] remaining t =
    Bitstream.Forward_reader.remaining t

  let[@inline] is_byte_aligned t =
    Bitstream.Forward_reader.is_byte_aligned t

  let[@inline] read_bits t n =
    wrap_truncated (fun () -> Bitstream.Forward_reader.read_bits t n)

  let[@inline] read_byte t =
    wrap_all (fun () -> Bitstream.Forward_reader.read_byte t)

  let rewind_bits t n =
    wrap_truncated (fun () -> Bitstream.Forward_reader.rewind_bits t n)

  let align t =
    Bitstream.Forward_reader.align t

  let byte_position t =
    wrap_all (fun () -> Bitstream.Forward_reader.byte_position t)

  let get_bytes t n =
    wrap_all (fun () -> Bitstream.Forward_reader.get_bytes t n)

  let advance t n =
    wrap_all (fun () -> Bitstream.Forward_reader.advance t n)

  let sub t n =
    wrap_all (fun () -> Bitstream.Forward_reader.sub t n)

  let remaining_bytes t =
    wrap_all (fun () -> Bitstream.Forward_reader.remaining_bytes t)
end

(** Backward bitstream reader - reads from end to start.
    Used for FSE and Huffman coded streams. *)
module Backward = struct
  type t = Bitstream.Backward_reader.t

  let create src ~pos ~len =
    wrap_all (fun () -> Bitstream.Backward_reader.of_bytes src ~pos ~len)

  let of_bytes src ~pos ~len =
    create src ~pos ~len

  let[@inline] remaining t =
    Bitstream.Backward_reader.remaining t

  let[@inline] read_bits t n =
    Bitstream.Backward_reader.read_bits t n

  let[@inline] is_empty t =
    Bitstream.Backward_reader.is_empty t
end

(** Read little-endian integers from bytes *)
let[@inline] get_u16_le src pos =
  Bytes.get_uint16_le src pos