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

//! What the format allows, and the pieces of it that both the frame header and
//! the subframes need to name.

const std = @import("std");

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

/// The four bytes every FLAC stream starts with (RFC 9639, Section 6).
pub const signature = "fLaC";

/// The bounds the format itself imposes. A decoder holds an input to these;
/// an encoder holds itself to them and to the streamable subset besides.
pub const limits = struct {
    /// Section 8.2: the streaminfo block stores block sizes in 16 bits and
    /// forbids anything under 16 samples.
    pub const min_block_size = 16;
    pub const max_block_size = 65535;

    /// Section 8.2, and the three channel bits of a frame header.
    pub const min_channels = 1;
    pub const max_channels = 8;

    /// Section 8.2. The frame header can only name a subset of these.
    pub const min_bits_per_sample = 4;
    pub const max_bits_per_sample = 32;

    /// Section 8.2: a 20-bit field, so no more.
    pub const max_sample_rate = (1 << 20) - 1;

    /// Section 8.2: a 36-bit field.
    pub const max_total_samples = (1 << 36) - 1;

    /// Section 9.2.5 and Section 9.2.6.
    pub const max_fixed_order = 4;
    pub const max_lpc_order = 32;

    /// Section 9.2.6: the precision is stored as four bits holding one less
    /// than it, and all ones is forbidden.
    pub const min_qlp_precision = 1;
    pub const max_qlp_precision = 15;

    /// Section 9.2.6: the prediction right shift is five bits, signed two's
    /// complement, and MUST NOT be negative -- so 15 is the largest there is,
    /// not 31 (Appendix B.4).
    pub const max_lpc_shift = 15;

    /// Section 9.2.7: four bits of partition order.
    pub const max_partition_order = 15;

    /// Section 8.1: a metadata block's length is a 24-bit field.
    pub const max_metadata_length = (1 << 24) - 1;

    /// Section 7. Everything above holds for any stream; these two hold only
    /// for the streamable subset.
    pub const subset_max_block_size = 16384;
    pub const subset_max_block_size_48k = 4608;
    pub const subset_max_lpc_order_48k = 12;
    pub const subset_max_partition_order = 8;
};

/// How the subframes of a frame map onto its channels (RFC 9639,
/// Section 9.1.3). Two-channel audio has three further options, each of which
/// codes one channel as the difference of the two -- the side channel, which
/// needs one bit more depth than the frame declares (Section 4.2).
pub const ChannelAssignment = union(enum) {
    /// One subframe per channel, each holding that channel's samples.
    independent: u4,
    /// Left in the first subframe, left minus right in the second.
    left_side,
    /// Left minus right in the first subframe, right in the second.
    right_side,
    /// (Left + right) >> 1 in the first subframe, left minus right in the
    /// second.
    mid_side,

    pub fn fromBits(v: u4) error{ReservedChannelAssignment}!ChannelAssignment {
        return switch (v) {
            0b0000...0b0111 => .{ .independent = v + 1 },
            0b1000 => .left_side,
            0b1001 => .right_side,
            0b1010 => .mid_side,
            else => error.ReservedChannelAssignment,
        };
    }

    pub fn toBits(a: ChannelAssignment) u4 {
        return switch (a) {
            .independent => |n| @intCast(n - 1),
            .left_side => 0b1000,
            .right_side => 0b1001,
            .mid_side => 0b1010,
        };
    }

    pub fn channelCount(a: ChannelAssignment) u4 {
        return switch (a) {
            .independent => |n| n,
            else => 2,
        };
    }

    /// Which of a frame's subframes carry a side channel, and so are coded one
    /// bit deeper than the frame header says.
    pub fn isSideChannel(a: ChannelAssignment, channel: usize) bool {
        return switch (a) {
            .independent => false,
            .left_side, .mid_side => channel == 1,
            .right_side => channel == 0,
        };
    }
};

/// What a subframe's six type bits say it is (RFC 9639, Section 9.2.1).
pub const SubframeType = union(enum) {
    /// One sample, repeated for the whole block.
    constant,
    /// Every sample, unencoded.
    verbatim,
    /// One of the five predictors the format defines, order 0 to 4.
    fixed: u3,
    /// A predictor the subframe carries with it, order 1 to 32.
    lpc: u6,

    pub fn fromBits(v: u6) error{ReservedSubframeType}!SubframeType {
        return switch (v) {
            0b000000 => .constant,
            0b000001 => .verbatim,
            0b001000...0b001100 => .{ .fixed = @intCast(v - 8) },
            0b100000...0b111111 => .{ .lpc = @intCast(v - 31) },
            else => error.ReservedSubframeType,
        };
    }

    pub fn toBits(t: SubframeType) u6 {
        return switch (t) {
            .constant => 0b000000,
            .verbatim => 0b000001,
            .fixed => |order| @as(u6, order) + 8,
            .lpc => |order| @as(u6, order) + 31,
        };
    }

    /// How many warm-up samples the subframe stores unencoded ahead of its
    /// residual, which is one per order of the predictor.
    pub fn predictorOrder(t: SubframeType) u6 {
        return switch (t) {
            .constant, .verbatim => 0,
            .fixed => |order| order,
            .lpc => |order| order,
        };
    }
};

/// The frame or sample number in a frame header, which is stored in the
/// variable-length code of UTF-8 extended to 36 bits and seven octets
/// (RFC 9639, Section 9.1.5). It counts frames or samples, not characters, so
/// a general-purpose UTF-8 codec is no use for it.
pub const coded_number = struct {
    pub const max = (1 << 36) - 1;

    pub const Error = error{MalformedCodedNumber};

    /// The octet count a value needs, which is the shortest form that holds
    /// it. Length 1 covers 7 bits and each further octet adds 5 then 6.
    pub fn length(value: u64) error{CodedNumberTooLarge}!u3 {
        return if (value < 0x80) 1 else if (value < 0x800) 2 else if (value < 0x10000)
            3
        else if (value < 0x200000)
            4
        else if (value < 0x4000000)
            5
        else if (value < 0x80000000)
            6
        else if (value <= max)
            7
        else
            error.CodedNumberTooLarge;
    }

    pub fn write(w: *bits.BitWriter, value: u64) (bits.BitWriter.Error || error{CodedNumberTooLarge})!void {
        const n = try length(value);
        if (n == 1) {
            try w.put(8, value);
            return;
        }
        // A lead octet of n ones, a zero, and what room is left; then n-1
        // continuation octets of 0b10 and six bits each, most significant
        // first.
        const payload_bits: u7 = @as(u7, n - 1) * 6;
        const lead_bits: u7 = 7 - n;
        try w.put(n, (@as(u64, 1) << @intCast(n)) - 1);
        try w.put(1, 0);
        try w.put(lead_bits, value >> @intCast(payload_bits));
        var shift: u7 = payload_bits;
        while (shift > 0) {
            shift -= 6;
            try w.put(2, 0b10);
            try w.put(6, value >> @intCast(shift));
        }
    }

    pub fn read(r: *bits.BitReader) (bits.BitReader.Error || Error)!u64 {
        const lead = try r.take(8);
        if (lead < 0x80) return lead;
        // Between one and seven leading ones; a lone continuation octet and an
        // octet of all ones are both outside the table.
        const ones: u4 = @intCast(@clz(~@as(u8, @intCast(lead))));
        if (ones < 2 or ones > 7) return error.MalformedCodedNumber;
        const lead_bits: u4 = @intCast(8 - ones - 1);
        var value: u64 = lead & ((@as(u64, 1) << lead_bits) - 1);
        for (1..ones) |_| {
            const cont = try r.take(8);
            if (cont & 0xc0 != 0x80) return error.MalformedCodedNumber;
            value = (value << 6) | (cont & 0x3f);
        }
        return value;
    }
};

const testing = std.testing;

test "the coded number for 51 billion samples" {
    // RFC 9639, Section 9.1.5 spells this one out octet by octet.
    const value: u64 = 51_000_000_000;
    const expected = [_]u8{ 0xfe, 0xaf, 0x9f, 0xb5, 0xa3, 0xb8, 0x80 };

    var buffer: [8]u8 = undefined;
    var output: std.Io.Writer = .fixed(&buffer);
    var w: bits.BitWriter = .init(&output);
    try coded_number.write(&w, value);
    try testing.expect(w.isByteAligned());
    try testing.expectEqualSlices(u8, &expected, buffer[0..7]);

    var input: std.Io.Reader = .fixed(&expected);
    var r: bits.BitReader = .init(&input);
    try testing.expectEqual(value, try coded_number.read(&r));
}

test "coded numbers round-trip at every length boundary" {
    const boundaries = [_]u64{
        0,         0x7f,       0x80,           0x7ff,       0x800,
        0xffff,    0x10000,    0x1fffff,       0x200000,    0x3ffffff,
        0x4000000, 0x7fffffff, 0x80000000,     0xfffffffff, 1,
        2,         127,        51_000_000_000,
    };
    for (boundaries) |value| {
        var buffer: [8]u8 = undefined;
        var output: std.Io.Writer = .fixed(&buffer);
        var w: bits.BitWriter = .init(&output);
        try coded_number.write(&w, value);
        try testing.expect(w.isByteAligned());
        try testing.expectEqual(@as(u64, try coded_number.length(value)), w.pos);

        var input: std.Io.Reader = .fixed(buffer[0..@intCast(w.pos)]);
        var r: bits.BitReader = .init(&input);
        try testing.expectEqual(value, try coded_number.read(&r));
    }
}

test "coded numbers reject octets outside the table" {
    for ([_]u8{ 0x80, 0xbf, 0xff }) |lead| {
        var input: std.Io.Reader = .fixed(&.{ lead, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80 });
        var r: bits.BitReader = .init(&input);
        try testing.expectError(error.MalformedCodedNumber, coded_number.read(&r));
    }

    // A lead octet promising a continuation that is not one.
    var input: std.Io.Reader = .fixed(&.{ 0xc2, 0x41 });
    var r: bits.BitReader = .init(&input);
    try testing.expectError(error.MalformedCodedNumber, coded_number.read(&r));
}

test "channel assignments and subframe types round-trip" {
    for (0..0b1011) |v| {
        const a = try ChannelAssignment.fromBits(@intCast(v));
        try testing.expectEqual(@as(u4, @intCast(v)), a.toBits());
    }
    for (0b1011..0b10000) |v| {
        try testing.expectError(error.ReservedChannelAssignment, ChannelAssignment.fromBits(@intCast(v)));
    }

    for (0..0b1000000) |v| {
        const t = SubframeType.fromBits(@intCast(v)) catch continue;
        try testing.expectEqual(@as(u6, @intCast(v)), t.toBits());
    }
    try testing.expectEqual(SubframeType{ .fixed = 4 }, try SubframeType.fromBits(0b001100));
    try testing.expectEqual(SubframeType{ .lpc = 1 }, try SubframeType.fromBits(0b100000));
    try testing.expectEqual(SubframeType{ .lpc = 32 }, try SubframeType.fromBits(0b111111));
}