#include <aacnext/libaac.h>
#include <cstdio>
#include "argparse.h"
#include "audiofile.h"
#include "indicator.h"
#define MINIMP4_IMPLEMENTATION
#include "minimp4.h"

using namespace indicators;

static int doWrite(int64_t offset, const void *buffer, size_t size, void *token)
{
    FILE *f = (FILE*)token;
    fseek(f, offset, SEEK_SET);
    return fwrite(buffer, 1, size, f) != size;
}

int main(int argc, char *argv[]) {
    argparse::ArgumentParser ap("aacenc");
    AudioFile<float> audioFile;
    AACSettings aac = {};

    ap.add_argument("input");
    ap.add_argument("output");

    ap.add_argument("-a", "--adts").help("ADTS format").flag();
    ap.add_argument("-b", "--bitrate").help("Bit rate").default_value(128).scan<'i', int>();
    ap.add_argument("-c", "--cutoff").help("Cutoff frequency").default_value(0).scan<'i', int>();
    ap.add_argument("-e", "--enhanced-sbr").help("Enhanced SBR").flag();
    ap.add_argument("-f", "--frame-size").help("Frame size").default_value(0).scan<'i', int>();
    ap.add_argument("-i", "--iq").help("Inverse quantization").default_value(2).scan<'i', int>();
    ap.add_argument("-p", "--profile").help("AAC profile").default_value(0).scan<'i', int>();
    ap.add_argument("-t").help("TNS").flag();

    try {
        ap.parse_args(argc, argv);
    } catch (const std::exception& err) {
        std::cerr << err.what() << std::endl;
        std::cerr << ap;
        return 1;
    }

    auto adts = ap.get<bool>("-a");
    auto inFile = ap.get<std::string>("input");
    auto outFile = ap.get<std::string>("output");
    auto bitRate = ap.get<int>("-b");
    auto cutOff = ap.get<int>("-c");
    auto eSBR = ap.get<bool>("-e");
    auto frameSize = ap.get<int>("-f");
    auto inverseQuant = ap.get<int>("-i");
    auto profile = ap.get<int>("-p");
    auto tns = ap.get<bool>("-t");

    if (!audioFile.load(inFile)) {
        return 1;
    }

    FILE * outData = fopen(outFile.c_str(), "wb");
    if (!outData) {
        std::cerr << "Unable to open " << outFile << ": " << std::strerror(errno) << std::endl;
        return 1;
    }

    std::cout << "Input file: " << inFile << std::endl;
    std::cout << "Output file: " << outFile << std::endl;
    audioFile.printSummary();
    std::cout << "AAC encode configuration:" << std::endl;
    std::cout << "Bitrate: " << bitRate << "kbps" << std::endl;
    std::cout << "Cutoff: " << cutOff << std::endl;
    std::cout << "Enhanced SBR: " << (eSBR ? "true" : "false") << std::endl;
    std::cout << "Frame Size: " << frameSize << std::endl;
    std::cout << "Inverse quantization: " << inverseQuant << std::endl;
    std::cout << "Profile: " << profile << std::endl;
    std::cout << "TNS: " << (tns ? "true" : "false") << std::endl;
    std::cout << "ADTS: " << (adts ? "true" : "false") << std::endl;

    aac.adts = adts;
    aac.bitrate = bitRate * 1000;
    aac.bitsPerSamples = 32;
    aac.cutoff = cutOff;
    aac.frameSize = frameSize;
    aac.iq = inverseQuant;
    aac.noChannels = audioFile.getNumChannels();
    aac.profile = (profiles)profile;
    aac.sampleRate = audioFile.getSampleRate();
    aac.eSBR = eSBR;
    aac.tns = tns;

    auto aacEncode = aac_encode_open(aac);
    if (aacEncode == NULL) {
        std::cerr << "Failed opening AAC encoder" << std::endl;
        return 1;
    }

    MP4E_mux_t *mux;
    int audio_track_id;
    uint32_t tick_duration = ((aacEncode->no_samples/audioFile.getNumChannels()) * 90000)/audioFile.getSampleRate();
    std::cout << "Block size: " << (aacEncode->no_samples/audioFile.getNumChannels())  << std::endl;
    std::cout << "Delay: " << (aacEncode->inputDelay) << std::endl;

    if (!adts) {
        mux = MP4E_open(0, 0, outData, doWrite);
        MP4E_set_text_comment(mux, "generated by libaac-next");

        MP4E_track_t tr;
        tr.track_media_kind = e_audio;
        tr.language[0] = 'u';
        tr.language[1] = 'n';
        tr.language[2] = 'd';
        tr.language[3] = 0;
        tr.object_type_indication = MP4_OBJECT_TYPE_AUDIO_ISO_IEC_14496_3;
        tr.time_scale = 90000;
        tr.default_duration = 0;
        tr.u.a.channelcount = audioFile.getNumChannels();
        audio_track_id = MP4E_add_track(mux, &tr);
        MP4E_set_dsi(mux, audio_track_id, aacEncode->asc, aacEncode->ascSize);

        printf("Extradata: ");
        for (uint32_t i = 0; i < aacEncode->ascSize; i++) {
            printf("%02x ", aacEncode->asc[i]);
        }
        printf("\n");
    }

    size_t numSamples = audioFile.getNumSamplesPerChannel();
    size_t sampleReadOffset = 0;

    float *samples = new float[aacEncode->no_samples];
    uint8_t *aacBuf = new uint8_t[aacEncode->max_out_bytes];
    uint32_t aacWriteSize;

    show_console_cursor(false);

    ProgressBar bar{
        option::BarWidth{50},
        option::Start{"["},
        option::Fill{"="},
        option::Lead{">"},
        option::Remainder{" "},
        option::End{"]"},
        option::PostfixText{"Processing samples"},
        option::ShowPercentage{true},
        option::FontStyles{std::vector<FontStyle>{FontStyle::bold}}
    };

    for (size_t sampleOffset = 0; sampleOffset < numSamples; sampleOffset++) {
        for (unsigned int ch = 0; ch < aac.noChannels; ch++) {
            samples[sampleReadOffset++] = audioFile.samples[ch][sampleOffset];
        }
        if (sampleReadOffset >= aacEncode->no_samples) {
            bar.set_progress(((double)sampleOffset / (double)numSamples) * 100);
            auto ret = aac_encode(aacEncode, (unsigned char *)samples, sampleReadOffset * 4, aacBuf, &aacWriteSize);
            if (ret == -1) {
                show_console_cursor(true);
                std::cerr << "Failed encoding to AAC encoder" << std::endl;
                return 1;
            }
            if (adts) {
                fwrite(aacBuf, 1, aacWriteSize, outData);
            } else {
                if (MP4E_STATUS_OK != MP4E_put_sample(mux, audio_track_id, aacBuf, aacWriteSize, tick_duration, MP4E_SAMPLE_RANDOM_ACCESS)) {
                    show_console_cursor(true);
                    std::cerr << "Failed muxing to MP4 muxer: " << ret << std::endl;
                    return 1;
                }
            }
            sampleReadOffset = 0;
        }
    }

    bar.set_progress(100);
    show_console_cursor(true);

    if (sampleReadOffset > 0) {
        auto ret = aac_encode(aacEncode, (unsigned char *)samples, sampleReadOffset * 4, aacBuf, &aacWriteSize);
        if (ret == -1) {
            std::cerr << "Failed encoding to AAC encoder (last)" << std::endl;
            return 1;
        } else if (ret == 0) {
            std::cerr << "Flush" << std::endl;
            auto ret = aac_encode(aacEncode, NULL, 0, aacBuf, &aacWriteSize);
            if (ret == -1) {
                std::cerr << "Failed encoding to AAC encoder (last flush)" << std::endl;
                return 1;
            }
        }
        if (adts) {
            fwrite(aacBuf, 1, aacWriteSize, outData);
        } else {
            if (MP4E_STATUS_OK != MP4E_put_sample(mux, audio_track_id, aacBuf, aacWriteSize, tick_duration, MP4E_SAMPLE_RANDOM_ACCESS)) {
                std::cerr << "Failed muxing to MP4 muxer" << std::endl;
                return 1;
            }
        }
    }

    for (int i = 0; i < 2; i++) {
        auto ret = aac_encode(aacEncode, NULL, 0, aacBuf, &aacWriteSize);
        if (ret == -1) {
            std::cerr << "Failed encoding to AAC encoder (padding)" << std::endl;
            return 1;
        }
        if (adts) {
            fwrite(aacBuf, 1, aacWriteSize, outData);
        } else {
            if (MP4E_STATUS_OK != MP4E_put_sample(mux, audio_track_id, aacBuf, aacWriteSize, tick_duration, MP4E_SAMPLE_RANDOM_ACCESS)) {
                std::cerr << "Failed muxing to MP4 muxer" << std::endl;
                return 1;
            }
        }
    }

    
    if (!adts) MP4E_close(mux);
    fclose(outData);
    aac_encode_close(aacEncode);
}