#pragma once
#include <string>
#include <vector>
#include <unordered_map>
#include <memory>
#include <fstream>
#include <cstring>
#include <cmath>
#include <mutex>
#include <iostream>
#include <algorithm>

#ifdef __APPLE__
#include <AudioToolbox/AudioToolbox.h>
#endif

// Lightweight audio engine with WAV loading and real-time mixing.
// Uses CoreAudio AudioQueue on macOS.
//
// Usage:
//   Audio audio;
//   audio.init();
//   audio.loadSound("jump", "assets/sounds/jump.wav");
//   audio.loadMusic("bgm", "assets/music/theme.wav");
//   audio.playSound("jump");
//   audio.playMusic("bgm");
//   audio.setMasterVolume(0.8f);
//   audio.update(); // call each frame
//   audio.shutdown();

class Audio {
public:
    // PCM audio clip loaded from WAV
    struct Clip {
        std::vector<int16_t> samples; // interleaved stereo
        int sampleRate = 44100;
        int channels = 2;
    };

    // Active playing voice
    struct Voice {
        const Clip* clip = nullptr;
        size_t position = 0;
        float volume = 1.0f;
        bool loop = false;
        bool active = false;
    };

    bool init(int sampleRate = 44100, int bufferSize = 4096) {
        sampleRate_ = sampleRate;
        bufferSize_ = bufferSize;
        mixBuffer_.resize(bufferSize * 2, 0); // stereo

#ifdef __APPLE__
        AudioStreamBasicDescription fmt{};
        fmt.mSampleRate = sampleRate;
        fmt.mFormatID = kAudioFormatLinearPCM;
        fmt.mFormatFlags = kLinearPCMFormatFlagIsSignedInteger | kLinearPCMFormatFlagIsPacked;
        fmt.mBitsPerChannel = 16;
        fmt.mChannelsPerFrame = 2;
        fmt.mBytesPerFrame = 4;
        fmt.mFramesPerPacket = 1;
        fmt.mBytesPerPacket = 4;

        OSStatus status = AudioQueueNewOutput(&fmt, audioCallback, this, nullptr,
                                               nullptr, 0, &queue_);
        if (status != noErr) {
            std::cerr << "Audio: failed to create AudioQueue (" << status << ")" << std::endl;
            return false;
        }

        // Allocate and enqueue buffers
        for (int i = 0; i < NUM_BUFFERS; ++i) {
            AudioQueueAllocateBuffer(queue_, bufferSize * 4, &buffers_[i]);
            buffers_[i]->mAudioDataByteSize = bufferSize * 4;
            memset(buffers_[i]->mAudioData, 0, bufferSize * 4);
            AudioQueueEnqueueBuffer(queue_, buffers_[i], 0, nullptr);
        }

        AudioQueueStart(queue_, nullptr);
        initialized_ = true;
        std::cout << "Audio initialized (CoreAudio, " << sampleRate << " Hz)" << std::endl;
        return true;
#else
        std::cout << "Audio: no backend for this platform" << std::endl;
        return false;
#endif
    }

    void shutdown() {
#ifdef __APPLE__
        if (queue_) {
            AudioQueueStop(queue_, true);
            for (int i = 0; i < NUM_BUFFERS; ++i)
                if (buffers_[i]) AudioQueueFreeBuffer(queue_, buffers_[i]);
            AudioQueueDispose(queue_, true);
            queue_ = nullptr;
        }
#endif
        initialized_ = false;
    }

    // Load a WAV file as a sound effect
    bool loadSound(const std::string& name, const std::string& path) {
        auto clip = loadWAV(path);
        if (!clip) return false;
        std::lock_guard<std::mutex> lock(mutex_);
        sounds_[name] = std::move(clip);
        std::cout << "Audio: loaded sound '" << name << "' from " << path << std::endl;
        return true;
    }

    // Load a WAV file as music (same format, just tagged differently)
    bool loadMusic(const std::string& name, const std::string& path) {
        return loadSound(name, path);
    }

    // Play a sound effect (fire and forget)
    void playSound(const std::string& name, float volume = 1.0f) {
        std::lock_guard<std::mutex> lock(mutex_);
        auto it = sounds_.find(name);
        if (it == sounds_.end()) return;

        // Find free voice or oldest
        Voice* best = nullptr;
        for (auto& v : sfxVoices_) {
            if (!v.active) { best = &v; break; }
        }
        if (!best) best = &sfxVoices_[0]; // steal oldest

        best->clip = it->second.get();
        best->position = 0;
        best->volume = volume * sfxVolume_;
        best->loop = false;
        best->active = true;
    }

    // Play music (looping, replaces current music)
    void playMusic(const std::string& name, float volume = 1.0f) {
        std::lock_guard<std::mutex> lock(mutex_);
        auto it = sounds_.find(name);
        if (it == sounds_.end()) return;

        musicVoice_.clip = it->second.get();
        musicVoice_.position = 0;
        musicVoice_.volume = volume * musicVolume_;
        musicVoice_.loop = true;
        musicVoice_.active = true;
    }

    void stopMusic() {
        std::lock_guard<std::mutex> lock(mutex_);
        musicVoice_.active = false;
    }

    void setMasterVolume(float v) { masterVolume_ = std::clamp(v, 0.0f, 1.0f); }
    void setSfxVolume(float v)    { sfxVolume_ = std::clamp(v, 0.0f, 1.0f); }
    void setMusicVolume(float v)  { musicVolume_ = std::clamp(v, 0.0f, 1.0f); }

    float getMasterVolume() const { return masterVolume_; }
    float getSfxVolume() const    { return sfxVolume_; }
    float getMusicVolume() const  { return musicVolume_; }

    bool isInitialized() const { return initialized_; }

private:
    static constexpr int NUM_BUFFERS = 3;
    static constexpr int MAX_SFX_VOICES = 16;

    int sampleRate_ = 44100;
    int bufferSize_ = 4096;
    bool initialized_ = false;

    float masterVolume_ = 0.8f;
    float sfxVolume_    = 1.0f;
    float musicVolume_  = 0.7f;

    std::mutex mutex_;
    std::unordered_map<std::string, std::unique_ptr<Clip>> sounds_;
    Voice sfxVoices_[MAX_SFX_VOICES]{};
    Voice musicVoice_{};
    std::vector<int16_t> mixBuffer_;

#ifdef __APPLE__
    AudioQueueRef queue_ = nullptr;
    AudioQueueBufferRef buffers_[NUM_BUFFERS]{};

    static void audioCallback(void* userData, AudioQueueRef queue, AudioQueueBufferRef buffer) {
        Audio* self = static_cast<Audio*>(userData);
        self->fillBuffer(static_cast<int16_t*>(buffer->mAudioData),
                         buffer->mAudioDataByteSize / sizeof(int16_t));
        AudioQueueEnqueueBuffer(queue, buffer, 0, nullptr);
    }
#endif

    void fillBuffer(int16_t* out, int totalSamples) {
        std::lock_guard<std::mutex> lock(mutex_);

        // Clear
        memset(out, 0, totalSamples * sizeof(int16_t));

        int frames = totalSamples / 2; // stereo

        // Mix music
        if (musicVoice_.active && musicVoice_.clip) {
            mixVoice(out, frames, musicVoice_, musicVolume_);
        }

        // Mix SFX voices
        for (auto& v : sfxVoices_) {
            if (v.active && v.clip) {
                mixVoice(out, frames, v, sfxVolume_);
            }
        }
    }

    void mixVoice(int16_t* out, int frames, Voice& voice, float categoryVol) {
        if (!voice.clip || !voice.active) return;
        const Clip& clip = *voice.clip;
        float vol = voice.volume * categoryVol * masterVolume_;

        for (int f = 0; f < frames; ++f) {
            if (voice.position >= clip.samples.size() / clip.channels) {
                if (voice.loop) {
                    voice.position = 0;
                } else {
                    voice.active = false;
                    break;
                }
            }

            if (clip.channels == 2) {
                size_t idx = voice.position * 2;
                int32_t l = out[f * 2]     + (int32_t)(clip.samples[idx]     * vol);
                int32_t r = out[f * 2 + 1] + (int32_t)(clip.samples[idx + 1] * vol);
                out[f * 2]     = (int16_t)std::clamp(l, (int32_t)-32768, (int32_t)32767);
                out[f * 2 + 1] = (int16_t)std::clamp(r, (int32_t)-32768, (int32_t)32767);
            } else {
                // Mono -> stereo
                size_t idx = voice.position;
                int16_t s = (int16_t)(clip.samples[idx] * vol);
                int32_t l = out[f * 2]     + s;
                int32_t r = out[f * 2 + 1] + s;
                out[f * 2]     = (int16_t)std::clamp(l, (int32_t)-32768, (int32_t)32767);
                out[f * 2 + 1] = (int16_t)std::clamp(r, (int32_t)-32768, (int32_t)32767);
            }
            voice.position++;
        }
    }

    // WAV file loader — parses RIFF/WAVE headers, extracts PCM data
    std::unique_ptr<Clip> loadWAV(const std::string& path) {
        std::ifstream file(path, std::ios::binary);
        if (!file.is_open()) {
            std::cerr << "Audio: cannot open " << path << std::endl;
            return nullptr;
        }

        // RIFF header
        char riff[4]; file.read(riff, 4);
        if (memcmp(riff, "RIFF", 4) != 0) {
            std::cerr << "Audio: not a RIFF file: " << path << std::endl;
            return nullptr;
        }

        uint32_t fileSize; file.read(reinterpret_cast<char*>(&fileSize), 4);
        char wave[4]; file.read(wave, 4);
        if (memcmp(wave, "WAVE", 4) != 0) {
            std::cerr << "Audio: not a WAVE file: " << path << std::endl;
            return nullptr;
        }

        uint16_t audioFormat = 0, numChannels = 0, bitsPerSample = 0;
        uint32_t sampleRate = 0, dataSize = 0;
        bool foundFmt = false, foundData = false;

        // Parse chunks
        while (file.good() && !(foundFmt && foundData)) {
            char chunkId[4]; file.read(chunkId, 4);
            uint32_t chunkSize; file.read(reinterpret_cast<char*>(&chunkSize), 4);

            if (memcmp(chunkId, "fmt ", 4) == 0) {
                file.read(reinterpret_cast<char*>(&audioFormat), 2);
                file.read(reinterpret_cast<char*>(&numChannels), 2);
                file.read(reinterpret_cast<char*>(&sampleRate), 4);
                uint32_t byteRate; file.read(reinterpret_cast<char*>(&byteRate), 4);
                uint16_t blockAlign; file.read(reinterpret_cast<char*>(&blockAlign), 2);
                file.read(reinterpret_cast<char*>(&bitsPerSample), 2);
                // Skip extra fmt bytes
                if (chunkSize > 16) file.seekg(chunkSize - 16, std::ios::cur);
                foundFmt = true;
            } else if (memcmp(chunkId, "data", 4) == 0) {
                dataSize = chunkSize;
                foundData = true;
            } else {
                file.seekg(chunkSize, std::ios::cur);
            }
        }

        if (!foundFmt || !foundData) {
            std::cerr << "Audio: missing fmt/data chunk in " << path << std::endl;
            return nullptr;
        }
        if (audioFormat != 1) {
            std::cerr << "Audio: only PCM WAV supported (got format " << audioFormat << ")" << std::endl;
            return nullptr;
        }

        auto clip = std::make_unique<Clip>();
        clip->sampleRate = sampleRate;
        clip->channels = numChannels;

        if (bitsPerSample == 16) {
            size_t numSamples = dataSize / 2;
            clip->samples.resize(numSamples);
            file.read(reinterpret_cast<char*>(clip->samples.data()), dataSize);
        } else if (bitsPerSample == 8) {
            // Convert 8-bit unsigned to 16-bit signed
            size_t numSamples = dataSize;
            std::vector<uint8_t> raw(numSamples);
            file.read(reinterpret_cast<char*>(raw.data()), dataSize);
            clip->samples.resize(numSamples);
            for (size_t i = 0; i < numSamples; ++i)
                clip->samples[i] = ((int16_t)raw[i] - 128) * 256;
        } else if (bitsPerSample == 24) {
            size_t numSamples = dataSize / 3;
            clip->samples.resize(numSamples);
            for (size_t i = 0; i < numSamples; ++i) {
                uint8_t b[3]; file.read(reinterpret_cast<char*>(b), 3);
                int32_t val = (b[2] << 24) | (b[1] << 16) | (b[0] << 8);
                clip->samples[i] = (int16_t)(val >> 16);
            }
        } else {
            std::cerr << "Audio: unsupported bit depth " << bitsPerSample << std::endl;
            return nullptr;
        }

        // Resample if needed (simple nearest-neighbor)
        if ((int)sampleRate != sampleRate_) {
            double ratio = (double)sampleRate_ / sampleRate;
            size_t newLen = (size_t)(clip->samples.size() / numChannels * ratio) * numChannels;
            std::vector<int16_t> resampled(newLen);
            for (size_t i = 0; i < newLen / numChannels; ++i) {
                size_t srcIdx = (size_t)(i / ratio);
                if (srcIdx >= clip->samples.size() / numChannels)
                    srcIdx = clip->samples.size() / numChannels - 1;
                for (int ch = 0; ch < numChannels; ++ch)
                    resampled[i * numChannels + ch] = clip->samples[srcIdx * numChannels + ch];
            }
            clip->samples = std::move(resampled);
            clip->sampleRate = sampleRate_;
        }

        return clip;
    }
};
