#include "ffmpeg_audio.h"

#include <stdint.h>
#include <stdio.h>
#include <stdlib.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) {
    const uint32_t byte_rate =
        (uint32_t)(sample_rate * channels * bits_per_sample / 8);
    const uint16_t block_align =
        (uint16_t)(channels * bits_per_sample / 8);

    fwrite("RIFF", 1, 4, f);
    write_u32_le(f, 36u + 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);               /* 1 = integer 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);
}

static int write_decoder_buffer(FILE *out,
                                fmpg_decoder *dec,
                                uint64_t *total_written) {
    const uint8_t *buf = fmpg_decoder_buffer(dec);
    const int size = fmpg_decoder_buffer_size(dec);

    if (!buf || size <= 0) {
        return 1;
    }

    if (fwrite(buf, 1, (size_t)size, out) != (size_t)size) {
        return 0;
    }

    *total_written += (uint64_t)size;
    return 1;
}

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

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

    int sample_rate;
    int channels;
    int bits_per_sample;
    int64_t duration_ms;
    int64_t duration_samples;
    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 = fmpg_init();
    if (!ac) {
        fprintf(stderr, "ac_init failed\n");
        return 1;
    }

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

    sample_rate = fmpg_audio_sample_rate(ac);
    channels = fmpg_audio_channels(ac);
    bits_per_sample = fmpg_audio_bits_per_sample(ac);
    duration_ms = fmpg_duration_ms(ac);
    duration_samples = fmpg_duration_samples(ac);

    if (sample_rate <= 0 || channels <= 0 || bits_per_sample != 32) {
        fprintf(stderr, "invalid audio parameters\n");
        fmpg_free(ac);
        return 1;
    }

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

    out = fopen(outfile, "wb");
    if (!out) {
        fprintf(stderr, "could not open output file: %s\n", outfile);
        fmpg_free_decoder(dec);
        fmpg_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,
                          sample_rate,
                          channels,
                          bits_per_sample,
                          0)) {
        fprintf(stderr, "could not write WAV header\n");
        fclose(out);
        fmpg_free_decoder(dec);
        fmpg_free(ac);
        return 1;
    }

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

        if (!pkg) {
            break;
        }

        /*
         * ac_read_package() now returns only packets from the internally
         * selected audio stream. No stream_index test is needed anymore.
         */
        if (fmpg_decode_package(pkg, dec)) {
            if (!write_decoder_buffer(out, dec, &total_written)) {
                fprintf(stderr, "could not write PCM data\n");
                fmpg_free_package(pkg);
                fclose(out);
                fmpg_free_decoder(dec);
                fmpg_free(ac);
                return 1;
            }
        }

        fmpg_free_package(pkg);
    }

    /* Drain delayed samples from the decoder and resampler. */
    while (fmpg_flush_decoder(dec)) {
        if (!write_decoder_buffer(out, dec, &total_written)) {
            fprintf(stderr, "could not write flushed PCM data\n");
            fclose(out);
            fmpg_free_decoder(dec);
            fmpg_free(ac);
            return 1;
        }
    }

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

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

    fclose(out);

    printf("wrote %s\n", outfile);
    printf("title:       %s\n", fmpg_file_title(ac));
    printf("album:       %s\n", fmpg_file_album(ac));
    printf("sample rate: %d\n", sample_rate);
    printf("channels:    %d\n", channels);
    printf("sample bits: %d\n", bits_per_sample);
    printf("duration ms: %lld\n", (long long)duration_ms);
    printf("duration smp:%lld\n", (long long)duration_samples);
    printf("decoded smp: %lld\n", (long long)fmpg_decoder_sample_position(dec));
    printf("data bytes:  %llu\n",
           (unsigned long long)total_written);

    fmpg_free_decoder(dec);
    fmpg_free(ac);

    return 0;
}