// Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
// See LICENSE.md for license information.
// SPDX-License-Identifier: MIT
//
// gateWav.c - WAV files for the gate tools (gateWav.h).
//
// The reading walks a WAV as Python's wave module did, since the graders
// were written against it and a recording that is not quite whole must
// yield the same audio: the RIFF chunk's body is walked chunk by chunk,
// every read capped at the RIFF's declared size and at the file's end, a
// chunk of odd size followed by a pad byte; the fmt chunk before the first
// data chunk gives the format, and the audio is that data chunk's whole
// frames, as far as the caps let them go.

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "gateText.h"
#include "gateWav.h"
#include "toolFile.h"

#define RIFF_HEADER       8u
#define CHUNK_ID_BYTES    4u
#define CHUNK_SIZE_BYTES  4u
#define FMT_BYTES         14u
#define FMT_BITS_BYTES    2u
#define FMT_EXT_BYTES     8u
#define FMT_GUID_BYTES    16u
#define FMT_CHANNELS_AT   2u
#define FMT_RATE_AT       4u
#define FORMAT_PCM        0x0001u
#define FORMAT_EXTENSIBLE 0xFFFEu
#define BITS_PER_BYTE     8u
#define SAMPLE_BYTES      2u
#define WAV_FMT_SIZE      16u
#define WAV_RIFF_EXTRA    36u

// The RIFF chunk being read: its declared size, and how much of it has
// been read (its body starts RIFF_HEADER bytes into the file).
typedef struct RiffT {
    const uint8_t *file;
    size_t         fileLen;
    uint64_t       size;
    uint64_t       pos;
} RiffT;

// A chunk inside it: its declared size, where its body starts in the RIFF,
// and how much of it has been read.
typedef struct ChunkT {
    char     id[CHUNK_ID_BYTES];
    uint64_t size;
    uint64_t offset;
    uint64_t pos;
} ChunkT;

static size_t   chunkRead(RiffT *riff, ChunkT *chunk, uint64_t n, const uint8_t **at);
static bool     chunkSkip(RiffT *riff, ChunkT *chunk);
static bool     fmtRead(RiffT *riff, ChunkT *chunk, const char *path, GateWavT *wav, ToolBufT *problem);
static uint32_t le16(const uint8_t *at);
static uint32_t le32(const uint8_t *at);
static size_t   riffRead(RiffT *riff, uint64_t n, const uint8_t **at);

// KSDATAFORMAT_SUBTYPE_PCM, the GUID of an extensible WAV's PCM samples.
static const uint8_t kPcmGuid[FMT_GUID_BYTES] = {
    0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x80, 0x00, 0x00, 0xAA, 0x00, 0x38, 0x9B, 0x71
};


// Up to n bytes of a chunk's body (then its pad byte, when its size is odd
// and its end is reached): how many, at *at.
static size_t chunkRead(RiffT *riff, ChunkT *chunk, uint64_t n, const uint8_t **at) {
    const uint8_t *pad;
    size_t         got;

    *at = NULL;
    if (chunk->pos >= chunk->size) {
        return 0u;
    }
    if (n > chunk->size - chunk->pos) {
        n = chunk->size - chunk->pos;
    }
    got         = riffRead(riff, n, at);
    chunk->pos += got;
    if (chunk->pos == chunk->size && (chunk->size & 1u) != 0u) {
        chunk->pos += riffRead(riff, 1u, &pad);
    }
    return got;
}


// Past the rest of a chunk (and its pad byte); false when that would leave
// the RIFF chunk.
static bool chunkSkip(RiffT *riff, ChunkT *chunk) {
    int64_t step;
    int64_t to;

    step = (int64_t)chunk->size - (int64_t)chunk->pos;
    if ((chunk->size & 1u) != 0u) {
        step++;
    }
    to = (int64_t)riff->pos + step;
    if (to < 0 || (uint64_t)to > riff->size) {
        return false;
    }
    riff->pos   = (uint64_t)to;
    chunk->pos += (uint64_t)step;
    return true;
}


// The fmt chunk: 16-bit PCM (plainly, or as the extensible format's PCM),
// its channels and rate; false, with what is wrong in problem, when it is
// not one this reads.
static bool fmtRead(RiffT *riff, ChunkT *chunk, const char *path, GateWavT *wav, ToolBufT *problem) {
    const uint8_t *at;
    uint32_t       format;
    uint32_t       bits;

    if (chunkRead(riff, chunk, FMT_BYTES, &at) < FMT_BYTES) {
        gateTextAdd(problem, "%s: fmt chunk cut short", path);
        return false;
    }
    format        = le16(at);
    wav->channels = le16(at + FMT_CHANNELS_AT);
    wav->rate     = le32(at + FMT_RATE_AT);
    if (format != FORMAT_PCM && format != FORMAT_EXTENSIBLE) {
        gateTextAdd(problem, "%s: format %u, not PCM", path, (unsigned)format);
        return false;
    }
    if (chunkRead(riff, chunk, FMT_BITS_BYTES, &at) < FMT_BITS_BYTES) {
        gateTextAdd(problem, "%s: fmt chunk cut short", path);
        return false;
    }
    bits = le16(at);
    if (format == FORMAT_EXTENSIBLE) {
        if (chunkRead(riff, chunk, FMT_EXT_BYTES, &at) < FMT_EXT_BYTES || chunkRead(riff, chunk, FMT_GUID_BYTES, &at) < FMT_GUID_BYTES) {
            gateTextAdd(problem, "%s: fmt chunk cut short", path);
            return false;
        }
        if (memcmp(at, kPcmGuid, FMT_GUID_BYTES) != 0) {
            gateTextAdd(problem, "%s: extensible format, not PCM", path);
            return false;
        }
    }
    // A sample takes whole bytes: 9 to 16 bits are stored in 16.
    if ((bits + BITS_PER_BYTE - 1u) / BITS_PER_BYTE != SAMPLE_BYTES) {
        gateTextAdd(problem, "%s: %u-bit samples, not 16-bit", path, (unsigned)bits);
        return false;
    }
    if (wav->channels == 0u) {
        gateTextAdd(problem, "%s: no channels", path);
        return false;
    }
    return true;
}


size_t gateWavFrames(const GateWavT *wav) {
    return wav->byteCount / (SAMPLE_BYTES * wav->channels);
}


void gateWavFree(GateWavT *wav) {
    free(wav->bytes);
    memset(wav, 0, sizeof(*wav));
}


double *gateWavMono(const GateWavT *wav) {
    double  *mono;
    size_t   frames;
    size_t   frame;
    uint32_t channel;
    double   total;

    frames = gateWavFrames(wav);
    mono   = (double *)malloc((frames > 0u ? frames : 1u) * sizeof(double));
    if (mono == NULL) {
        toolFail("out of memory");
    }
    for (frame = 0u; frame < frames; frame++) {
        total = 0.0;
        for (channel = 0u; channel < wav->channels; channel++) {
            total += gateWavSample(wav, frame, channel);
        }
        mono[frame] = total / (double)wav->channels;
    }
    return mono;
}


bool gateWavRead(const char *path, GateWavT *wav, ToolBufT *problem) {
    uint8_t       *file;
    size_t         fileLen;
    const uint8_t *at;
    const char    *wrong;
    RiffT          riff;
    ChunkT         chunk;
    ChunkT         data;
    bool           haveFmt;
    bool           haveData;
    uint64_t       frames;

    memset(wav, 0, sizeof(*wav));
    memset(&data, 0, sizeof(data));
    if (!gateTextReadFile(path, &file, &fileLen, problem)) {
        return false;
    }
    riff.file    = file;
    riff.fileLen = fileLen;
    riff.size    = fileLen >= RIFF_HEADER ? le32(file + CHUNK_ID_BYTES) : 0u;
    riff.pos     = 0u;
    wrong        = NULL;
    if (fileLen < RIFF_HEADER || memcmp(file, "RIFF", CHUNK_ID_BYTES) != 0) {
        wrong = "not a RIFF file";
    } else if (riffRead(&riff, CHUNK_ID_BYTES, &at) < CHUNK_ID_BYTES || memcmp(at, "WAVE", CHUNK_ID_BYTES) != 0) {
        wrong = "not a WAVE file";
    }
    haveFmt  = false;
    haveData = false;
    while (wrong == NULL && !haveData) {
        if (riffRead(&riff, CHUNK_ID_BYTES, &at) < CHUNK_ID_BYTES) {
            break;
        }
        memcpy(chunk.id, at, CHUNK_ID_BYTES);
        if (riffRead(&riff, CHUNK_SIZE_BYTES, &at) < CHUNK_SIZE_BYTES) {
            break;
        }
        chunk.size   = le32(at);
        chunk.offset = riff.pos;
        chunk.pos    = 0u;
        if (memcmp(chunk.id, "fmt ", CHUNK_ID_BYTES) == 0) {
            if (!fmtRead(&riff, &chunk, path, wav, problem)) {
                free(file);
                return false;
            }
            haveFmt = true;
        } else if (memcmp(chunk.id, "data", CHUNK_ID_BYTES) == 0) {
            data     = chunk;
            haveData = true;
            continue;
        }
        if (!chunkSkip(&riff, &chunk)) {
            wrong = "a chunk runs past the end of the RIFF";
        }
    }
    if (wrong == NULL && !haveFmt) {
        wrong = "no fmt chunk before its data";
    } else if (wrong == NULL && !haveData) {
        wrong = "no data chunk";
    } else if (wrong == NULL && data.offset > riff.size) {
        wrong = "a chunk runs past the end of the RIFF";
    }
    if (wrong != NULL) {
        gateTextAdd(problem, "%s: %s", path, wrong);
        free(file);
        return false;
    }
    // The data's whole frames, as far as the RIFF and the file go.
    riff.pos       = data.offset;
    data.pos       = 0u;
    frames         = data.size / ((uint64_t)wav->channels * SAMPLE_BYTES);
    wav->byteCount = chunkRead(&riff, &data, frames * wav->channels * SAMPLE_BYTES, &at);
    wav->bytes     = (uint8_t *)malloc(wav->byteCount > 0u ? wav->byteCount : 1u);
    if (wav->bytes == NULL) {
        toolFail("out of memory");
    }
    if (wav->byteCount > 0u) {
        memcpy(wav->bytes, at, wav->byteCount);
    }
    free(file);
    return true;
}


double gateWavSample(const GateWavT *wav, size_t frame, uint32_t channel) {
    const uint8_t *at;

    at = wav->bytes + (frame * wav->channels + channel) * SAMPLE_BYTES;
    return (double)(int16_t)le16(at);
}


bool gateWavWrite(const char *path, uint32_t channels, uint32_t rate, const uint8_t *data, size_t len) {
    ToolBufT out;
    bool     ok;

    memset(&out, 0, sizeof(out));
    toolBufAdd(&out, "RIFF", CHUNK_ID_BYTES);
    toolBufLe32(&out, (uint32_t)(WAV_RIFF_EXTRA + len));
    toolBufAdd(&out, "WAVEfmt ", 2u * CHUNK_ID_BYTES);
    toolBufLe32(&out, WAV_FMT_SIZE);
    toolBufLe16(&out, FORMAT_PCM);
    toolBufLe16(&out, channels);
    toolBufLe32(&out, rate);
    toolBufLe32(&out, channels * rate * SAMPLE_BYTES);
    toolBufLe16(&out, channels * SAMPLE_BYTES);
    toolBufLe16(&out, SAMPLE_BYTES * BITS_PER_BYTE);
    toolBufAdd(&out, "data", CHUNK_ID_BYTES);
    toolBufLe32(&out, (uint32_t)len);
    toolBufAdd(&out, data, len);
    ok = toolWriteFile(path, out.data, out.len);
    toolBufFree(&out);
    return ok;
}


static uint32_t le16(const uint8_t *at) {
    return (uint32_t)at[0] | ((uint32_t)at[1] << BITS_PER_BYTE);
}


static uint32_t le32(const uint8_t *at) {
    return le16(at) | (le16(at + SAMPLE_BYTES) << (2u * BITS_PER_BYTE));
}


// Up to n bytes of the RIFF chunk's body, capped at its size and the
// file's end (then its pad byte, when its size is odd and its end is
// reached): how many, at *at.
static size_t riffRead(RiffT *riff, uint64_t n, const uint8_t **at) {
    uint64_t offset;
    uint64_t avail;

    *at = NULL;
    if (riff->pos >= riff->size) {
        return 0u;
    }
    if (n > riff->size - riff->pos) {
        n = riff->size - riff->pos;
    }
    offset = RIFF_HEADER + riff->pos;
    avail  = offset < riff->fileLen ? riff->fileLen - offset : 0u;
    if (n > avail) {
        n = avail;
    }
    *at        = riff->file + offset;
    riff->pos += n;
    if (riff->pos == riff->size && (riff->size & 1u) != 0u && offset + n < riff->fileLen) {
        riff->pos++;
    }
    return (size_t)n;
}