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

//! Enough of RIFF/WAVE to get uncompressed samples in and out.
//!
//! This is not part of the FLAC format and nothing else in the library depends
//! on it. It is here because a caller encoding audio almost always has it in a
//! WAVE file, and writing the same two hundred lines again is nobody's idea of
//! a good time.
//!
//! What is understood: `WAVE_FORMAT_PCM` and `WAVE_FORMAT_EXTENSIBLE` holding
//! PCM, at 8, 16, 24 or 32 bits per sample. Eight-bit WAVE samples are
//! unsigned by definition and are shifted into the signed range on the way in
//! and back on the way out; every other width is signed little-endian already.
//! Chunks that are not `fmt ` or `data` are skipped.
//!
//! What is not: floating-point samples, A-law and mu-law, RF64 for files past
//! four gigabytes, and an extensible header whose valid bit count differs from
//! its container's -- the container's width is what is used.

const std = @import("std");

pub const ReadError = std.Io.Reader.Error || error{
    NotWave,
    UnsupportedWaveFormat,
    MalformedWave,
};

pub const WriteError = std.Io.Writer.Error || error{WaveTooLarge};

const format_pcm: u16 = 0x0001;
const format_extensible: u16 = 0xfffe;

/// What a WAVE file says its samples are.
pub const Format = struct {
    sample_rate: u32,
    channels: u8,
    /// 8, 16, 24 or 32.
    bits_per_sample: u8,
    /// Interchannel samples, worked out from the size of the data chunk.
    total_samples: u64 = 0,

    pub fn bytesPerSample(f: Format) u8 {
        return f.bits_per_sample / 8;
    }

    pub fn blockAlign(f: Format) u16 {
        return @as(u16, f.channels) * f.bytesPerSample();
    }
};

/// The size of the header `writeHeader` produces: `RIFF`, a 16-byte `fmt `
/// chunk, and the `data` chunk header.
pub const header_length = 44;

/// Where the `RIFF` chunk's size field is, so that a caller who could not know
/// it in advance can seek back and put it right.
pub const riff_size_offset = 4;
/// Where the `data` chunk's size field is, likewise.
pub const data_size_offset = 40;

pub const Reader = struct {
    input: *std.Io.Reader,
    format: Format,
    /// Interchannel samples not yet read. Meaningless when `unknown_length`.
    remaining: u64,
    /// Set where the data chunk does not state a usable size, which is what a
    /// WAVE file written to a pipe looks like: the size fields cannot be
    /// filled in after the fact, so they are left at zero or at all ones.
    /// Reading then runs to the end of the stream instead.
    unknown_length: bool,

    pub fn init(input: *std.Io.Reader) ReadError!Reader {
        const riff = input.takeArray(4) catch return error.NotWave;
        if (!std.mem.eql(u8, riff, "RIFF")) return error.NotWave;
        _ = input.takeInt(u32, .little) catch return error.NotWave;
        const wave = input.takeArray(4) catch return error.NotWave;
        if (!std.mem.eql(u8, wave, "WAVE")) return error.NotWave;

        var format: ?Format = null;
        while (true) {
            const id = input.takeArray(4) catch |err| switch (err) {
                error.EndOfStream => return error.MalformedWave,
                error.ReadFailed => return error.ReadFailed,
            };
            const size = try input.takeInt(u32, .little);

            if (std.mem.eql(u8, id, "fmt ")) {
                format = try readFormat(input, size);
            } else if (std.mem.eql(u8, id, "data")) {
                const f = format orelse return error.MalformedWave;
                const align_bytes = f.blockAlign();
                if (align_bytes == 0) return error.MalformedWave;
                const unknown = size == 0 or size == std.math.maxInt(u32);
                var complete = f;
                complete.total_samples = if (unknown) 0 else size / align_bytes;
                return .{
                    .input = input,
                    .format = complete,
                    .remaining = complete.total_samples,
                    .unknown_length = unknown,
                };
            } else {
                // Chunks are padded to an even length, and the pad byte is not
                // counted in the size.
                try input.discardAll(size + (size & 1));
            }
        }
    }

    fn readFormat(input: *std.Io.Reader, size: u32) ReadError!Format {
        if (size < 16) return error.MalformedWave;
        var tag = try input.takeInt(u16, .little);
        const channels = try input.takeInt(u16, .little);
        const sample_rate = try input.takeInt(u32, .little);
        _ = try input.takeInt(u32, .little); // bytes per second, which is redundant
        _ = try input.takeInt(u16, .little); // block align, likewise
        const bits_per_sample = try input.takeInt(u16, .little);
        var left = size - 16;

        if (tag == format_extensible) {
            // Two bytes of extension size and twenty-two of extension.
            if (left < 24) return error.MalformedWave;
            _ = try input.takeInt(u16, .little); // extension size
            _ = try input.takeInt(u16, .little); // valid bits per sample
            _ = try input.takeInt(u32, .little); // channel mask
            // The first two bytes of the sub-format GUID hold the tag the
            // header would have carried if it were not extensible.
            tag = try input.takeInt(u16, .little);
            try input.discardAll(14);
            left -= 24;
        }
        try input.discardAll(left + (size & 1));

        if (tag != format_pcm) return error.UnsupportedWaveFormat;
        if (channels == 0 or channels > 255) return error.UnsupportedWaveFormat;
        switch (bits_per_sample) {
            8, 16, 24, 32 => {},
            else => return error.UnsupportedWaveFormat,
        }
        return .{
            .sample_rate = sample_rate,
            .channels = @intCast(channels),
            .bits_per_sample = @intCast(bits_per_sample),
        };
    }

    /// Read interleaved samples into `out`, which must have room for a whole
    /// number of interchannel samples. Returns how many values were written,
    /// which is short only at the end of the data.
    ///
    /// A partial interchannel sample at the end of the stream is not a sample
    /// and is dropped; a file that ends inside one is truncated, and where the
    /// data chunk stated its size that is an error rather than an ending.
    pub fn read(r: *Reader, out: []i32) ReadError!usize {
        const channels = r.format.channels;
        const width = r.format.bytesPerSample();
        const wanted: u64 = if (r.unknown_length)
            out.len / channels
        else
            @min(out.len / channels, r.remaining);

        var produced: usize = 0;
        var complete: u64 = 0;
        while (complete < wanted) : (complete += 1) {
            const started = produced;
            for (0..channels) |_| {
                const raw = r.input.take(width) catch |err| switch (err) {
                    error.EndOfStream => {
                        if (!r.unknown_length) return error.MalformedWave;
                        return started;
                    },
                    error.ReadFailed => return error.ReadFailed,
                };
                out[produced] = decodeSample(raw, width);
                produced += 1;
            }
        }
        if (!r.unknown_length) r.remaining -= complete;
        return produced;
    }
};

fn decodeSample(raw: []const u8, width: u8) i32 {
    return switch (width) {
        // Eight-bit WAVE samples are unsigned, centred on 128.
        1 => @as(i32, raw[0]) - 128,
        2 => std.mem.readInt(i16, raw[0..2], .little),
        3 => blk: {
            const value: u32 = @as(u32, raw[0]) | (@as(u32, raw[1]) << 8) | (@as(u32, raw[2]) << 16);
            break :blk @as(i32, @bitCast(value << 8)) >> 8;
        },
        4 => std.mem.readInt(i32, raw[0..4], .little),
        else => unreachable,
    };
}

/// Write a 44-byte header. `format.total_samples` may be zero if it is not yet
/// known, in which case both size fields say zero and a caller that can seek
/// puts them right with `patchSizes`.
pub fn writeHeader(w: *std.Io.Writer, format: Format) WriteError!void {
    const data_bytes = format.total_samples * format.blockAlign();
    if (data_bytes > std.math.maxInt(u32) - header_length) return error.WaveTooLarge;
    const size: u32 = @intCast(data_bytes);

    try w.writeAll("RIFF");
    try w.writeInt(u32, if (size == 0) 0 else size + header_length - 8, .little);
    try w.writeAll("WAVE");

    try w.writeAll("fmt ");
    try w.writeInt(u32, 16, .little);
    try w.writeInt(u16, format_pcm, .little);
    try w.writeInt(u16, format.channels, .little);
    try w.writeInt(u32, format.sample_rate, .little);
    try w.writeInt(u32, format.sample_rate * format.blockAlign(), .little);
    try w.writeInt(u16, format.blockAlign(), .little);
    try w.writeInt(u16, format.bits_per_sample, .little);

    try w.writeAll("data");
    try w.writeInt(u32, size, .little);
}

/// The two size fields, for a caller writing them back over a header whose
/// sizes were not known when it was written.
pub const Sizes = struct {
    riff: u32,
    data: u32,
};

pub fn sizes(format: Format) WriteError!Sizes {
    const data_bytes = format.total_samples * format.blockAlign();
    if (data_bytes > std.math.maxInt(u32) - header_length) return error.WaveTooLarge;
    return .{ .riff = @intCast(data_bytes + header_length - 8), .data = @intCast(data_bytes) };
}

/// Write interleaved samples in the file's own width.
pub fn writeSamples(w: *std.Io.Writer, format: Format, samples: []const i32) WriteError!void {
    switch (format.bytesPerSample()) {
        1 => for (samples) |sample| try w.writeByte(@intCast(sample + 128)),
        2 => for (samples) |sample| try w.writeInt(i16, @intCast(sample), .little),
        3 => for (samples) |sample| {
            const value: u32 = @bitCast(sample);
            try w.writeAll(&.{
                @truncate(value),
                @truncate(value >> 8),
                @truncate(value >> 16),
            });
        },
        4 => for (samples) |sample| try w.writeInt(i32, sample, .little),
        else => unreachable,
    }
}

const testing = std.testing;

test "a header this module wrote is one it can read" {
    const format: Format = .{
        .sample_rate = 44100,
        .channels = 2,
        .bits_per_sample = 24,
        .total_samples = 3,
    };
    const samples = [_]i32{ 0, -1, 8388607, -8388608, 1234567, -7654321 };

    var buffer: [header_length + 18]u8 = undefined;
    var w: std.Io.Writer = .fixed(&buffer);
    try writeHeader(&w, format);
    try testing.expectEqual(@as(usize, header_length), w.end);
    try writeSamples(&w, format, &samples);

    var input: std.Io.Reader = .fixed(&buffer);
    var r: Reader = try .init(&input);
    try testing.expectEqual(format, r.format);

    var out: [6]i32 = undefined;
    try testing.expectEqual(@as(usize, 6), try r.read(&out));
    try testing.expectEqualSlices(i32, &samples, &out);
    try testing.expectEqual(@as(u64, 0), r.remaining);
}

test "every width round-trips at its extremes" {
    for ([_]u8{ 8, 16, 24, 32 }) |bits_per_sample| {
        const high: i32 = @intCast((@as(i64, 1) << @intCast(bits_per_sample - 1)) - 1);
        const low: i32 = @intCast(-(@as(i64, 1) << @intCast(bits_per_sample - 1)));
        const samples = [_]i32{ 0, 1, -1, high, low, high - 1, low + 1 };
        const format: Format = .{
            .sample_rate = 8000,
            .channels = 1,
            .bits_per_sample = bits_per_sample,
            .total_samples = samples.len,
        };

        var buffer: [header_length + 32]u8 = undefined;
        var w: std.Io.Writer = .fixed(&buffer);
        try writeHeader(&w, format);
        try writeSamples(&w, format, &samples);

        var input: std.Io.Reader = .fixed(buffer[0..w.end]);
        var r: Reader = try .init(&input);
        var out: [samples.len]i32 = undefined;
        _ = try r.read(&out);
        try testing.expectEqualSlices(i32, &samples, &out);
    }
}

test "an extensible header is understood, and chunks in between are skipped" {
    // A `LIST` chunk of odd length before `fmt `, an extensible format chunk,
    // and then the data: all three things a real file does that the simple
    // path does not cover.
    var out: std.Io.Writer.Allocating = .init(testing.allocator);
    defer out.deinit();
    const w = &out.writer;

    try w.writeAll("RIFF");
    try w.writeInt(u32, 0, .little);
    try w.writeAll("WAVE");

    try w.writeAll("LIST");
    try w.writeInt(u32, 5, .little);
    try w.writeAll("INFO\x00");
    try w.writeByte(0); // the pad byte an odd chunk needs

    try w.writeAll("fmt ");
    try w.writeInt(u32, 40, .little);
    try w.writeInt(u16, format_extensible, .little);
    try w.writeInt(u16, 2, .little);
    try w.writeInt(u32, 48000, .little);
    try w.writeInt(u32, 48000 * 4, .little);
    try w.writeInt(u16, 4, .little);
    try w.writeInt(u16, 16, .little);
    try w.writeInt(u16, 22, .little);
    try w.writeInt(u16, 16, .little);
    try w.writeInt(u32, 3, .little);
    try w.writeInt(u16, format_pcm, .little);
    try w.splatByteAll(0, 14);

    try w.writeAll("data");
    try w.writeInt(u32, 8, .little);
    for ([_]i16{ 1, -1, 32767, -32768 }) |sample| try w.writeInt(i16, sample, .little);

    var input: std.Io.Reader = .fixed(out.written());
    var r: Reader = try .init(&input);
    try testing.expectEqual(@as(u32, 48000), r.format.sample_rate);
    try testing.expectEqual(@as(u8, 2), r.format.channels);
    try testing.expectEqual(@as(u8, 16), r.format.bits_per_sample);
    try testing.expectEqual(@as(u64, 2), r.format.total_samples);

    var samples: [4]i32 = undefined;
    try testing.expectEqual(@as(usize, 4), try r.read(&samples));
    try testing.expectEqualSlices(i32, &.{ 1, -1, 32767, -32768 }, &samples);
}

test "a data chunk with no usable size reads to the end of the stream" {
    // What a WAVE file written to a pipe looks like: the size fields could not
    // be filled in afterwards, so they say all ones.
    var out: std.Io.Writer.Allocating = .init(testing.allocator);
    defer out.deinit();
    const w = &out.writer;
    try writeHeader(w, .{ .sample_rate = 44100, .channels = 2, .bits_per_sample = 16 });
    // Overwrite the data size with the placeholder a streaming writer uses.
    std.mem.writeInt(u32, out.written()[data_size_offset..][0..4], std.math.maxInt(u32), .little);
    for ([_]i16{ 1, -1, 2, -2, 3, -3 }) |sample| try w.writeInt(i16, sample, .little);
    // And half an interchannel sample, which is not one.
    try w.writeInt(i16, 99, .little);

    var input: std.Io.Reader = .fixed(out.written());
    var r: Reader = try .init(&input);
    try testing.expect(r.unknown_length);
    try testing.expectEqual(@as(u64, 0), r.format.total_samples);

    var samples: [16]i32 = undefined;
    try testing.expectEqual(@as(usize, 6), try r.read(&samples));
    try testing.expectEqualSlices(i32, &.{ 1, -1, 2, -2, 3, -3 }, samples[0..6]);
    try testing.expectEqual(@as(usize, 0), try r.read(&samples));
}

test "a file that ends inside its stated data is truncated" {
    var out: std.Io.Writer.Allocating = .init(testing.allocator);
    defer out.deinit();
    try writeHeader(&out.writer, .{
        .sample_rate = 44100,
        .channels = 2,
        .bits_per_sample = 16,
        .total_samples = 4,
    });
    try out.writer.writeInt(i16, 1, .little);

    var input: std.Io.Reader = .fixed(out.written());
    var r: Reader = try .init(&input);
    var samples: [8]i32 = undefined;
    try testing.expectError(error.MalformedWave, r.read(&samples));
}

test "files that are not WAVE files are refused" {
    var not_riff: std.Io.Reader = .fixed("fLaC\x00\x00\x00\x22");
    try testing.expectError(error.NotWave, Reader.init(&not_riff));

    var not_wave: std.Io.Reader = .fixed("RIFF\x08\x00\x00\x00AVI ");
    try testing.expectError(error.NotWave, Reader.init(&not_wave));

    // A format this module does not decode: IEEE float.
    var floating: std.Io.Writer.Allocating = .init(testing.allocator);
    defer floating.deinit();
    const w = &floating.writer;
    try w.writeAll("RIFF");
    try w.writeInt(u32, 36, .little);
    try w.writeAll("WAVE");
    try w.writeAll("fmt ");
    try w.writeInt(u32, 16, .little);
    try w.writeInt(u16, 3, .little);
    try w.writeInt(u16, 1, .little);
    try w.writeInt(u32, 44100, .little);
    try w.writeInt(u32, 44100 * 4, .little);
    try w.writeInt(u16, 4, .little);
    try w.writeInt(u16, 32, .little);
    var input: std.Io.Reader = .fixed(floating.written());
    try testing.expectError(error.UnsupportedWaveFormat, Reader.init(&input));
}