// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! Reading and writing a big-endian bit stream, which is what a FLAC frame is:
//! only the first subframe of a frame starts on a byte boundary, and
//! everything after it is packed against whatever came before (RFC 9639,
//! Section 9.2).
//!
//! Both types carry the frame's two CRCs (see `crc.zig`) and update them as
//! bytes go past. That is why `BitReader` pulls exactly the bytes it needs and
//! not a machine word at a time: at any byte-aligned point the CRCs then
//! account for precisely the bytes consumed so far, which is what the frame
//! header's check and the frame footer's check are each defined over. Where
//! the format wants a snapshot -- after the header, and after the padding that
//! byte-aligns the last subframe -- the caller takes one and reads the stored
//! check afterwards, which does not disturb the value already taken.

const std = @import("std");
const assert = std.debug.assert;

const crc = @import("crc.zig");

/// The widest field either type will move in one call. The accumulator holds
/// 64 bits and up to 7 of them are left over from the previous call, so 57 is
/// what remains. Nothing in a FLAC frame is wider: the largest is a 33-bit
/// unencoded sample, which is a 32-bit side channel.
pub const max_field_bits = 57;

fn lowMask(n: u7) u64 {
    if (n >= 64) return ~@as(u64, 0);
    return (@as(u64, 1) << @intCast(n)) - 1;
}

pub const BitReader = struct {
    input: *std.Io.Reader,

    /// Bits pulled but not yet consumed, right-aligned, oldest bit highest.
    acc: u64 = 0,
    /// How many of `acc`'s low bits are pending. Always below 8 on return
    /// from any public function.
    nbits: u7 = 0,

    /// Bytes pulled from `input`. At a byte-aligned point this is exactly the
    /// number of bytes consumed.
    pos: u64 = 0,

    crc8: crc.Crc8 = .init(),
    crc16: crc.Crc16 = .init(),

    pub const Error = std.Io.Reader.Error || error{
        /// A unary code ran past any length a valid stream can contain.
        MalformedBitstream,
    };

    pub fn init(input: *std.Io.Reader) BitReader {
        return .{ .input = input };
    }

    /// Start both checks over. The decoder calls this having consumed nothing
    /// of the frame, so that both cover the sync code as the format requires.
    pub fn resetCrc(r: *BitReader) void {
        r.crc8 = .init();
        r.crc16 = .init();
    }

    fn pull(r: *BitReader) Error!void {
        const b = try r.input.takeByte();
        r.crc8.update(&.{b});
        r.crc16.update(&.{b});
        r.acc = (r.acc << 8) | b;
        r.nbits += 8;
        r.pos += 1;
    }

    /// The next `n` bits as an unsigned number, most significant bit first.
    pub fn take(r: *BitReader, n: u7) Error!u64 {
        assert(n <= max_field_bits);
        if (n == 0) return 0;
        while (r.nbits < n) try r.pull();
        const rest: u7 = r.nbits - n;
        const value = (r.acc >> @intCast(rest)) & lowMask(n);
        r.nbits = rest;
        r.acc &= lowMask(rest);
        return value;
    }

    /// The next `n` bits as a two's complement signed number, which is how
    /// FLAC stores samples, predictor coefficients and the prediction shift.
    pub fn takeSigned(r: *BitReader, n: u7) Error!i64 {
        if (n == 0) return 0;
        const value = try r.take(n);
        const sign_shift: u6 = @intCast(64 - @as(u7, n));
        return @as(i64, @bitCast(value << sign_shift)) >> sign_shift;
    }

    /// The count of zero bits before the next one bit, having consumed both.
    /// This is the most-significant half of a Rice code word and the way the
    /// number of wasted bits is stored.
    pub fn takeUnary(r: *BitReader) Error!u32 {
        var zeros: u64 = 0;
        while (true) {
            if (r.nbits == 0) try r.pull();
            if (r.acc == 0) {
                zeros += r.nbits;
                r.nbits = 0;
                if (zeros > std.math.maxInt(u32)) return error.MalformedBitstream;
                continue;
            }
            // The highest set bit of what is pending is the terminating one.
            const highest: u7 = @intCast(63 - @clz(r.acc));
            zeros += r.nbits - 1 - highest;
            r.nbits = highest;
            r.acc &= lowMask(highest);
            if (zeros > std.math.maxInt(u32)) return error.MalformedBitstream;
            return @intCast(zeros);
        }
    }

    pub fn isByteAligned(r: *const BitReader) bool {
        return r.nbits == 0;
    }

    /// Discard bits up to the next byte boundary. The bits themselves have
    /// already been counted towards both CRCs, which is what the frame footer
    /// requires of the padding it follows.
    pub fn alignToByte(r: *BitReader) void {
        r.nbits = 0;
        r.acc = 0;
    }
};

pub const BitWriter = struct {
    output: *std.Io.Writer,

    /// Bits produced but not yet whole bytes, right-aligned.
    acc: u64 = 0,
    nbits: u7 = 0,

    /// Bytes handed to `output`.
    pos: u64 = 0,

    crc8: crc.Crc8 = .init(),
    crc16: crc.Crc16 = .init(),

    pub const Error = std.Io.Writer.Error;

    pub fn init(output: *std.Io.Writer) BitWriter {
        return .{ .output = output };
    }

    pub fn resetCrc(w: *BitWriter) void {
        w.crc8 = .init();
        w.crc16 = .init();
    }

    fn emit(w: *BitWriter, b: u8) Error!void {
        w.crc8.update(&.{b});
        w.crc16.update(&.{b});
        w.pos += 1;
        try w.output.writeByte(b);
    }

    /// The low `n` bits of `value`, most significant bit first.
    pub fn put(w: *BitWriter, n: u7, value: u64) Error!void {
        assert(n <= max_field_bits);
        if (n == 0) return;
        w.acc = (w.acc << @intCast(n)) | (value & lowMask(n));
        w.nbits += n;
        while (w.nbits >= 8) {
            w.nbits -= 8;
            try w.emit(@truncate(w.acc >> @intCast(w.nbits)));
        }
        w.acc &= lowMask(w.nbits);
    }

    /// A signed number in `n` bits, two's complement. The caller is
    /// responsible for the value fitting: one that does not would have its
    /// surplus bits dropped and be read back as some other number, so it is
    /// an assertion rather than a truncation.
    pub fn putSigned(w: *BitWriter, n: u7, value: i64) Error!void {
        assert(n == 0 or n > 63 or
            (value >= -(@as(i64, 1) << @intCast(n - 1)) and value < (@as(i64, 1) << @intCast(n - 1))));
        try w.put(n, @bitCast(value));
    }

    /// `zeros` zero bits followed by a one bit.
    pub fn putUnary(w: *BitWriter, zeros: u32) Error!void {
        var left = zeros;
        while (left >= 32) : (left -= 32) try w.put(32, 0);
        try w.put(@as(u7, @intCast(left)) + 1, 1);
    }

    pub fn isByteAligned(w: *const BitWriter) bool {
        return w.nbits == 0;
    }

    /// Zero bits up to the next byte boundary, which is what a frame puts
    /// between its last subframe and its footer.
    pub fn padToByte(w: *BitWriter) Error!void {
        if (w.nbits != 0) try w.put(8 - w.nbits, 0);
    }
};

const testing = std.testing;

test "reading fields the specification breaks down" {
    // The frame header of RFC 9639's first example file, field by field as
    // Table 29 gives it.
    var input: std.Io.Reader = .fixed(&.{ 0xff, 0xf8, 0x69, 0x18, 0x00, 0x00, 0xbf });
    var r: BitReader = .init(&input);

    try testing.expectEqual(@as(u64, 0b111111111111100), try r.take(15));
    try testing.expectEqual(@as(u64, 0), try r.take(1));
    try testing.expectEqual(@as(u64, 0b0110), try r.take(4));
    try testing.expectEqual(@as(u64, 0b1001), try r.take(4));
    try testing.expectEqual(@as(u64, 0b0001), try r.take(4));
    try testing.expectEqual(@as(u64, 0b100), try r.take(3));
    try testing.expectEqual(@as(u64, 0), try r.take(1));
    try testing.expect(r.isByteAligned());
    try testing.expectEqual(@as(u64, 0x00), try r.take(8));
    try testing.expectEqual(@as(u64, 0x00), try r.take(8));

    // Six bytes have been consumed, and the check over them is the one the
    // seventh byte holds.
    try testing.expectEqual(@as(u64, 6), r.pos);
    try testing.expectEqual(@as(u8, 0xbf), r.crc8.final());
}

test "signed fields sign-extend" {
    var input: std.Io.Reader = .fixed(&.{ 0xff, 0x80, 0x7f });
    var r: BitReader = .init(&input);
    try testing.expectEqual(@as(i64, -1), try r.takeSigned(8));
    try testing.expectEqual(@as(i64, -128), try r.takeSigned(8));
    try testing.expectEqual(@as(i64, 127), try r.takeSigned(8));
}

test "unary codes spanning byte boundaries" {
    // 0b00000000 0b00000000 0b00100000: eighteen zeros, a one, then five bits
    // left over in the byte.
    var input: std.Io.Reader = .fixed(&.{ 0x00, 0x00, 0x20, 0xff });
    var r: BitReader = .init(&input);
    try testing.expectEqual(@as(u32, 18), try r.takeUnary());
    try testing.expectEqual(@as(u64, 0), try r.take(5));
    try testing.expectEqual(@as(u32, 0), try r.takeUnary());
}

test "the Rice code word from Section 9.2.7.2" {
    // Rice parameter 3 and a folded residual of 38 make 0b00001110.
    var buffer: [8]u8 = undefined;
    var output: std.Io.Writer = .fixed(&buffer);
    var w: BitWriter = .init(&output);
    try w.putUnary(4);
    try w.put(3, 6);
    try testing.expect(w.isByteAligned());
    try testing.expectEqual(@as(u8, 0b00001110), buffer[0]);

    var input: std.Io.Reader = .fixed(buffer[0..1]);
    var r: BitReader = .init(&input);
    try testing.expectEqual(@as(u32, 4), try r.takeUnary());
    try testing.expectEqual(@as(u64, 6), try r.take(3));
}

test "writing back what was read" {
    var prng: std.Random.DefaultPrng = .init(0x466c6143);
    const random = prng.random();

    var widths: [512]u7 = undefined;
    var values: [512]u64 = undefined;
    for (&widths, &values) |*width, *value| {
        width.* = random.intRangeAtMost(u7, 0, max_field_bits);
        value.* = random.int(u64) & lowMask(width.*);
    }

    var buffer: [4096]u8 = undefined;
    var output: std.Io.Writer = .fixed(&buffer);
    var w: BitWriter = .init(&output);
    for (widths, values) |width, value| try w.put(width, value);
    try w.padToByte();

    var input: std.Io.Reader = .fixed(buffer[0..@intCast(w.pos)]);
    var r: BitReader = .init(&input);
    for (widths, values) |width, value| {
        try testing.expectEqual(value, try r.take(width));
    }

    // Reader and writer agree byte for byte, so their checks agree too.
    try testing.expectEqual(w.crc16.final(), r.crc16.final());
}