#include <iostream>
#include <fstream>
#include <algorithm>
#include <cmath>
#include <vector>
#include <sstream>
#include <cstdint>
#include <cstring>
#include <string>

const int SAMPLE_RATE = 44100;
const double M_TAU = 6.283185307179586476925286766559;

// Enum for waveform types
enum Waveform {
    SINE,
    TRIANGLE,
    SAWTOOTH,
    SQUARE,
    NOISE,
    DPCM
};

// Structure to hold the waveform parameters
struct Wave {
    Waveform type;
    int start_ms;
    int duration_ms;
    int frequency;
    double volume;
    double dutyCycle;   // For SQUARE
    double sweepDepth;  // Vibrato or slide
    double sweepRate;
    double attack;
    double decay;
    double sustain = 1.0;
    double release = 0.0;
    std::string samplePath;
    std::vector<int16_t> sampleData;
};

bool loadWavFile(const std::string& path, std::vector<int16_t>& outSamples, int& outDurationMs) {
    std::ifstream file(path, std::ios::binary);
    if (!file.is_open()) return false;

    file.seekg(44); // skip header (basic .wav, no format parsing)
    outSamples.clear();

    int16_t sample;
    while (file.read(reinterpret_cast<char*>(&sample), sizeof(sample))) {
        outSamples.push_back(sample);
    }

    int sampleRate = 44100;  // hardcoded for now
    outDurationMs = static_cast<int>((outSamples.size() * 1000.0) / sampleRate);

    return true;
}

// Read the parameters from the text file
bool loadWaveParams(const std::string& filename, std::vector<Wave>& waves, int& totalDuration) {
    std::ifstream file(filename);
    if (!file.is_open()) {
        std::cerr << "Error opening file!" << std::endl;
        return false;
    }

    std::string line;
    bool inCommentBlock = false;

    while (std::getline(file, line)) {
        // Handle multiline /* comments */
        if (inCommentBlock) {
            size_t end = line.find("*/");
            if (end != std::string::npos) {
                inCommentBlock = false;
                line = line.substr(end + 2);
            }
            else {
                continue;
            }
        }

        while (true) {
            size_t start = line.find("/*");
            if (start != std::string::npos) {
                size_t end = line.find("*/", start + 2);
                if (end != std::string::npos) {
                    line.erase(start, end - start + 2);  // Remove /* ... */
                }
                else {
                    line = line.substr(0, start);  // Cut off everything after /*
                    inCommentBlock = true;
                    break;
                }
            }
            else {
                break;
            }
        }

        // Trim leading whitespace
        line.erase(0, line.find_first_not_of(" \t\r\n"));
        if (line.empty() || line[0] == '#') continue;

        // Handle TOTAL_DURATION
        if (line.rfind("TOTAL_DURATION", 0) == 0) {
            std::istringstream durationStream(line);
            std::string keyword;
            durationStream >> keyword >> totalDuration;
            continue;
        }

        std::istringstream ss(line);
        std::string waveType;
        ss >> waveType;

        if (waveType.empty()) continue;

        Wave wave;
        wave.type = SINE;

        if (waveType == "SINE")         wave.type = SINE;
        else if (waveType == "SAWTOOTH") wave.type = SAWTOOTH;
        else if (waveType == "SQUARE")   wave.type = SQUARE;
        else if (waveType == "TRIANGLE") wave.type = TRIANGLE;
        else if (waveType == "NOISE")    wave.type = NOISE;
        else if (waveType == "DPCM")     wave.type = DPCM;
        else {
            std::cout << "Unknown waveform type: " << waveType << std::endl;
            continue;
        }

        int start_ms, duration_ms;
        ss >> start_ms;
        wave.start_ms = start_ms;

        if (wave.type == DPCM) {
            std::string token;

            // Skip any comment block
            while (ss >> token && token == "/*") {
                while (ss >> token && token != "*/");
            }

            if (token.empty()) {
                std::cout << "Missing sample path for DPCM at " << start_ms << "ms\n";
                continue;
            }

            wave.samplePath = token;

            // Optional volume
            if (!(ss >> wave.volume)) wave.volume = 1.0;

            // Optional ADSR envelope for DPCM
            if (!(ss >> wave.attack)) wave.attack = 0.0;
            if (!(ss >> wave.decay)) wave.decay = 0.0;
            if (!(ss >> wave.sustain)) wave.sustain = 1.0;
            if (!(ss >> wave.release)) wave.release = 0.0;

            // Not used for DPCM but set anyway
            wave.frequency = 0;
            wave.dutyCycle = 0.5;
            wave.sweepDepth = 0.0;
            wave.sweepRate = 0.0;

            int actualLength = 0;
            if (!loadWavFile(wave.samplePath, wave.sampleData, actualLength)) {
                std::cout << "Failed to load sample: " << wave.samplePath << std::endl;
            }

            wave.duration_ms = actualLength;
            wave.start_ms = start_ms;

            waves.push_back(wave);
            continue;
        }

        ss >> duration_ms;
        wave.duration_ms = duration_ms;

        // Standard wave parameters
        double volume = 1.0;
        double duty = 0.5;
        double sweepDepth = 0.0, sweepRate = 0.0;
        double attack = 0.01, decay = 0.2;
        int frequency;

        ss >> frequency;
        ss >> volume;
        if (!(ss >> duty)) duty = 0.5;
        if (!(ss >> sweepDepth)) sweepDepth = 0.0;
        if (!(ss >> sweepRate)) sweepRate = 0.0;
        if (!(ss >> attack)) attack = 0.01;
        if (!(ss >> decay)) decay = 0.2;

        wave.frequency = frequency;
        wave.volume = volume;
        wave.dutyCycle = duty;
        wave.sweepDepth = sweepDepth;
        wave.sweepRate = sweepRate;
        wave.attack = attack;
        wave.decay = decay;

        waves.push_back(wave);
    }

    return true;
}

// Waveform generation functions
double sineWave(double frequency, double time) {
    return std::sin(M_TAU * frequency * time);
}

double triangleWave(double frequency, double time) {
    double period = 1.0 / frequency;
    double phase = fmod(time, period) / period;  // Normalized phase (0 to 1)
    if (phase < 0.5) {
        return 4 * phase - 1;  // Rising part of the triangle
    }
    else {
        return 3 - 4 * phase;  // Falling part of the triangle
    }
}

double sawtoothWave(double frequency, double time) {
    double period = 1.0 / frequency;
    double phase = fmod(time, period) / period;  // Normalized phase (0 to 1)
    return 2 * phase - 1;  // Ramp from -1 to 1
}

double squareWave(double frequency, double time, double duty = 0.5) {
    double period = 1.0 / frequency;
    double phase = fmod(time, period) / period;  // 0 to 1
    return phase < duty ? 1.0 : -1.0;
}

// Mix the waveforms with volume control
double generateWave(Waveform waveform, double frequency, double time, double duty = 0.5) {
    switch (waveform) {
    case SINE:     return sineWave(frequency, time);
    case TRIANGLE: return triangleWave(frequency, time);
    case SAWTOOTH: return sawtoothWave(frequency, time);
    case SQUARE:   return squareWave(frequency, time, duty);
    default:       return sineWave(frequency, time);
    }
}

double applyVibrato(double baseFreq, double time, double depth = 5.0, double rate = 5.0) {
    return baseFreq + depth * std::sin(M_TAU * rate * time);
}

double applyEnvelope(double time, double startTime, double duration, double volume, double attack, double decay) {
    double elapsed = time - startTime;
    if (elapsed < 0.0) return 0.0;

    if (elapsed < attack) {
        // Attack phase: ramp up from 0 to volume
        return volume * (elapsed / attack);
    }

    if (elapsed < (attack + decay)) {
        // Decay phase: ramp down to 0
        return volume * (1.0 - ((elapsed - attack) / decay));
    }

    // After envelope ends
    return 0.0;
}

double noiseWave() {
    static uint16_t lfsr = 0xACE1u;
    static unsigned bit;

    // 16-bit Galois LFSR
    bit = ((lfsr >> 0) ^ (lfsr >> 2) ^ (lfsr >> 3) ^ (lfsr >> 5)) & 1;
    lfsr = (lfsr >> 1) | (bit << 15);
    return ((lfsr & 1) ? 1.0 : -1.0);  // White-ish noise
}

std::vector<int16_t> generateAudio(const std::vector<Wave>& waves, int totalDuration) {
    std::vector<int16_t> samples(SAMPLE_RATE * totalDuration / 1000, 0); // Initialize with zeros

    for (const auto& wave : waves) {
        int startSample = wave.start_ms * SAMPLE_RATE / 1000;
        int endSample = (wave.start_ms + wave.duration_ms) * SAMPLE_RATE / 1000;

        if (wave.type == DPCM) {
            // DPCM Sample Playback (raw PCM sample, assumed 16-bit mono)
            for (int i = 0; i < wave.sampleData.size(); ++i) {
                int t = startSample + i;
                if (t >= samples.size()) break;

                double envelopeVol = applyEnvelope(
                    t / (double)SAMPLE_RATE,
                    wave.start_ms / 1000.0,
                    wave.duration_ms / 1000.0,
                    1.0, wave.attack, wave.decay
                );

                samples[t] += static_cast<int16_t>(wave.sampleData[i] * envelopeVol);
            }
            continue;  // Skip standard waveform generation for DPCM
        }

        for (int t = startSample; t < endSample; ++t) {
            if (t >= samples.size()) break;

            double time = static_cast<double>(t) / SAMPLE_RATE;
            double startTime = wave.start_ms / 1000.0;
            double relTime = time - startTime;
            if (relTime < 0) continue;

            // Frequency sweep / vibrato
            double modFreq = wave.frequency + wave.sweepDepth * std::sin(M_TAU * wave.sweepRate * time);

            // Generate waveform
            double rawSample = 0.0;
            if (wave.type == NOISE) {
                rawSample = noiseWave(); // LFSR noise
            }
            else {
                rawSample = generateWave(wave.type, modFreq, time, wave.dutyCycle);
            }

            // Envelope application
            double envelopeVolume = applyEnvelope(time, startTime, wave.duration_ms / 1000.0, wave.volume, wave.attack, wave.decay);

            // Final sample
            double finalSample = rawSample * envelopeVolume;

            // Mix into audio buffer
            samples[t] += static_cast<int16_t>(std::clamp(finalSample * 32767.0, -32768.0, 32767.0));
        }
    }

    return samples;
}

// Write WAV file
void writeWav(const std::string& filename, const std::vector<int16_t>& samples) {
    std::ofstream file(filename, std::ios::binary);

    // WAV Header
    struct WAVHeader {
        char riff[4]; // "RIFF"
        uint32_t fileSize;
        char wave[4]; // "WAVE"
        char fmt[4];  // "fmt "
        uint32_t fmtSize;
        uint16_t audioFormat;
        uint16_t numChannels;
        uint32_t sampleRate;
        uint32_t byteRate;
        uint16_t blockAlign;
        uint16_t bitsPerSample;
        char data[4]; // "data"
        uint32_t dataSize;
    };

    WAVHeader header;
    memcpy(header.riff, "RIFF", 4);
    header.fileSize = 36 + samples.size() * sizeof(int16_t);
    memcpy(header.wave, "WAVE", 4);
    memcpy(header.fmt, "fmt ", 4);
    header.fmtSize = 16;
    header.audioFormat = 1;  // PCM
    header.numChannels = 1;  // Mono
    header.sampleRate = SAMPLE_RATE;
    header.byteRate = SAMPLE_RATE * sizeof(int16_t);
    header.blockAlign = sizeof(int16_t);
    header.bitsPerSample = 16;
    memcpy(header.data, "data", 4);
    header.dataSize = samples.size() * sizeof(int16_t);

    file.write(reinterpret_cast<char*>(&header), sizeof(header));
    file.write(reinterpret_cast<const char*>(samples.data()), samples.size() * sizeof(int16_t));
}

int main(int argc, char* argv[]) {
    std::string out_name;
    if (argc == 2)
        out_name = argv[1];
    else
        out_name = "output.wav";

    std::vector<Wave> waves;
    int totalDuration = 0;

    // Load the wave parameters from a text file
    if (!loadWaveParams("waves.txt", waves, totalDuration)) {
        std::cerr << "Failed to load wave parameters!" << std::endl;
        return 1;
    }

    // Generate the audio by mixing the waveforms
    std::vector<int16_t> samples = generateAudio(waves, totalDuration);

    // Write the audio to a WAV file
    writeWav(out_name, samples);

    std::cout << "WAV file generated: " << out_name << std::endl;
    return 0;
}