#include "bass.h"        // BASS core API
#include "bassmidi.h"    // BASSMIDI extension API
#include <iostream>
#include <fstream>
#include <cstring>
#include <iomanip>
#include <chrono>

void WriteWAVHeader(std::ofstream &file, DWORD length) {
    DWORD dataChunkSize = length;
    DWORD fmtChunkSize = 16;
    DWORD sampleRate = 44100;
    WORD bitsPerSample = 16;
    WORD channels = 2;
    DWORD byteRate = sampleRate * channels * bitsPerSample / 8;
    WORD blockAlign = channels * bitsPerSample / 8;
    DWORD totalSize = 4 + (8 + fmtChunkSize) + (8 + dataChunkSize);

    file.write("RIFF", 4);
    file.write(reinterpret_cast<const char*>(&totalSize), 4);
    file.write("WAVE", 4);

    // fmt chunk
    file.write("fmt ", 4);
    file.write(reinterpret_cast<const char*>(&fmtChunkSize), 4);
    WORD formatTag = 1; // PCM
    file.write(reinterpret_cast<const char*>(&formatTag), 2);
    file.write(reinterpret_cast<const char*>(&channels), 2);
    file.write(reinterpret_cast<const char*>(&sampleRate), 4);
    file.write(reinterpret_cast<const char*>(&byteRate), 4);
    file.write(reinterpret_cast<const char*>(&blockAlign), 2);
    file.write(reinterpret_cast<const char*>(&bitsPerSample), 2);

    // data chunk
    file.write("data", 4);
    file.write(reinterpret_cast<const char*>(&dataChunkSize), 4);
}

int main(int argc, char** argv) {
    std::cout << "Program start\n";
    if (argc != 4) {
        std::cerr << "Wrong arguments!\n";
        return -1;
    }
    bool synthToFile = std::strcmp(argv[3], "true") == 0;
    
    // 1) Initialize BASS: default device, 44100 Hz, no flags
    if (!BASS_Init(-1, 44100, 0, NULL, NULL)) {
        std::cerr << "BASS_Init error: " << BASS_ErrorGetCode() << std::endl;
        return 1;
    }

    // 2) Load the SF2PACK soundfont
    HSOUNDFONT sfont = BASS_MIDI_FontInit(argv[1], 0);
    if (!sfont) {
        std::cerr << "BASS_MIDI_FontInit error: " << BASS_ErrorGetCode() << std::endl;
        BASS_Free();
        return 1;
    }

    // Create decoding stream
    DWORD flags = BASS_SAMPLE_FLOAT | BASS_MIDI_NOFX;
    if (synthToFile)
        flags |= BASS_STREAM_DECODE;

    HSTREAM midiStream = BASS_MIDI_StreamCreateFile(FALSE, argv[2], 0, 0, flags, 44100);
    if (!midiStream) {
        std::cerr << "BASS_MIDI_StreamCreateFile error: " << BASS_ErrorGetCode() << std::endl;
        BASS_MIDI_FontFree(sfont);
        BASS_Free();
        return 1;
    }

    // Assign soundfont
    BASS_MIDI_FONT fontConfig{ sfont, -1, 0 };
    if (!BASS_MIDI_StreamSetFonts(midiStream, &fontConfig, 1)) {
        std::cerr << "BASS_MIDI_StreamSetFonts error: " << BASS_ErrorGetCode() << std::endl;
        BASS_StreamFree(midiStream);
        BASS_MIDI_FontFree(sfont);
        BASS_Free();
        return 1;
    }

    if (!synthToFile) {
        std::cout << "Playing live...\n";
        if (!BASS_ChannelPlay(midiStream, FALSE)) {
            std::cerr << "BASS_ChannelPlay error: " << BASS_ErrorGetCode() << std::endl;
            BASS_StreamFree(midiStream);
            BASS_MIDI_FontFree(sfont);
            BASS_Free();
            return 1;
        }

        std::cout << "Press Enter to stop...\n";
        std::cin.get();  // wait for user input
    } else {
        std::ofstream wav("output.wav", std::ios::binary);
        if (!wav) {
            std::cerr << "Can't open output.wav for writing" << std::endl;
            return 1;
        }

        WriteWAVHeader(wav, 0);
        DWORD totalBytes = 0;
        const int bufferSize = 44100 * 2 * sizeof(float);
        float buffer[44100 * 2];
        QWORD totalRead = 0;
        QWORD totalStreamLength = BASS_ChannelGetLength(midiStream, BASS_POS_BYTE);

        std::cout << "Synthesis start\n";
        while (true) {
            int bytes = BASS_ChannelGetData(midiStream, buffer, bufferSize);
            if (bytes <= 0) break;

            short pcm[44100 * 2];
            for (int i = 0; i < bytes / 4; ++i) {
                float sample = buffer[i] * 32767.0f;
                if (sample > 32767.0f) sample = 32767.0f;
                if (sample < -32768.0f) sample = -32768.0f;
                pcm[i] = static_cast<short>(sample);
            }

            wav.write(reinterpret_cast<const char*>(pcm), bytes / 2);
            totalBytes += bytes / 2;
            totalRead += bytes;

            float percent = (totalStreamLength > 0) ? (100.0f * totalRead / totalStreamLength) : 0.0f;
            printf("\rProgress: %.2f%%", percent);
        }

        wav.seekp(0);
        WriteWAVHeader(wav, totalBytes);
        wav.close();
        std::cout << "\nSynthesis end\n";
        std::cout << "Exported to output.wav" << std::endl;
    }

    // 6) Clean up
    BASS_StreamFree(midiStream);     // free stream
    BASS_MIDI_FontFree(sfont);       // unload soundfont
    BASS_Free();                     // shut down BASS

    return 0;
}