#include "ffmpeg_audio.h"

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

#ifdef WIN32
#define fprintf fprintf_s
#endif

static void write_u16_le(FILE *f, uint16_t v) {
    fputc((int)(v & 0xff), f);
    fputc((int)((v >> 8) & 0xff), f);
}

static void write_u32_le(FILE *f, uint32_t v) {
    fputc((int)(v & 0xff), f);
    fputc((int)((v >> 8) & 0xff), f);
    fputc((int)((v >> 16) & 0xff), f);
    fputc((int)((v >> 24) & 0xff), f);
}

static int write_wav_header(FILE *f,
                            int sample_rate,
                            int channels,
                            int bits_per_sample,
                            uint32_t data_size) {
    uint32_t byte_rate =
        (uint32_t)(sample_rate * channels * bits_per_sample / 8);
    uint16_t block_align =
        (uint16_t)(channels * bits_per_sample / 8);

    fwrite("RIFF", 1, 4, f);
    write_u32_le(f, 36 + data_size);
    fwrite("WAVE", 1, 4, f);

    fwrite("fmt ", 1, 4, f);
    write_u32_le(f, 16);              /* fmt chunk size */
    write_u16_le(f, 1);               /* PCM */
    write_u16_le(f, (uint16_t)channels);
    write_u32_le(f, (uint32_t)sample_rate);
    write_u32_le(f, byte_rate);
    write_u16_le(f, block_align);
    write_u16_le(f, (uint16_t)bits_per_sample);

    fwrite("data", 1, 4, f);
    write_u32_le(f, data_size);

    return ferror(f) ? 0 : 1;
}

static int rewrite_wav_header(FILE *f,
                              int sample_rate,
                              int channels,
                              int bits_per_sample,
                              uint32_t data_size) {
    if (fseek(f, 0, SEEK_SET) != 0) {
        return 0;
    }

    return write_wav_header(f,
                            sample_rate,
                            channels,
                            bits_per_sample,
                            data_size);
}

int main(int argc, char **argv) {
    const char *infile;
    const char *outfile;

    fmpg_instance * ac = NULL;
    fmpg_decoder * dec = NULL;
    FILE *out = NULL;

    int stream;
    fmpg_audio_info info;
    uint64_t total_written = 0;

    if (argc != 3) {
        fprintf(stderr, "usage: %s input-audio output.wav\n", argv[0]);
        return 1;
    }

    infile = argv[1];
    outfile = argv[2];

    ac = ac_init();
    if (!ac) {
        fprintf(stderr, "ac_init failed\n");
        return 1;
    }

    if (!ac_open_file(ac, infile)) {
        fprintf(stderr, "could not open input file: %s\n", infile);
        ac_free(ac);
        return 1;
    }

    stream = ac_get_default_audio_stream(ac);
    if (stream < 0) {
        fprintf(stderr, "no audio stream found\n");
        ac_free(ac);
        return 1;
    }

    memset(&info, 0, sizeof(info));
    if (!ac_get_audio_info(ac, stream, &info)) {
        fprintf(stderr, "could not get audio info\n");
        ac_free(ac);
        return 1;
    }

    dec = ac_create_decoder(ac, stream);
    if (!dec) {
        fprintf(stderr, "could not create decoder\n");
        ac_free(ac);
        return 1;
    }

    out = fopen(outfile, "wb");
    if (!out) {
        fprintf(stderr, "could not open output file: %s\n", outfile);
        ac_free_decoder(dec);
        ac_free(ac);
        return 1;
    }

    /*
     * We do not know the final WAV data size yet.
     * Write a placeholder header first and patch it at the end.
     */
    if (!write_wav_header(out,
                          info.sample_rate,
                          info.channels,
                          32,
                          0)) {
        fprintf(stderr, "could not write WAV header\n");
        fclose(out);
        ac_free_decoder(dec);
        ac_free(ac);
        return 1;
    }

    for (;;) {
        fmpg_package * pkg = ac_read_package(ac);

        if (!pkg) {
            break;
        }

        if (ac_package_stream_index(pkg) == stream) {
            if (ac_decode_package(pkg, dec)) {
                const uint8_t *buf = ac_decoder_buffer(dec);
                int size = ac_decoder_buffer_size(dec);

                if (buf && size > 0) {
                    fwrite(buf, 1, (size_t)size, out);
                    total_written += (uint64_t)size;
                }
            }
        }

        ac_free_package(pkg);
    }

    /*
     * Drain delayed samples from the decoder and resampler.
     */
    while (ac_flush_decoder(dec)) {
        const uint8_t *buf = ac_decoder_buffer(dec);
        int size = ac_decoder_buffer_size(dec);

        if (buf && size > 0) {
            fwrite(buf, 1, (size_t)size, out);
            total_written += (uint64_t)size;
        }
    }

    if (total_written > 0xffffffffu) {
        fprintf(stderr,
                "warning: WAV data larger than 4GB; classic WAV header "
                "will overflow\n");
    }

    if (!rewrite_wav_header(out,
                            info.sample_rate,
                            info.channels,
                            32,
                            (uint32_t)total_written)) {
        fprintf(stderr, "could not rewrite WAV header\n");
    }

    fclose(out);
    ac_free_decoder(dec);
    ac_free(ac);

    printf("wrote %s\n", outfile);
    printf("sample rate: %d\n", info.sample_rate);
    printf("channels:    %d\n", info.channels);
    printf("sample bits: %d\n", 32);
    printf("data bytes:  %llu\n",
           (unsigned long long)total_written);

    return 0;
}