#define NOMINMAX
#include <windows.h>

#include <audioclient.h>
#include <audioclientactivationparams.h>
#include <fcntl.h>
#include <io.h>
#include <mmdeviceapi.h>
#include <wrl.h>
#include <wrl/implements.h>

#include "capture_target.h"
#include "process_audio.h"

#include <array>
#include <cstdint>
#include <iostream>
#include <limits>
#include <string>
#include <vector>

using Microsoft::WRL::ComPtr;
using Microsoft::WRL::FtmBase;
using Microsoft::WRL::Make;
using Microsoft::WRL::RuntimeClass;
using Microsoft::WRL::RuntimeClassFlags;
using Microsoft::WRL::ClassicCom;

namespace piik::capture {
namespace {
class ActivationHandler final : public RuntimeClass<
    RuntimeClassFlags<ClassicCom>, FtmBase, IActivateAudioInterfaceCompletionHandler> {
 public:
  ActivationHandler() : completed_(CreateEventW(nullptr, FALSE, FALSE, nullptr)) {}
  ~ActivationHandler() { if (completed_) CloseHandle(completed_); }

  STDMETHODIMP ActivateCompleted(IActivateAudioInterfaceAsyncOperation* operation) override {
    ComPtr<IUnknown> activated;
    HRESULT activationResult = E_UNEXPECTED;
    result_ = operation->GetActivateResult(&activationResult, &activated);
    if (SUCCEEDED(result_)) result_ = activationResult;
    if (SUCCEEDED(result_)) result_ = activated.As(&client_);
    SetEvent(completed_);
    return S_OK;
  }

  HRESULT Result() const { return result_; }
  ComPtr<IAudioClient> Client() const { return client_; }
  HANDLE Completed() const { return completed_; }

 private:
  HANDLE completed_ = nullptr;
  HRESULT result_ = E_UNEXPECTED;
  ComPtr<IAudioClient> client_;
};

HRESULT ActivateProcessLoopback(DWORD pid, bool exclude,
                                ComPtr<IAudioClient>* client) {
  AUDIOCLIENT_ACTIVATION_PARAMS parameters{};
  parameters.ActivationType = AUDIOCLIENT_ACTIVATION_TYPE_PROCESS_LOOPBACK;
  parameters.ProcessLoopbackParams.TargetProcessId = pid;
  parameters.ProcessLoopbackParams.ProcessLoopbackMode =
      exclude ? PROCESS_LOOPBACK_MODE_EXCLUDE_TARGET_PROCESS_TREE
              : PROCESS_LOOPBACK_MODE_INCLUDE_TARGET_PROCESS_TREE;
  PROPVARIANT variant{};
  variant.vt = VT_BLOB;
  variant.blob.cbSize = sizeof(parameters);
  variant.blob.pBlobData = reinterpret_cast<BYTE*>(&parameters);

  // The completion can outlive a timed-out caller; it owns its event handle.
  auto handler = Make<ActivationHandler>();
  if (!handler) return E_OUTOFMEMORY;
  if (!handler->Completed()) return HRESULT_FROM_WIN32(GetLastError());
  ComPtr<IActivateAudioInterfaceAsyncOperation> operation;
  HRESULT result = ActivateAudioInterfaceAsync(VIRTUAL_AUDIO_DEVICE_PROCESS_LOOPBACK,
                                                __uuidof(IAudioClient), &variant,
                                                handler.Get(), &operation);
  if (FAILED(result)) return result;
  if (WaitForSingleObject(handler->Completed(), 10'000) != WAIT_OBJECT_0) return HRESULT_FROM_WIN32(WAIT_TIMEOUT);
  result = handler->Result();
  if (SUCCEEDED(result)) *client = handler->Client();
  return result;
}

HRESULT ActivateEndpoint(ComPtr<IAudioClient>* client, bool microphone = false, const std::wstring& device_id = L"") {
  ComPtr<IMMDeviceEnumerator> enumerator;
  HRESULT result = CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr,
                                    CLSCTX_ALL, IID_PPV_ARGS(&enumerator));
  ComPtr<IMMDevice> device;
  if (SUCCEEDED(result)) {
    result = device_id.empty()
        ? enumerator->GetDefaultAudioEndpoint(microphone ? eCapture : eRender, microphone ? eCommunications : eConsole, &device)
        : enumerator->GetDevice(device_id.c_str(), &device);
    if (SUCCEEDED(result) && microphone) {
      ComPtr<IMMEndpoint> endpoint;
      EDataFlow flow = eAll;
      result = device.As(&endpoint);
      if (SUCCEEDED(result)) result = endpoint->GetDataFlow(&flow);
      if (SUCCEEDED(result) && flow != eCapture) result = E_INVALIDARG;
    }
  }
  if (SUCCEEDED(result)) {
    result = device->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr,
                              reinterpret_cast<void**>(client->GetAddressOf()));
  }
  return result;
}

HRESULT CaptureAudioFrames(ComPtr<IAudioClient> client, HANDLE process, bool loopback,
                             HANDLE stop_event, const StopProbe& stop_probe,
                             const ReadyWriter& ready_writer,
                             const PCMWriter& writer) {
  if (stop_event == nullptr || !stop_probe || !ready_writer || !writer) {
    return HRESULT_FROM_WIN32(ERROR_INVALID_PARAMETER);
  }
  HANDLE sampleReady = CreateEventW(nullptr, FALSE, FALSE, nullptr);
  if (sampleReady == nullptr) {
    const HRESULT result = HRESULT_FROM_WIN32(GetLastError());
    return result;
  }
  HRESULT result = S_OK;
  // Bind process audio to the selected lifetime, including while draining PCM.
  // Exclusion must never become unfiltered system capture after its target dies.
  const auto processAlive = [&]() -> HRESULT {
    if (!process) return S_OK;
    const DWORD wait = WaitForSingleObject(process, 0);
    if (wait == WAIT_TIMEOUT) return S_OK;
    return HRESULT_FROM_WIN32(wait == WAIT_OBJECT_0 ? ERROR_PROCESS_ABORTED : GetLastError());
  };
  ComPtr<IAudioCaptureClient> capture;
  WAVEFORMATEX format{};
  format.wFormatTag = WAVE_FORMAT_PCM;
  format.nChannels = kAudioChannels;
  format.nSamplesPerSec = kAudioSampleRate;
  format.wBitsPerSample = kAudioBytesPerSample * 8;
  format.nBlockAlign = kAudioChannels * kAudioBytesPerSample;
  format.nAvgBytesPerSec = kAudioSampleRate * format.nBlockAlign;
  if (SUCCEEDED(result)) {
    result = client->Initialize(AUDCLNT_SHAREMODE_SHARED,
        (loopback ? AUDCLNT_STREAMFLAGS_LOOPBACK : 0) | AUDCLNT_STREAMFLAGS_EVENTCALLBACK |
            AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM | AUDCLNT_STREAMFLAGS_SRC_DEFAULT_QUALITY,
        0, 0, &format, nullptr);
  }
  if (SUCCEEDED(result)) result = client->GetService(IID_PPV_ARGS(&capture));
  if (SUCCEEDED(result)) result = client->SetEventHandle(sampleReady);
  if (SUCCEEDED(result)) result = client->Start();
  if (SUCCEEDED(result)) result = processAlive();
  if (SUCCEEDED(result)) result = ready_writer();
  if (SUCCEEDED(result)) {
    std::array<BYTE, kAudioBytesPerChunk> silence{};
    result = writer(0, silence.data(), static_cast<DWORD>(silence.size()));
  }

  std::vector<BYTE> pending;
  size_t consumed = 0;
  UINT64 nextTimestamp = 0;
  const HANDLE process_waits[] = {process, sampleReady, stop_event};
  const HANDLE system_waits[] = {sampleReady, stop_event};
  while (SUCCEEDED(result)) {
    if (stop_probe()) {
      result = S_OK;
      break;
    }
    const DWORD wait = process != nullptr
                           ? WaitForMultipleObjects(3, process_waits, FALSE,
                                                    1'000)
                           : WaitForMultipleObjects(2, system_waits, FALSE,
                                                    1'000);
    if (process != nullptr && wait == WAIT_OBJECT_0) {
      result = HRESULT_FROM_WIN32(ERROR_PROCESS_ABORTED);
      break;
    }
    if (wait == WAIT_TIMEOUT) continue;
    const DWORD sample_index = process != nullptr ? WAIT_OBJECT_0 + 1
                                                  : WAIT_OBJECT_0;
    const DWORD stop_index = process != nullptr ? WAIT_OBJECT_0 + 2
                                                : WAIT_OBJECT_0 + 1;
    if (wait == stop_index) {
      result = S_OK;
      break;
    }
    if (wait != sample_index) {
      result = HRESULT_FROM_WIN32(GetLastError());
      break;
    }
    UINT32 frames = 0;
    while (SUCCEEDED(result = capture->GetNextPacketSize(&frames)) && frames > 0) {
      if (FAILED(result = processAlive())) break;
      BYTE* data = nullptr;
      DWORD flags = 0;
      UINT64 devicePosition = 0;
      UINT64 qpcPosition = 0;
      result = capture->GetBuffer(&data, &frames, &flags, &devicePosition, &qpcPosition);
      if (FAILED(result)) break;
      const size_t bytes = static_cast<size_t>(frames) * format.nBlockAlign;
      if ((flags & AUDCLNT_BUFFERFLAGS_TIMESTAMP_ERROR) != 0) {
        capture->ReleaseBuffer(frames);
        continue;
      }
      if ((flags & AUDCLNT_BUFFERFLAGS_DATA_DISCONTINUITY) != 0) {
        pending.clear();
        consumed = 0;
      }
      if (pending.size() == consumed) {
        pending.clear();
        consumed = 0;
        nextTimestamp = qpcPosition;
      }
      if ((flags & AUDCLNT_BUFFERFLAGS_SILENT) != 0) {
        pending.insert(pending.end(), bytes, 0);
      } else {
        pending.insert(pending.end(), data, data + bytes);
      }
      capture->ReleaseBuffer(frames);
      while (pending.size() - consumed >= kAudioBytesPerChunk) {
        if (FAILED(result = processAlive())) break;
        result = writer(nextTimestamp, pending.data() + consumed,
                        kAudioBytesPerChunk);
        if (FAILED(result)) break;
        consumed += kAudioBytesPerChunk;
        nextTimestamp += kAudioChunkDuration100ns;
      }
      if (FAILED(result)) break;
      if (consumed > kAudioBytesPerChunk * 4) {
        pending.erase(pending.begin(), pending.begin() + static_cast<ptrdiff_t>(consumed));
        consumed = 0;
      }
    }
  }
  if (client) client->Stop();
  CloseHandle(sampleReady);
  return result;
}

}  // namespace

bool ProcessAudioAvailable(bool exclude) {
  ComPtr<IAudioClient> client;
  const HRESULT result =
      ActivateProcessLoopback(GetCurrentProcessId(), exclude, &client);
  return SUCCEEDED(result) && client != nullptr;
}

bool SystemAudioAvailable() {
  ComPtr<IAudioClient> client;
  return SUCCEEDED(ActivateEndpoint(&client)) && client != nullptr;
}

HRESULT CaptureProcessAudio(DWORD pid, UINT64 expectedCreationTime, bool exclude,
                            HANDLE stop_event, const StopProbe& stop_probe,
                            const ReadyWriter& ready_writer,
                            const PCMWriter& writer) {
  HRESULT com_result = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
  if (FAILED(com_result)) return com_result;
  // Validate and watch the same handle so PID reuse cannot change the target.
  HANDLE process = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE, FALSE, pid);
  HRESULT result = process ? ValidateProcessTarget(process, expectedCreationTime)
                           : HRESULT_FROM_WIN32(GetLastError());
  ComPtr<IAudioClient> client;
  if (SUCCEEDED(result)) {
    result = ActivateProcessLoopback(pid, exclude, &client);
  }
  if (SUCCEEDED(result)) {
    result = CaptureAudioFrames(client, process, true, stop_event, stop_probe,
                                  ready_writer, writer);
  }
  if (process != nullptr) CloseHandle(process);
  client.Reset();
  CoUninitialize();
  return result;
}

HRESULT CaptureSystemAudio(HANDLE stop_event, const StopProbe& stop_probe,
                           const ReadyWriter& ready_writer,
                           const PCMWriter& writer) {
  HRESULT com_result = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
  if (FAILED(com_result)) return com_result;
  ComPtr<IAudioClient> client;
  HRESULT result = ActivateEndpoint(&client);
  if (SUCCEEDED(result)) {
    result = CaptureAudioFrames(client, nullptr, true, stop_event, stop_probe,
                                  ready_writer, writer);
  }
  client.Reset();
  CoUninitialize();
  return result;
}

HRESULT CaptureMicrophone(const std::wstring& device_id, HANDLE stop_event, const StopProbe& stop_probe,
                           const ReadyWriter& ready_writer, const PCMWriter& writer) {
  HRESULT result = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
  if (FAILED(result)) return result;
  ComPtr<IAudioClient> client;
  result = ActivateEndpoint(&client, true, device_id);
  if (SUCCEEDED(result)) {
    result = CaptureAudioFrames(client, nullptr, false, stop_event, stop_probe, ready_writer, writer);
  }
  client.Reset();
  CoUninitialize();
  return result;
}

}  // namespace piik::capture