/*
 * Beat Rat -- WAV reader/writer
 *
 * Deliberately hand-written rather than pulling in libsndfile: WAV is a
 * simple enough format that we'd rather own this code outright. Only
 * handles 16-bit PCM, which is all Beat Rat uses internally right now.
 *
 * Byte order note: multi-byte fields are written/read explicitly as
 * little-endian regardless of host order, but raw sample data is
 * fread/fwritten directly as int16_t, which assumes a little-endian
 * host. That's true for every realistic Beat Rat target (x86, x86_64,
 * and default-endian arm64), so this is a deliberate simplification,
 * not an oversight.
 */

#include "wav.h"
#include "log.h"

#include <string.h>

static void write_u32le(FILE *fp, uint32_t v) {
    unsigned char b[4] = { v & 0xFF, (v >> 8) & 0xFF, (v >> 16) & 0xFF, (v >> 24) & 0xFF };
    fwrite(b, 1, 4, fp);
}

static void write_u16le(FILE *fp, uint16_t v) {
    unsigned char b[2] = { v & 0xFF, (v >> 8) & 0xFF };
    fwrite(b, 1, 2, fp);
}

static int read_u32le(FILE *fp, uint32_t *out) {
    unsigned char b[4];
    if (fread(b, 1, 4, fp) != 4) return -1;
    *out = (uint32_t)b[0] | ((uint32_t)b[1] << 8) | ((uint32_t)b[2] << 16) | ((uint32_t)b[3] << 24);
    return 0;
}

static int read_u16le(FILE *fp, uint16_t *out) {
    unsigned char b[2];
    if (fread(b, 1, 2, fp) != 2) return -1;
    *out = (uint16_t)b[0] | ((uint16_t)b[1] << 8);
    return 0;
}

int wav_writer_open(WavWriter *w, const char *path, uint32_t sample_rate,
                     uint16_t channels, uint16_t bits_per_sample) {
    w->fp = fopen(path, "wb");
    if (!w->fp) {
        LOG_ERROR("could not open '%s' for writing", path);
        return -1;
    }
    w->sample_rate = sample_rate;
    w->channels = channels;
    w->bits_per_sample = bits_per_sample;
    w->data_bytes_written = 0;

    fwrite("RIFF", 1, 4, w->fp);
    w->riff_size_field_pos = ftell(w->fp);
    write_u32le(w->fp, 0); /* patched in wav_writer_close */
    fwrite("WAVE", 1, 4, w->fp);

    fwrite("fmt ", 1, 4, w->fp);
    write_u32le(w->fp, 16); /* PCM fmt chunk is always 16 bytes */
    write_u16le(w->fp, 1);  /* AudioFormat = 1 (PCM) */
    write_u16le(w->fp, channels);
    write_u32le(w->fp, sample_rate);
    uint32_t byte_rate = sample_rate * channels * (bits_per_sample / 8);
    write_u32le(w->fp, byte_rate);
    uint16_t block_align = (uint16_t)(channels * (bits_per_sample / 8));
    write_u16le(w->fp, block_align);
    write_u16le(w->fp, bits_per_sample);

    fwrite("data", 1, 4, w->fp);
    w->data_size_field_pos = ftell(w->fp);
    write_u32le(w->fp, 0); /* patched in wav_writer_close */

    LOG_INFO("opened '%s' for writing (%uHz, %u ch, %u-bit)", path, sample_rate, channels, bits_per_sample);
    return 0;
}

int wav_writer_write(WavWriter *w, const int16_t *samples, size_t count) {
    size_t written = fwrite(samples, sizeof(int16_t), count, w->fp);
    w->data_bytes_written += (uint32_t)(written * sizeof(int16_t));
    if (written != count) {
        LOG_ERROR("short write to WAV file (%zu of %zu samples)", written, count);
        return -1;
    }
    return 0;
}

int wav_writer_close(WavWriter *w) {
    if (!w->fp) return -1;

    uint32_t data_bytes = w->data_bytes_written;
    uint32_t riff_bytes = 36 + data_bytes;

    fseek(w->fp, w->riff_size_field_pos, SEEK_SET);
    write_u32le(w->fp, riff_bytes);
    fseek(w->fp, w->data_size_field_pos, SEEK_SET);
    write_u32le(w->fp, data_bytes);

    fclose(w->fp);
    w->fp = NULL;
    LOG_INFO("closed WAV writer, %u data bytes total", data_bytes);
    return 0;
}

int wav_reader_open(WavReader *r, const char *path) {
    r->fp = fopen(path, "rb");
    if (!r->fp) {
        LOG_ERROR("could not open '%s' for reading", path);
        return -1;
    }

    char magic[4];
    if (fread(magic, 1, 4, r->fp) != 4 || memcmp(magic, "RIFF", 4) != 0) {
        LOG_ERROR("'%s' is not a RIFF file", path);
        fclose(r->fp);
        return -1;
    }
    uint32_t riff_size;
    read_u32le(r->fp, &riff_size); /* not needed, just consume it */

    if (fread(magic, 1, 4, r->fp) != 4 || memcmp(magic, "WAVE", 4) != 0) {
        LOG_ERROR("'%s' is not a WAVE file", path);
        fclose(r->fp);
        return -1;
    }

    int found_fmt = 0, found_data = 0;

    while (!found_data) {
        char chunk_id[4];
        uint32_t chunk_size;
        if (fread(chunk_id, 1, 4, r->fp) != 4) break;
        if (read_u32le(r->fp, &chunk_size) != 0) break;

        if (memcmp(chunk_id, "fmt ", 4) == 0) {
            uint16_t audio_format = 0, channels = 0, block_align = 0, bits_per_sample = 0;
            uint32_t sample_rate = 0, byte_rate = 0;
            int ok = 1;
            ok = ok && (read_u16le(r->fp, &audio_format) == 0);
            ok = ok && (read_u16le(r->fp, &channels) == 0);
            ok = ok && (read_u32le(r->fp, &sample_rate) == 0);
            ok = ok && (read_u32le(r->fp, &byte_rate) == 0);
            ok = ok && (read_u16le(r->fp, &block_align) == 0);
            ok = ok && (read_u16le(r->fp, &bits_per_sample) == 0);
            (void)byte_rate; (void)block_align;

            if (!ok) {
                LOG_ERROR("'%s' has a truncated fmt chunk", path);
                fclose(r->fp);
                return -1;
            }
            if (audio_format != 1 || bits_per_sample != 16) {
                LOG_ERROR("'%s' is not 16-bit PCM (format=%u, bits=%u) -- unsupported",
                           path, audio_format, bits_per_sample);
                fclose(r->fp);
                return -1;
            }
            r->sample_rate = sample_rate;
            r->channels = channels;
            r->bits_per_sample = bits_per_sample;

            /* Skip any extra bytes beyond the 16 standard fmt fields
             * (e.g. WAVE_FORMAT_EXTENSIBLE headers). */
            long consumed = 16;
            if ((long)chunk_size > consumed)
                fseek(r->fp, chunk_size - consumed, SEEK_CUR);
            found_fmt = 1;
        } else if (memcmp(chunk_id, "data", 4) == 0) {
            r->data_bytes_total = chunk_size;
            r->data_bytes_read = 0;
            found_data = 1; /* file is now positioned at the start of sample data */
        } else {
            /* Unknown chunk (LIST, fact, etc.) -- skip it. */
            fseek(r->fp, chunk_size, SEEK_CUR);
        }

        /* RIFF chunks are word-aligned; skip the pad byte if size is odd. */
        if (chunk_size % 2 != 0)
            fseek(r->fp, 1, SEEK_CUR);
    }

    if (!found_fmt || !found_data) {
        LOG_ERROR("'%s' is missing a fmt or data chunk", path);
        fclose(r->fp);
        return -1;
    }

    LOG_INFO("opened '%s' for reading (%uHz, %u ch, %u data bytes)",
              path, r->sample_rate, r->channels, r->data_bytes_total);
    return 0;
}

size_t wav_reader_read(WavReader *r, int16_t *samples, size_t count) {
    uint32_t remaining_bytes = r->data_bytes_total - r->data_bytes_read;
    size_t remaining_samples = remaining_bytes / sizeof(int16_t);
    size_t n = count < remaining_samples ? count : remaining_samples;
    if (n == 0) return 0;

    size_t actual = fread(samples, sizeof(int16_t), n, r->fp);
    r->data_bytes_read += (uint32_t)(actual * sizeof(int16_t));
    return actual;
}

int wav_reader_close(WavReader *r) {
    if (!r->fp) return -1;
    fclose(r->fp);
    r->fp = NULL;
    return 0;
}