#include "wasapi_loopback_capturer.h"
#include <audioclient.h>
#include <comdef.h>
#include <mmdeviceapi.h>
#include <windows.h>
#include <QDebug>

#undef min

class WASAPILoopbackCapturerPrivate
{
public:
    IAudioCapturer::Type deviceType;
    WASAPILoopbackCapturerPrivate()
    {
        CoInitialize(nullptr);
    }
    ~WASAPILoopbackCapturerPrivate()
    {
        cleanupSilencePlayer();
        
        if (pwfx)
            CoTaskMemFree(pwfx);
        if (pCaptureClient)
            pCaptureClient->Release();
        if (pAudioClient)
            pAudioClient->Release();
        if (pDevice)
            pDevice->Release();
        if (pEnumerator)
            pEnumerator->Release();
        CoUninitialize();
    }

    bool init(IAudioCapturer::Type type)
    {
        HRESULT hr;
        hr = CoCreateInstance(__uuidof(MMDeviceEnumerator),
                              nullptr,
                              CLSCTX_ALL,
                              __uuidof(IMMDeviceEnumerator),
                              (void**)&pEnumerator);
        if (FAILED(hr)) {
            qDebug() << "Failed to create MMDeviceEnumerator";
            return false;
        }

        deviceType = type;
        EDataFlow dataFlow = (deviceType == IAudioCapturer::Type::Speaker) ? eRender : eCapture;

        hr = pEnumerator->GetDefaultAudioEndpoint(dataFlow, eConsole, &pDevice);
        if (FAILED(hr)) {
            qDebug() << "Failed to get default audio endpoint";
            return false;
        }

        hr = pDevice->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr, (void**)&pAudioClient);
        if (FAILED(hr)) {
            qDebug() << "Failed to activate audio client";
            return false;
        }
        return true;
    }

    // 初始化音频客户端
    bool initializeAudioClient()
    {
        AUDCLNT_SHAREMODE shareMode = AUDCLNT_SHAREMODE_SHARED;
        DWORD streamFlags = AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM
                            | AUDCLNT_STREAMFLAGS_SRC_DEFAULT_QUALITY;
        if (deviceType == IAudioCapturer::Type::Speaker) {
            streamFlags |= AUDCLNT_STREAMFLAGS_LOOPBACK;
        }
        REFERENCE_TIME hnsBufferDuration = 0;

        HRESULT hr
            = pAudioClient->Initialize(shareMode, streamFlags, hnsBufferDuration, 0, pwfx, nullptr);
        if (FAILED(hr)) {
            qDebug() << "Failed to initialize audio client, HRESULT:" << QString::number(hr, 16);
            return false;
        }
        qDebug() << "Audio client initialized successfully";
        return true;
    }

    bool setupCaptureClient()
    {
        if (!pAudioClient) {
            qDebug() << "Audio client is null, cannot get capture client";
            return false;
        }
        
        HRESULT hr = pAudioClient->GetService(__uuidof(IAudioCaptureClient), (void**) &pCaptureClient);
        if (FAILED(hr)) {
            qDebug() << "Failed to get capture client, HRESULT:" << QString::number(hr, 16) 
                     << "AudioClient:" << pAudioClient << "CaptureClient:" << pCaptureClient;
            return false;
        }
        
        qDebug() << "Capture client obtained successfully:" << pCaptureClient;
        return true;
    }

    bool startCapture()
    {
        if (!pAudioClient) {
            qDebug() << "Audio client is null, cannot start capture";
            return false;
        }
        
        HRESULT hr = pAudioClient->Start();
        if (FAILED(hr)) {
            qDebug() << "Failed to start audio client, HRESULT:" << QString::number(hr, 16);
            return false;
        }
        
        qDebug() << "Audio client started successfully";
        return true;
    }

    void stopCapture()
    {
        if (pAudioClient) {
            pAudioClient->Stop();
        }
    }

    // 获取音频格式并设置到AudioFormat结构中
    bool setupAudioFormat(AudioFormat& audioFormat)
    {
        if (!pwfx) {
            return false;
        }

        // 设置音频格式
        audioFormat.sampleRate = pwfx->nSamplesPerSec;
        audioFormat.channels = pwfx->nChannels;
        audioFormat.bitsPerSample = pwfx->wBitsPerSample;
        audioFormat.blockAlign = pwfx->nBlockAlign;
        audioFormat.avgBytesPerSec = pwfx->nAvgBytesPerSec;

        // 如果是浮点格式，保持32bit，让混音器处理格式转换
        if (pwfx->wFormatTag == WAVE_FORMAT_IEEE_FLOAT) {
            qDebug() << "Keeping 32t float format for mixer processing";
        }

        return true;
    }

    // 处理音频数据
    int processAudioData(std::vector<char>& buffer, std::mutex& mutex)
    {
        UINT32 packetLength = 0;
        HRESULT hr = pCaptureClient->GetNextPacketSize(&packetLength);
        if (FAILED(hr)) {
            qDebug() << "Failed to get next packet size";
            return -1;
        }
        
        if (packetLength == 0) {
            return 0;
        }
        
        BYTE* pData = nullptr;
        UINT32 numFrames = 0;
        DWORD flags = 0;
        hr = pCaptureClient->GetBuffer(&pData, &numFrames, &flags, nullptr, nullptr);
        if (FAILED(hr)) {
            qDebug() << "Failed to get buffer";
            return -1;
        }
        
        int bytes = numFrames * pwfx->nBlockAlign;

        if (pwfx->wFormatTag == WAVE_FORMAT_IEEE_FLOAT) {
            // 32bit浮点格式，直接传递原始数据，让混音器处理格式转换
            std::lock_guard<std::mutex> lock(mutex);
            // 限制缓冲区大小，避免内存无限增长
            const size_t maxBufferSize = 1024 * 1024; // 1MB
            int bytesToAdd = numFrames * pwfx->nBlockAlign;
            if (buffer.size() + bytesToAdd > maxBufferSize) {
                buffer.erase(buffer.begin(), buffer.begin() + (buffer.size() - bytesToAdd));
            }
            buffer.insert(buffer.end(),
                          (char*) pData,
                          (char*) pData + bytesToAdd);
        } else {
            // PCM 直接拷贝
            std::lock_guard<std::mutex> lock(mutex);
            // 限制缓冲区大小，避免内存无限增长
            const size_t maxBufferSize = 1024 * 1024; // 1MB
            int bytesToAdd = numFrames * pwfx->nBlockAlign;
            if (buffer.size() + bytesToAdd > maxBufferSize) {
                buffer.erase(buffer.begin(), buffer.begin() + (buffer.size() - bytesToAdd));
            }
            buffer.insert(buffer.end(),
                          (char*) pData,
                          (char*) pData + bytesToAdd);
        }

        pCaptureClient->ReleaseBuffer(numFrames);
        return bytes;
    }

    IMMDeviceEnumerator* pEnumerator = nullptr;
    IMMDevice* pDevice = nullptr;
    IAudioClient* pAudioClient = nullptr;
    IAudioCaptureClient* pCaptureClient = nullptr;
    WAVEFORMATEX* pwfx = nullptr;

    WAVEFORMATEXTENSIBLE _formatex;
    
    // 静音播放器相关成员
    std::atomic<bool> m_silencePlayerRunning{false};
    std::thread m_silencePlayerThread;
    IMMDevice* pSilenceDevice = nullptr;
    IAudioClient* pSilenceAudioClient = nullptr;
    IAudioRenderClient* pSilenceRenderClient = nullptr;
    
    // 静音播放器方法
    bool initializeSilencePlayer();
    void cleanupSilencePlayer();
    void silencePlayerThreadFunc();
};

WASAPILoopbackCapturer::WASAPILoopbackCapturer(QObject* parent)
    : QObject(parent)
    , d(new WASAPILoopbackCapturerPrivate)
{

}

WASAPILoopbackCapturer::~WASAPILoopbackCapturer()
{
    Stop();
    delete d;
}

bool WASAPILoopbackCapturer::Init(Type deviceType)
{
    m_deviceType = deviceType;
    if (!d->init(deviceType)) {
        qDebug() << "Failed to initialize WASAPI components";
        return false;
    }

    // 尝试使用默认格式
    AudioFormat defaultFormat = IAudioCapturer::GetDefaultFormat();
    WAVEFORMATEX* formatToUse = (WAVEFORMATEX*)CoTaskMemAlloc(sizeof(WAVEFORMATEX));
    if (!formatToUse) {
        qDebug() << "Failed to allocate memory for WAVEFORMATEX";
        return false;
    }
    formatToUse->wFormatTag = WAVE_FORMAT_PCM;
    formatToUse->nChannels = defaultFormat.channels;
    formatToUse->nSamplesPerSec = defaultFormat.sampleRate;
    formatToUse->wBitsPerSample = defaultFormat.bitsPerSample;
    formatToUse->nBlockAlign = (formatToUse->nChannels * formatToUse->wBitsPerSample) / 8;
    formatToUse->nAvgBytesPerSec = formatToUse->nSamplesPerSec * formatToUse->nBlockAlign;
    formatToUse->cbSize = 0;

    d->pwfx = formatToUse;

    // 尝试使用默认格式初始化，如果失败则回退
    if (!d->initializeAudioClient()) {
        qWarning("Default audio format not supported, falling back to GetMixFormat.");

        CoTaskMemFree(d->pwfx);
        d->pwfx = nullptr;

        // 使用GetMixFormat获取系统首选格式
        HRESULT hr = d->pAudioClient->GetMixFormat(&d->pwfx);
        if (FAILED(hr)) {
            qDebug() << "Failed to GetMixFormat on fallback, HRESULT:" << QString::number(hr, 16);
            return false;
        }

        // 再次尝试初始化
        if (!d->initializeAudioClient()) {
            qDebug() << "Failed to initialize audio client on fallback.";
            CoTaskMemFree(d->pwfx);
            d->pwfx = nullptr;
            return false;
        }
    }

    // 设置音频格式
    if (!d->setupAudioFormat(m_audioFormat)) {
        qDebug() << "Failed to setup audio format";
        return false;
    }

    qDebug() << "WASAPI Loopback Capturer initialized successfully";
    return true;
}

bool WASAPILoopbackCapturer::Start()
{
    if (m_running)
        return false;
        
    // 如果是扬声器设备，启动静音播放器确保音频引擎活跃
    if (m_deviceType == Type::Speaker) {
        if (!d->initializeSilencePlayer()) {
            qDebug() << "Failed to initialize silence player";
            return false;
        }
    }
        
    if (!d->setupCaptureClient()) {
        return false;
    }
    
    if (!d->startCapture()) {
        return false;
    }
    
    m_running = true;
    m_captureThread = std::thread(&WASAPILoopbackCapturer::captureThreadFunc, this);
    return true;
}

void WASAPILoopbackCapturer::Stop()
{
    if (!m_running)
        return;
    m_running = false;
    if (m_captureThread.joinable())
        m_captureThread.join();
    d->stopCapture();
}

const AudioFormat& WASAPILoopbackCapturer::GetFormat() const
{
    return m_audioFormat;
}

int WASAPILoopbackCapturer::readAudioData(char* buf, int maxLen)
{
    std::unique_lock<std::mutex> lock(m_mutex);
    
    // 按帧对齐读取，确保不会破坏音频帧
    int blockAlign = m_audioFormat.blockAlign;
    if (blockAlign <= 0) {
        return 0;
    }
    
    // 计算可以读取的完整帧数
    int availableFrames = m_buffer.size() / blockAlign;
    int requestedFrames = maxLen / blockAlign;
    int framesToRead = std::min(availableFrames, requestedFrames);
    
    if (framesToRead > 0) {
        int bytesToRead = framesToRead * blockAlign;
        memcpy(buf, m_buffer.data(), bytesToRead);
        m_buffer.erase(m_buffer.begin(), m_buffer.begin() + bytesToRead);
        return bytesToRead;
    }
    
    return 0;
}

void WASAPILoopbackCapturer::captureThreadFunc()
{
    qDebug() << "WASAPI Loopback capture started successfully";

    while (m_running) {
        int result = d->processAudioData(m_buffer, m_mutex);
        if (result < 0) {
            qDebug() << "Error processing audio data";
            break;
        }
        if (result == 0) {
            Sleep(10);
        }
    }

    qDebug() << "WASAPI Loopback capture stopped";
}

// 静音播放器实现
bool WASAPILoopbackCapturerPrivate::initializeSilencePlayer()
{
    HRESULT hr;
    
    // 获取默认音频渲染设备
    hr = pEnumerator->GetDefaultAudioEndpoint(eRender, eConsole, &pSilenceDevice);
    if (FAILED(hr)) {
        qDebug() << "Failed to get default audio render endpoint for silence player";
        return false;
    }
    
    // 激活音频客户端
    hr = pSilenceDevice->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr, (void**)&pSilenceAudioClient);
    if (FAILED(hr)) {
        qDebug() << "Failed to activate audio client for silence player";
        return false;
    }
    
    // 获取混音格式
    WAVEFORMATEX* pSilenceFormat = nullptr;
    hr = pSilenceAudioClient->GetMixFormat(&pSilenceFormat);
    if (FAILED(hr)) {
        qDebug() << "Failed to get mix format for silence player";
        return false;
    }
    
    // 初始化音频客户端（共享模式）
    hr = pSilenceAudioClient->Initialize(
        AUDCLNT_SHAREMODE_SHARED,
        0,
        10000000,  // 1秒缓冲区
        0,
        pSilenceFormat,
        nullptr
    );
    
    CoTaskMemFree(pSilenceFormat);
    
    if (FAILED(hr)) {
        qDebug() << "Failed to initialize audio client for silence player";
        return false;
    }
    
    // 获取渲染客户端
    hr = pSilenceAudioClient->GetService(__uuidof(IAudioRenderClient), (void**)&pSilenceRenderClient);
    if (FAILED(hr)) {
        qDebug() << "Failed to get render client for silence player";
        return false;
    }
    
    // 启动音频客户端
    hr = pSilenceAudioClient->Start();
    if (FAILED(hr)) {
        qDebug() << "Failed to start audio client for silence player";
        return false;
    }
    
    // 启动静音播放线程
    m_silencePlayerRunning = true;
    m_silencePlayerThread = std::thread(&WASAPILoopbackCapturerPrivate::silencePlayerThreadFunc, this);
    
    qDebug() << "Silence player initialized successfully";
    return true;
}

void WASAPILoopbackCapturerPrivate::cleanupSilencePlayer()
{
    // 停止静音播放线程
    if (m_silencePlayerRunning) {
        m_silencePlayerRunning = false;
        if (m_silencePlayerThread.joinable()) {
            m_silencePlayerThread.join();
        }
    }
    
    // 清理 WASAPI 资源
    if (pSilenceAudioClient) {
        pSilenceAudioClient->Stop();
        pSilenceAudioClient->Release();
        pSilenceAudioClient = nullptr;
    }
    
    if (pSilenceRenderClient) {
        pSilenceRenderClient->Release();
        pSilenceRenderClient = nullptr;
    }
    
    if (pSilenceDevice) {
        pSilenceDevice->Release();
        pSilenceDevice = nullptr;
    }
}

void WASAPILoopbackCapturerPrivate::silencePlayerThreadFunc()
{
    qDebug() << "Silence player thread started";
    
    UINT32 bufferFrameCount;
    HRESULT hr = pSilenceAudioClient->GetBufferSize(&bufferFrameCount);
    if (FAILED(hr)) {
        qDebug() << "Failed to get buffer size for silence player";
        return;
    }
    
    while (m_silencePlayerRunning) {
        UINT32 numFramesPadding;
        hr = pSilenceAudioClient->GetCurrentPadding(&numFramesPadding);
        if (FAILED(hr)) {
            qDebug() << "Failed to get current padding for silence player";
            break;
        }
        
        UINT32 numFramesAvailable = bufferFrameCount - numFramesPadding;
        
        if (numFramesAvailable > 0) {
            BYTE* pData;
            hr = pSilenceRenderClient->GetBuffer(numFramesAvailable, &pData);
            if (SUCCEEDED(hr)) {
                // 填充静音数据（全零）
                memset(pData, 0, numFramesAvailable * sizeof(float) * 2); // 假设立体声
                hr = pSilenceRenderClient->ReleaseBuffer(numFramesAvailable, 0);
                if (FAILED(hr)) {
                    qDebug() << "Failed to release buffer for silence player";
                }
            }
        }
        
        Sleep(10); // 10ms 间隔
    }
    
    qDebug() << "Silence player thread stopped";
}