#include <windows.h>
#include <stdio.h>
#include <mmreg.h>
#include <mmdeviceapi.h>
#include <functiondiscoverykeys.h>
#include <Audioclient.h>
#include <process.h>
#include <avrt.h>

#include "kbwasapi_device.h"
#include "kbwasapi_enum.h"
#include "kbwasapi_config.h"

//AVRT.DLL関係
//AvSetMmThreadCharacteristics
//AvRevertMmThreadCharacteristics
//AvSetMmThreadPriority
//を使うと WinXP で kbwasapi.kpi の LoadLibrary に失敗するので
//AVRT.DLL を動的にロードする
//(リンカオプションで AVRT.DLL を遅延ロードした方が簡単だが)
//(kbwasapi.kpi の LoadLibrary に成功したところで XP 環境で
// WASAPI が使えないことに変わりはない)
typedef HANDLE(WINAPI* fpAvSetMmThreadCharacteristics)(LPCWSTR TaskName,
                                                       LPDWORD TaskIndex);
typedef BOOL(WINAPI * fpAvRevertMmThreadCharacteristics)(HANDLE AvrtHandle);
typedef BOOL(WINAPI * fpAvSetMmThreadPriority)(HANDLE AvrtHandle,
                                               AVRT_PRIORITY Priority); 
static fpAvSetMmThreadCharacteristics    _AvSetMmThreadCharacteristics;
static fpAvRevertMmThreadCharacteristics _AvRevertMmThreadCharacteristics;
static fpAvSetMmThreadPriority           _AvSetMmThreadPriority;

HINSTANCE KbWasapiDevice::s_hAvrtDll = NULL;
///////////////////////////////////////////////////////////////////////////////
BOOL KbWasapiDevice::load_avrt(void)
{
    if(s_hAvrtDll){
        return TRUE;
    }
    s_hAvrtDll = LoadLibrary(L"avrt.dll");
    if (s_hAvrtDll) {
        _AvSetMmThreadCharacteristics =
            (fpAvSetMmThreadCharacteristics)GetProcAddress(s_hAvrtDll,
                "AvSetMmThreadCharacteristicsW");
        _AvRevertMmThreadCharacteristics =
            (fpAvRevertMmThreadCharacteristics)GetProcAddress(s_hAvrtDll,
                "AvRevertMmThreadCharacteristics");
        _AvSetMmThreadPriority = (fpAvSetMmThreadPriority)GetProcAddress
        (s_hAvrtDll, "AvSetMmThreadPriority");
        if (!_AvSetMmThreadCharacteristics ||
            !_AvRevertMmThreadCharacteristics ||
            !_AvSetMmThreadPriority) {
            _AvSetMmThreadCharacteristics = NULL;
            _AvRevertMmThreadCharacteristics = NULL;
            _AvSetMmThreadPriority = NULL;
            FreeLibrary(s_hAvrtDll);
            s_hAvrtDll = NULL;
        }
    }
    return s_hAvrtDll != NULL;
}
void KbWasapiDevice::unload_avrt(void)
{
    if(s_hAvrtDll){
        FreeLibrary(s_hAvrtDll);
        s_hAvrtDll = NULL;
        _AvSetMmThreadCharacteristics = NULL;
        _AvRevertMmThreadCharacteristics = NULL;
        _AvSetMmThreadPriority = NULL;
    }
}
///////////////////////////////////////////////////////////////////////////////
KbWasapiDevice::KbWasapiDevice(void)
{
    InitializeCriticalSection(&m_cs);
    m_pDevice = NULL;
    m_pAudioClient = NULL;
    m_pRenderClient = NULL;
    m_pAudioClock = NULL;
    m_nSampleRate = 0;
    m_nChannels = 0;
    m_nBps = 0;
    m_nBufferSizeSample = 0;
    m_nBlockAlign = 0;
    m_hEvent = NULL;
    m_hEventStop = NULL;
    m_hEventCanOutput = NULL;
    m_hThread = NULL;
    m_nBufferFrame = 0;
    m_nMode = KbWasapiEnumerator::MODE_EVENT;
    m_dwTime = 0;
    m_nFinished = 0;
    m_qwTotalWritten = 0;
    m_qwProcessedSample = 0;
    m_qwProcessedMax = 0;
    m_qwPrevPosition = 0;
    m_qwLatencySample = 0;
}
///////////////////////////////////////////////////////////////////////////////
KbWasapiDevice::~KbWasapiDevice(void)
{
    if(m_pAudioClock)m_pAudioClock->Release();
    if(m_pRenderClient)m_pRenderClient->Release();
    if(m_pAudioClient)m_pAudioClient->Release();
    if(m_pDevice)m_pDevice->Release();
    if(m_hEvent)CloseHandle(m_hEvent);
    if(m_hEventStop)CloseHandle(m_hEventStop);
    if(m_hEventCanOutput)CloseHandle(m_hEventCanOutput);
    if(m_hThread)CloseHandle(m_hThread);
    DeleteCriticalSection(&m_cs);
}
///////////////////////////////////////////////////////////////////////////////
unsigned __stdcall KbWasapiDevice::ThreadProc(void *p)
{
    KbWasapiDevice *pThis = (KbWasapiDevice*)p;
    pThis->ThreadProc();
    return 0;
}
///////////////////////////////////////////////////////////////////////////////
void __fastcall KbWasapiDevice::ThreadProc(void)
{
    DWORD dwTaskIndex = 0;
    HANDLE hTask = NULL;
    if(load_avrt()){
        hTask = _AvSetMmThreadCharacteristics(L"Pro Audio", &dwTaskIndex);
        if (hTask) {
            _AvSetMmThreadPriority(hTask, AVRT_PRIORITY_HIGH/*AVRT_PRIORITY_CRITICAL*/);
        }
    }
    m_pAudioClient->Start();
    HANDLE hEvents[] = {m_hEvent, m_hEventStop};
    while(1){
        int nWriteSample;
        if(IsEventMode()){//event
            DWORD dwWaitRet = WaitForMultipleObjects(2, hEvents, FALSE, INFINITE);
            if(dwWaitRet == WAIT_OBJECT_0 + 1){
                break;//Stop が呼ばれた
            }
            nWriteSample = m_nBufferFrame;
        }
        else{//push
            UINT uCurrentPadding;
            m_pAudioClient->GetCurrentPadding(&uCurrentPadding);
            nWriteSample = m_nBufferFrame - uCurrentPadding;
            if(nWriteSample <= 0){
            //if(nWriteSample < m_nBufferFrame / 2){
                DWORD dwWaitTime = (double)m_nBufferFrame/(double)m_nSampleRate * 1000.0;
                if(WaitForSingleObject(m_hEventStop, dwWaitTime/2) == WAIT_OBJECT_0){
                    break;
                }
                continue;
            }
        }
        Lock();
        BYTE *pData = NULL;
        if(FAILED(m_pRenderClient->GetBuffer(nWriteSample, &pData))){
            Unlock();
            break;
        }
        int nReadBytes = m_RingBuffer.Read(pData, nWriteSample * m_nBlockAlign);
        if(nReadBytes < nWriteSample * m_nBlockAlign){
            SetEvent(m_hEventCanOutput);
            while(!m_nFinished && nReadBytes < nWriteSample * m_nBlockAlign){
                Unlock();
                if(WaitForSingleObject(m_hEventStop, 1) == WAIT_OBJECT_0){
                //演奏停止
                    goto EXIT;
                }
                Lock();
                nReadBytes += m_RingBuffer.Read(pData + nReadBytes, nWriteSample * m_nBlockAlign - nReadBytes);
            }
            ZeroMemory(pData + nReadBytes, nWriteSample * m_nBlockAlign - nReadBytes);
        }
        m_qwProcessedSample = m_qwProcessedMax;
        m_qwProcessedMax += nWriteSample;
        m_pRenderClient->ReleaseBuffer(nWriteSample, 0);
        m_dwTime = timeGetTime();
        if(!m_nFinished){
            if(m_RingBuffer.GetRemain() >= m_nBufferSizeSample * m_nBlockAlign){
                SetEvent(m_hEventCanOutput);
            }
        }
        else if(nReadBytes != m_nBufferFrame * m_nBlockAlign){
            //終了済みの場合はバッファが空になるまで待つ
            if(m_nFinished == 1){
                m_dwTime = timeGetTime();
                m_nFinished = 2;
#if 1 && defined(_DEBUG)
                wchar_t szTime[256];
                wsprintf(szTime, L"m_nFinished=2\n", m_dwTime);
                OutputDebugString(szTime);
#endif
            }
            SetEvent(m_hEventCanOutput);
        }
        Unlock();
    }
EXIT:
    if(hTask){
        _AvRevertMmThreadCharacteristics(hTask);
    }
}
///////////////////////////////////////////////////////////////////////////////
void WINAPI KbWasapiDevice::Wait(void)
{//次の書き込みが可能になるまで待つ
    if(!m_hEventCanOutput){
        return;
    }
    HANDLE hEvents[] = {m_hEventCanOutput, m_hThread};
    WaitForMultipleObjects(2, hEvents, FALSE, INFINITE);
    while(m_nFinished > 1){
        Sleep(1);
        if(GetPosition() == m_qwTotalWritten){
            break;
        }
    }
}
///////////////////////////////////////////////////////////////////////////////
BOOL WINAPI KbWasapiDevice::Output(BYTE *pOutput, int nSizeSample)
{
    Lock();
    m_RingBuffer.Write(pOutput, nSizeSample * m_nBlockAlign);
    m_qwTotalWritten += nSizeSample;
    if(!m_hThread){
        int nBufferFrame = IsEventMode() ? m_nBufferFrame : m_nBufferSizeSample;
        if(nSizeSample == m_nBufferSizeSample){
            if(m_RingBuffer.GetWritten() < nBufferFrame * m_nBlockAlign){
                SetEvent(m_hEventCanOutput);
                Unlock();
                return TRUE;
            }
        }
        BYTE *pData = NULL;
        if(!FAILED(m_pRenderClient->GetBuffer(nBufferFrame, &pData))){
            m_RingBuffer.Read(pData, nBufferFrame * m_nBlockAlign);
            m_qwProcessedMax = nBufferFrame;
            m_pRenderClient->ReleaseBuffer(nBufferFrame, 0);
            m_dwTime = timeGetTime();
        }
        if(!m_hEventCanOutput){
            m_hEventCanOutput = CreateEvent(NULL, TRUE, FALSE, NULL);
        }
        else{
            ResetEvent(m_hEventCanOutput);
        }
        UINT thread_id = 0;
        m_hThread = (HANDLE)_beginthreadex(NULL, 0, ThreadProc, this, CREATE_SUSPENDED, &thread_id);
        SetThreadPriority(m_hThread, THREAD_PRIORITY_ABOVE_NORMAL);
        ResumeThread(m_hThread);
        //m_pAudioClient->Start();//ThreadProc 内で呼び出す
    }
    else{
        ResetEvent(m_hEventCanOutput);
    }
    if(nSizeSample < m_nBufferSizeSample && !m_nFinished){
        m_nFinished = 1;
    }
    else{
        if(m_RingBuffer.GetWritten() < m_nBufferFrame * m_nBlockAlign){
            SetEvent(m_hEventCanOutput);
        }
    }
    Unlock();
    return TRUE;
}
///////////////////////////////////////////////////////////////////////////////
void WINAPI KbWasapiDevice::Stop(void)
{
    if(m_hThread){
        SetEvent(m_hEventStop);
        WaitForSingleObject(m_hThread, INFINITE);
        CloseHandle(m_hThread);
        ResetEvent(m_hEventStop);
        m_hThread = NULL;
        m_pAudioClient->Stop();
        m_pAudioClient->Reset();
        m_RingBuffer.Reset();
        m_qwProcessedSample = 0;
        m_qwProcessedMax = 0;
        m_dwTime = 0;
        m_nFinished = 0;
        m_qwPrevPosition = 0;
        m_qwTotalWritten = 0;
    }
}
///////////////////////////////////////////////////////////////////////////////
void WINAPI KbWasapiDevice::Pause(void)
{
    Lock();
    m_pAudioClient->Stop();
    m_dwTime = 0;
    m_qwPrevPosition = 0;
    Unlock();
}
///////////////////////////////////////////////////////////////////////////////
void WINAPI KbWasapiDevice::Restart(void)
{
    m_pAudioClient->Start();
}
///////////////////////////////////////////////////////////////////////////////
UINT64 WINAPI KbWasapiDevice::GetPosition(void)
{//デバイスの再生位置をサンプル単位で返す
    UINT64 ret;
    Lock();
    if(m_pAudioClock){
        UINT64 pos = 0;
        //UINT64 pos2  = 0;
        m_pAudioClock->GetPosition(&pos, NULL);//&pos2);
        if(pos > m_qwTotalWritten){
            pos = m_qwTotalWritten;
        }
        ret = pos;
    }
    else{
        ret = m_qwProcessedSample;
        if(m_dwTime){
            DWORD dwMs = timeGetTime() - m_dwTime;//m_qwProcessedSample が更新されてからの経過時間
            UINT64 d = (UINT64)dwMs * m_nSampleRate / 1000;//↑のサンプル単位
            ret += d;
            if(ret > m_qwLatencySample){
                ret -= m_qwLatencySample;
            }
            else{
                ret = 0;
            }
            if(ret > m_qwProcessedMax){
                ret = m_qwProcessedMax;
            }
            if(ret < m_qwPrevPosition){
                ret = m_qwPrevPosition;
            }
            if(ret > m_qwTotalWritten){
                ret = m_qwTotalWritten;
            }
            m_qwPrevPosition = ret;
        }
    }
    Unlock();
    return ret;
}
///////////////////////////////////////////////////////////////////////////////
BOOL WINAPI KbWasapiDevice::SetVolume(int nVolume)
{
    return false;
}
///////////////////////////////////////////////////////////////////////////////
BOOL WINAPI KbWasapiDevice::SetRate(int nRate)
{
    return false;
}
///////////////////////////////////////////////////////////////////////////////
BOOL __fastcall KbWasapiDevice::Open(const wchar_t *cszDeviceName, //デバイス名
                                     const wchar_t *cszId,         //識別文字列
                                     const KPI_DEVICE_FORMAT *cpFormat, //再生周波数等
                                     int nBufferSizeSample,        //再生バッファサイズ（サンプル単位）
                                     IKpiConfig *pConfig)          //プラグイン独自の設定
{
    KbWasapiEnumerator wasapiEnumerator(pConfig);
    IMMDevice *pDevice = NULL;
    IAudioClient *pAudioClient = NULL;
    IAudioRenderClient *pRenderClient = NULL;
    wasapiEnumerator.GetMmDevice(cszDeviceName, cszId, &pDevice, &m_nMode);
    if(!pDevice){
        return FALSE;
    }
    HRESULT hr = pDevice->Activate(__uuidof(IAudioClient), CLSCTX_INPROC_SERVER,
                                   NULL, (void**)&pAudioClient);
    if(FAILED(hr)){
        pDevice->Release();
        return FALSE;
    }
    int nSampleRate = cpFormat->dwSampleRate;
    int nChannels = cpFormat->dwChannels;
    int nBps = cpFormat->nBitsPerSample;
    WAVEFORMATEXTENSIBLE wfx = {0};
    wfx.Format.wFormatTag      = WAVE_FORMAT_EXTENSIBLE;
    wfx.Format.nChannels       = nChannels;
    wfx.Format.nSamplesPerSec  = nSampleRate;
    wfx.Format.nBlockAlign     = (abs(nBps) / 8) * nChannels;
    wfx.Format.nAvgBytesPerSec = nSampleRate * wfx.Format.nBlockAlign;
    wfx.Format.wBitsPerSample  = abs(nBps);
    wfx.Format.cbSize          = 22;
    wfx.dwChannelMask = SPEAKER_FRONT_LEFT|SPEAKER_FRONT_RIGHT;
    wfx.Samples.wValidBitsPerSample = abs(nBps);
    if(nBps < 0){
        wfx.SubFormat = KSDATAFORMAT_SUBTYPE_IEEE_FLOAT;
    }
    else{
        wfx.SubFormat = KSDATAFORMAT_SUBTYPE_PCM;
    }
    hr = pAudioClient->IsFormatSupported(AUDCLNT_SHAREMODE_EXCLUSIVE, (const WAVEFORMATEX*)&wfx, NULL);
    if(hr == AUDCLNT_E_UNSUPPORTED_FORMAT){
        if(nBps == 32){
            wfx.Samples.wValidBitsPerSample = 24;//入力は32bitだが実質的な出力は24bit
            hr = pAudioClient->IsFormatSupported(AUDCLNT_SHAREMODE_EXCLUSIVE, (const WAVEFORMATEX*)&wfx, NULL);
        }
    }
    if(FAILED(hr)){
        pAudioClient->Release();
        pDevice->Release();
        return FALSE;
    }
    UINT uBufferFrame = 0;
    REFERENCE_TIME devicePeriod  = 0, minPeriod = 0;
    pAudioClient->GetDevicePeriod(&devicePeriod, &minPeriod);
    
    REFERENCE_TIME bufferDuration = devicePeriod;
    DWORD dwStreamFlags = AUDCLNT_STREAMFLAGS_NOPERSIST;
    if(IsEventMode()){
        dwStreamFlags |= AUDCLNT_STREAMFLAGS_EVENTCALLBACK;
    }
    {
        KbWasapiConfig cfg(pConfig, cszDeviceName, cszId);
        const REFERENCE_TIME rtDuration = cfg.GetBufferDuration() * 10000;//30 * 10000;
        bufferDuration = rtDuration;
    }
    //{
    //    REFERENCE_TIME rtBufSize = (_int64)nBufferSizeSample * 1000*10000 / nSampleRate;
    //    bufferDuration = rtBufSize;
    //}
    hr = pAudioClient->Initialize(AUDCLNT_SHAREMODE_EXCLUSIVE, 
                                  dwStreamFlags, 
                                  bufferDuration, 
                                  bufferDuration, 
                                  (const WAVEFORMATEX*)&wfx, 
                                  NULL);
    if(hr == AUDCLNT_E_BUFFER_SIZE_NOT_ALIGNED){
        pAudioClient->GetBufferSize(&uBufferFrame);
        bufferDuration = (REFERENCE_TIME)((10000.0 * 1000 / nSampleRate * uBufferFrame) + 0.5);
        pAudioClient->Release();
        hr = pDevice->Activate(__uuidof(IAudioClient), CLSCTX_INPROC_SERVER,
                                           NULL, (void**)&pAudioClient);
        if(FAILED(hr)){
            pDevice->Release();
            return FALSE;
        }
        dwStreamFlags |= AUDCLNT_STREAMFLAGS_RATEADJUST;
        hr = pAudioClient->Initialize(AUDCLNT_SHAREMODE_EXCLUSIVE, 
                                      dwStreamFlags, 
                                      bufferDuration, 
                                      bufferDuration, 
                                      (const WAVEFORMATEX*)&wfx, 
                                      NULL);
    }
    if(FAILED(hr)){
        pAudioClient->Release();
        pDevice->Release();
        return FALSE;
    }
    pAudioClient->GetBufferSize(&uBufferFrame);
    hr = pAudioClient->GetService(
                         __uuidof(IAudioRenderClient),
                         (void**)&pRenderClient);
    if(FAILED(hr)){
        pAudioClient->Release();
        pDevice->Release();
        return FALSE;
    }
    REFERENCE_TIME Latency = 0;
    pAudioClient->GetStreamLatency(&Latency);
    m_qwLatencySample = Latency * nSampleRate / (1000 * 10000);
    //成功
    m_hEventStop = CreateEvent(NULL, TRUE, FALSE, NULL);
    
    m_pAudioClient = pAudioClient;
    m_pRenderClient = pRenderClient;
    m_pDevice = pDevice;

    m_nSampleRate = nSampleRate;
    m_nChannels = nChannels;
    m_nBps = nBps;
    m_nBlockAlign = nChannels * abs(nBps) / 8;
    m_nBufferSizeSample = nBufferSizeSample;
    m_nBufferFrame = uBufferFrame;
    int nSizeSample = m_nBufferSizeSample;
    if(nSizeSample < m_nBufferFrame){
        nSizeSample = m_nBufferFrame;
    }
    m_RingBuffer.SetSize(nSizeSample * m_nBlockAlign * 2);
    
    if(IsEventMode()){
        m_hEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
        pAudioClient->SetEventHandle(m_hEvent);
    }
    pAudioClient->GetService(_uuidof(IAudioClock), (void**)&m_pAudioClock);
    return TRUE;
}
