(** Tests for Bitstream library *)

let test_forward_reader_bytes () =
  let data = Bytes.of_string "\x12\x34\x56\x78" in
  let r = Bitstream.Forward_reader.of_bytes data in
  Alcotest.(check int) "byte 0" 0x12 (Bitstream.Forward_reader.read_byte r);
  Alcotest.(check int) "byte 1" 0x34 (Bitstream.Forward_reader.read_byte r);
  Alcotest.(check int) "byte 2" 0x56 (Bitstream.Forward_reader.read_byte r);
  Alcotest.(check int) "byte 3" 0x78 (Bitstream.Forward_reader.read_byte r);
  (* Reading past end raises End_of_stream *)
  Alcotest.check_raises "past end" Bitstream.End_of_stream (fun () ->
    ignore (Bitstream.Forward_reader.read_byte r))

let test_forward_reader_bits () =
  (* 0x12 = 0001_0010, 0x34 = 0011_0100 in little-endian bits:
     Reading 4 bits: 0010 = 2
     Reading 4 bits: 0001 = 1
     Reading 8 bits: 0011_0100 = 0x34 *)
  let data = Bytes.of_string "\x12\x34" in
  let r = Bitstream.Forward_reader.of_bytes data in
  Alcotest.(check int) "4 bits" 0x2 (Bitstream.Forward_reader.read_bits r 4);
  Alcotest.(check int) "4 bits" 0x1 (Bitstream.Forward_reader.read_bits r 4);
  Alcotest.(check int) "8 bits" 0x34 (Bitstream.Forward_reader.read_bits r 8)

let test_forward_reader_mixed () =
  let data = Bytes.of_string "\xFF\x00\xAB\xCD" in
  let r = Bitstream.Forward_reader.of_bytes data in
  Alcotest.(check int) "16 bits LE" 0x00FF (Bitstream.Forward_reader.read_bits r 16);
  Alcotest.(check int) "remaining bits" 16 (Bitstream.Forward_reader.remaining r);
  Alcotest.(check int) "remaining bytes" 2 (Bitstream.Forward_reader.remaining_bytes r);
  let rest = Bitstream.Forward_reader.get_bytes r 2 in
  Alcotest.(check string) "get_bytes" "\xAB\xCD" (Bytes.to_string rest)

let test_forward_reader_32bit () =
  let data = Bytes.of_string "\x78\x56\x34\x12" in
  let r = Bitstream.Forward_reader.of_bytes data in
  Alcotest.(check int) "32 bits LE" 0x12345678 (Bitstream.Forward_reader.read_bits r 32)

let test_forward_reader_rewind () =
  let data = Bytes.of_string "\x12\x34\x56" in
  let r = Bitstream.Forward_reader.of_bytes data in
  let v1 = Bitstream.Forward_reader.read_bits r 12 in
  Alcotest.(check int) "first read 12 bits" 0x412 v1;  (* 0x12 + lower 4 of 0x34 *)
  Bitstream.Forward_reader.rewind_bits r 4;
  Alcotest.(check int) "remaining after rewind" 16 (Bitstream.Forward_reader.remaining r);
  let v2 = Bitstream.Forward_reader.read_bits r 8 in
  Alcotest.(check int) "read after rewind" 0x34 v2  (* now at byte 1 *)

let test_forward_reader_align () =
  let data = Bytes.of_string "\xFF\xAA\xBB" in
  let r = Bitstream.Forward_reader.of_bytes data in
  let _ = Bitstream.Forward_reader.read_bits r 3 in
  Alcotest.(check bool) "not aligned" false (Bitstream.Forward_reader.is_byte_aligned r);
  Bitstream.Forward_reader.align r;
  Alcotest.(check bool) "aligned" true (Bitstream.Forward_reader.is_byte_aligned r);
  Alcotest.(check int) "after align" 0xAA (Bitstream.Forward_reader.read_byte r)

let test_forward_reader_sub () =
  let data = Bytes.of_string "\x12\x34\x56\x78\x9A" in
  let r = Bitstream.Forward_reader.of_bytes data in
  let _ = Bitstream.Forward_reader.read_byte r in
  let sub = Bitstream.Forward_reader.sub r 2 in
  Alcotest.(check int) "sub byte 0" 0x34 (Bitstream.Forward_reader.read_byte sub);
  Alcotest.(check int) "sub byte 1" 0x56 (Bitstream.Forward_reader.read_byte sub);
  Alcotest.(check int) "parent continues" 0x78 (Bitstream.Forward_reader.read_byte r)

let test_forward_writer_bytes () =
  let buf = Bytes.create 8 in
  let w = Bitstream.Forward_writer.of_bytes buf in
  Bitstream.Forward_writer.write_byte w 0x12;
  Bitstream.Forward_writer.write_byte w 0x34;
  let len = Bitstream.Forward_writer.finalize w in
  Alcotest.(check int) "length" 2 len;
  Alcotest.(check int) "byte 0" 0x12 (Bytes.get_uint8 buf 0);
  Alcotest.(check int) "byte 1" 0x34 (Bytes.get_uint8 buf 1)

let test_forward_writer_bits () =
  let buf = Bytes.create 8 in
  let w = Bitstream.Forward_writer.of_bytes buf in
  Bitstream.Forward_writer.write_bits w 0x2 4;  (* lower 4 bits *)
  Bitstream.Forward_writer.write_bits w 0x1 4;  (* upper 4 bits *)
  Bitstream.Forward_writer.write_bits w 0x34 8;
  let len = Bitstream.Forward_writer.finalize w in
  Alcotest.(check int) "length" 2 len;
  Alcotest.(check int) "byte 0" 0x12 (Bytes.get_uint8 buf 0);
  Alcotest.(check int) "byte 1" 0x34 (Bytes.get_uint8 buf 1)

let test_forward_writer_32bit () =
  let buf = Bytes.create 8 in
  let w = Bitstream.Forward_writer.of_bytes buf in
  Bitstream.Forward_writer.write_bits w 0x12345678 32;
  let len = Bitstream.Forward_writer.finalize w in
  Alcotest.(check int) "length" 4 len;
  Alcotest.(check int) "byte 0" 0x78 (Bytes.get_uint8 buf 0);
  Alcotest.(check int) "byte 1" 0x56 (Bytes.get_uint8 buf 1);
  Alcotest.(check int) "byte 2" 0x34 (Bytes.get_uint8 buf 2);
  Alcotest.(check int) "byte 3" 0x12 (Bytes.get_uint8 buf 3)

let test_forward_roundtrip () =
  (* Write various bit patterns, then read them back *)
  let buf = Bytes.create 64 in
  let w = Bitstream.Forward_writer.of_bytes buf in
  Bitstream.Forward_writer.write_bits w 0b101 3;
  Bitstream.Forward_writer.write_bits w 0b11001 5;
  Bitstream.Forward_writer.write_bits w 0xABCD 16;
  Bitstream.Forward_writer.write_bits w 0b1111 4;
  let len = Bitstream.Forward_writer.finalize w in

  let r = Bitstream.Forward_reader.create buf ~pos:0 ~len in
  Alcotest.(check int) "3 bits" 0b101 (Bitstream.Forward_reader.read_bits r 3);
  Alcotest.(check int) "5 bits" 0b11001 (Bitstream.Forward_reader.read_bits r 5);
  Alcotest.(check int) "16 bits" 0xABCD (Bitstream.Forward_reader.read_bits r 16);
  Alcotest.(check int) "4 bits" 0b1111 (Bitstream.Forward_reader.read_bits r 4)

let test_backward_roundtrip () =
  (* Backward streams are read in REVERSE order of writing.
     This matches FSE: encode in reverse order, decode in forward order. *)
  let w = Bitstream.Backward_writer.create 64 in
  Bitstream.Backward_writer.write_bits w 0b101 3;
  Bitstream.Backward_writer.write_bits w 0b11001 5;
  Bitstream.Backward_writer.write_bits w 0xAB 8;
  let data = Bitstream.Backward_writer.finalize w in

  let r = Bitstream.Backward_reader.of_bytes data ~pos:0 ~len:(Bytes.length data) in
  (* Read in REVERSE order (last written = first read) *)
  Alcotest.(check int) "8 bits (written last)" 0xAB (Bitstream.Backward_reader.read_bits r 8);
  Alcotest.(check int) "5 bits" 0b11001 (Bitstream.Backward_reader.read_bits r 5);
  Alcotest.(check int) "3 bits (written first)" 0b101 (Bitstream.Backward_reader.read_bits r 3)

let test_backward_reader_peek () =
  (* For backward streams, bits are read from MSB to LSB within accumulated data.
     0x5A = 0101_1010 binary. Reading 4 bits at a time gives:
     - First 4 bits (high nibble): 0101 = 0x5
     - Last 4 bits (low nibble): 1010 = 0xA *)
  let w = Bitstream.Backward_writer.create 64 in
  Bitstream.Backward_writer.write_bits w 0x5A 8;
  let data = Bitstream.Backward_writer.finalize w in

  let r = Bitstream.Backward_reader.of_bytes data ~pos:0 ~len:(Bytes.length data) in
  Alcotest.(check int) "peek 4 (high nibble)" 0x5 (Bitstream.Backward_reader.peek_bits r 4);
  Alcotest.(check int) "peek 4 again" 0x5 (Bitstream.Backward_reader.peek_bits r 4);
  Alcotest.(check int) "read 4" 0x5 (Bitstream.Backward_reader.read_bits r 4);
  Alcotest.(check int) "read 4 (low nibble)" 0xA (Bitstream.Backward_reader.read_bits r 4)

let test_backward_is_empty () =
  let w = Bitstream.Backward_writer.create 64 in
  Bitstream.Backward_writer.write_bits w 0xFF 8;
  let data = Bitstream.Backward_writer.finalize w in

  let r = Bitstream.Backward_reader.of_bytes data ~pos:0 ~len:(Bytes.length data) in
  Alcotest.(check bool) "not empty" false (Bitstream.Backward_reader.is_empty r);
  let _ = Bitstream.Backward_reader.read_bits r 8 in
  Alcotest.(check bool) "empty after read" true (Bitstream.Backward_reader.is_empty r)

let test_backward_empty_stream () =
  (* Empty stream should raise End_of_stream *)
  Alcotest.check_raises "empty stream" Bitstream.End_of_stream (fun () ->
    ignore (Bitstream.Backward_reader.of_bytes (Bytes.empty) ~pos:0 ~len:0))

let test_backward_invalid_padding () =
  (* Zero byte has no padding marker - should raise *)
  let data = Bytes.of_string "\x00" in
  Alcotest.check_raises "zero padding" (Bitstream.Corrupted_stream "invalid padding marker") (fun () ->
    ignore (Bitstream.Backward_reader.of_bytes data ~pos:0 ~len:1))

let test_edge_cases () =
  (* Zero bits *)
  let buf = Bytes.create 8 in
  let w = Bitstream.Forward_writer.of_bytes buf in
  Bitstream.Forward_writer.write_bits w 0 0;
  let len = Bitstream.Forward_writer.finalize w in
  Alcotest.(check int) "zero bits" 0 len;

  (* Read zero bits *)
  let data = Bytes.of_string "\xFF" in
  let r = Bitstream.Forward_reader.of_bytes data in
  Alcotest.(check int) "read 0 bits" 0 (Bitstream.Forward_reader.read_bits r 0);
  Alcotest.(check int) "byte still available" 0xFF (Bitstream.Forward_reader.read_byte r)

let test_not_aligned_errors () =
  let data = Bytes.of_string "\xFF\xAA" in
  let r = Bitstream.Forward_reader.of_bytes data in
  let _ = Bitstream.Forward_reader.read_bits r 3 in

  Alcotest.check_raises "read_byte not aligned"
    (Bitstream.Invalid_state "read_byte: not byte aligned")
    (fun () -> ignore (Bitstream.Forward_reader.read_byte r));

  Alcotest.check_raises "byte_position not aligned"
    (Bitstream.Invalid_state "byte_position: not byte aligned")
    (fun () -> ignore (Bitstream.Forward_reader.byte_position r));

  Alcotest.check_raises "remaining_bytes not aligned"
    (Bitstream.Invalid_state "remaining_bytes: not byte aligned")
    (fun () -> ignore (Bitstream.Forward_reader.remaining_bytes r))

let tests = [
  "forward reader bytes", `Quick, test_forward_reader_bytes;
  "forward reader bits", `Quick, test_forward_reader_bits;
  "forward reader mixed", `Quick, test_forward_reader_mixed;
  "forward reader 32bit", `Quick, test_forward_reader_32bit;
  "forward reader rewind", `Quick, test_forward_reader_rewind;
  "forward reader align", `Quick, test_forward_reader_align;
  "forward reader sub", `Quick, test_forward_reader_sub;
  "forward writer bytes", `Quick, test_forward_writer_bytes;
  "forward writer bits", `Quick, test_forward_writer_bits;
  "forward writer 32bit", `Quick, test_forward_writer_32bit;
  "forward roundtrip", `Quick, test_forward_roundtrip;
  "backward roundtrip", `Quick, test_backward_roundtrip;
  "backward reader peek", `Quick, test_backward_reader_peek;
  "backward is_empty", `Quick, test_backward_is_empty;
  "backward empty stream", `Quick, test_backward_empty_stream;
  "backward invalid padding", `Quick, test_backward_invalid_padding;
  "edge cases", `Quick, test_edge_cases;
  "not aligned errors", `Quick, test_not_aligned_errors;
]

let () =
  Alcotest.run "Bitstream" [
    "bitstream", tests;
  ]