#include "backends/wasapi_format_probe.h"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <iterator>
#include <sstream>
#include <utility>

#include "core/dop_generator.h"

#if defined(_WIN32)
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#include <audioclient.h>
#include <initguid.h>
#include <mmdeviceapi.h>
#include <mmreg.h>
#endif

namespace aptt {
namespace {

constexpr std::size_t kMaxCandidateFormats = 128;
constexpr std::int32_t kHresultInvalidArg =
    static_cast<std::int32_t>(0x80070057u);
constexpr std::int32_t kHresultNotImplemented =
    static_cast<std::int32_t>(0x80004001u);

#if defined(_WIN32)

constexpr CLSID kClsidMmDeviceEnumerator = {
    0xbcde0395,
    0xe52f,
    0x467c,
    {0x8e, 0x3d, 0xc4, 0x57, 0x92, 0x91, 0x69, 0x2e},
};

constexpr GUID kWaveSubtypePcm = {
    0x00000001,
    0x0000,
    0x0010,
    {0x80, 0x00, 0x00, 0xaa, 0x00, 0x38, 0x9b, 0x71},
};

constexpr GUID kWaveSubtypeIeeeFloat = {
    0x00000003,
    0x0000,
    0x0010,
    {0x80, 0x00, 0x00, 0xaa, 0x00, 0x38, 0x9b, 0x71},
};

std::int64_t HResultCode(HRESULT result) {
  return static_cast<std::int32_t>(result);
}

template <typename T>
class ComPtr {
 public:
  ComPtr() = default;
  ComPtr(const ComPtr&) = delete;
  ComPtr& operator=(const ComPtr&) = delete;

  ~ComPtr() {
    if (ptr_ != nullptr) {
      ptr_->Release();
    }
  }

  T** put() { return &ptr_; }
  T* get() const { return ptr_; }
  T* operator->() const { return ptr_; }
  explicit operator bool() const { return ptr_ != nullptr; }

 private:
  T* ptr_ = nullptr;
};

class ScopedCom {
 public:
  ScopedCom() {
    result_ = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
    if (result_ == S_OK || result_ == S_FALSE) {
      needs_uninitialize_ = true;
    } else if (result_ == RPC_E_CHANGED_MODE) {
      result_ = S_OK;
      needs_uninitialize_ = false;
    }
  }

  ScopedCom(const ScopedCom&) = delete;
  ScopedCom& operator=(const ScopedCom&) = delete;

  ~ScopedCom() {
    if (needs_uninitialize_) {
      CoUninitialize();
    }
  }

  HRESULT result() const { return result_; }

 private:
  HRESULT result_ = E_FAIL;
  bool needs_uninitialize_ = false;
};

std::wstring Utf8ToWide(const std::string& value) {
  if (value.empty()) {
    return {};
  }
  const int required = MultiByteToWideChar(
      CP_UTF8, 0, value.c_str(), -1, nullptr, 0);
  if (required <= 1) {
    return {};
  }
  std::wstring converted(static_cast<std::size_t>(required - 1), L'\0');
  MultiByteToWideChar(
      CP_UTF8, 0, value.c_str(), -1, converted.data(), required);
  return converted;
}

SampleType SampleTypeFromWaveFormat(const WAVEFORMATEX& format) {
  if (format.wFormatTag == WAVE_FORMAT_IEEE_FLOAT) {
    return SampleType::kFloat32;
  }
  if (format.wFormatTag == WAVE_FORMAT_PCM && format.wBitsPerSample == 16) {
    return SampleType::kInt16;
  }
  if (format.wFormatTag == WAVE_FORMAT_PCM && format.wBitsPerSample == 24) {
    return SampleType::kInt24Packed;
  }
  if (format.wFormatTag == WAVE_FORMAT_PCM && format.wBitsPerSample == 32) {
    return SampleType::kInt32;
  }
  if (format.wFormatTag == WAVE_FORMAT_EXTENSIBLE &&
      format.cbSize >= sizeof(WAVEFORMATEXTENSIBLE) - sizeof(WAVEFORMATEX)) {
    const auto& extensible =
        reinterpret_cast<const WAVEFORMATEXTENSIBLE&>(format);
    if (IsEqualGUID(extensible.SubFormat, kWaveSubtypeIeeeFloat)) {
      return SampleType::kFloat32;
    }
    if (IsEqualGUID(extensible.SubFormat, kWaveSubtypePcm)) {
      if (extensible.Samples.wValidBitsPerSample == 16) {
        return SampleType::kInt16;
      }
      if (extensible.Samples.wValidBitsPerSample == 24 &&
          format.wBitsPerSample == 32) {
        return SampleType::kInt24In32;
      }
      if (extensible.Samples.wValidBitsPerSample == 24) {
        return SampleType::kInt24Packed;
      }
      if (extensible.Samples.wValidBitsPerSample == 32) {
        return SampleType::kInt32;
      }
    }
  }
  return SampleType::kFloat32;
}

FormatDescriptor FormatDescriptorFromWaveFormat(const WAVEFORMATEX& format) {
  FormatDescriptor descriptor;
  descriptor.signal_mode = SignalMode::kPcm;
  descriptor.sample_rate_hz = format.nSamplesPerSec;
  descriptor.sample_type = SampleTypeFromWaveFormat(format);
  descriptor.valid_bits_per_sample = format.wBitsPerSample;
  descriptor.container_bits_per_sample = format.wBitsPerSample;
  descriptor.channels.channel_count = format.nChannels;

  if (format.wFormatTag == WAVE_FORMAT_EXTENSIBLE &&
      format.cbSize >= sizeof(WAVEFORMATEXTENSIBLE) - sizeof(WAVEFORMATEX)) {
    const auto& extensible =
        reinterpret_cast<const WAVEFORMATEXTENSIBLE&>(format);
    descriptor.valid_bits_per_sample =
        extensible.Samples.wValidBitsPerSample;
    descriptor.container_bits_per_sample = format.wBitsPerSample;
    descriptor.channels.channel_mask = extensible.dwChannelMask;
  }
  return descriptor;
}

WAVEFORMATEXTENSIBLE WaveFormatFromDescriptor(const FormatDescriptor& descriptor) {
  WAVEFORMATEXTENSIBLE wave{};
  wave.Format.wFormatTag = WAVE_FORMAT_EXTENSIBLE;
  wave.Format.nChannels = static_cast<WORD>(descriptor.channels.channel_count);
  wave.Format.nSamplesPerSec = descriptor.sample_rate_hz;
  wave.Format.wBitsPerSample =
      static_cast<WORD>(descriptor.container_bits_per_sample);
  wave.Format.nBlockAlign =
      static_cast<WORD>((wave.Format.nChannels * wave.Format.wBitsPerSample) / 8);
  wave.Format.nAvgBytesPerSec =
      wave.Format.nSamplesPerSec * wave.Format.nBlockAlign;
  wave.Format.cbSize = sizeof(WAVEFORMATEXTENSIBLE) - sizeof(WAVEFORMATEX);
  wave.Samples.wValidBitsPerSample =
      static_cast<WORD>(descriptor.valid_bits_per_sample);
  wave.dwChannelMask = descriptor.channels.channel_mask;
  if (descriptor.sample_type == SampleType::kFloat32 ||
      descriptor.sample_type == SampleType::kFloat64) {
    wave.SubFormat = kWaveSubtypeIeeeFloat;
  } else {
    wave.SubFormat = kWaveSubtypePcm;
  }
  return wave;
}

bool SameFormatDescriptor(const FormatDescriptor& left,
                          const FormatDescriptor& right) {
  return left.signal_mode == right.signal_mode &&
         left.sample_rate_hz == right.sample_rate_hz &&
         left.sample_type == right.sample_type &&
         left.valid_bits_per_sample == right.valid_bits_per_sample &&
         left.container_bits_per_sample == right.container_bits_per_sample &&
         left.channels.channel_count == right.channels.channel_count &&
         left.channels.channel_mask == right.channels.channel_mask &&
         left.dsd_rate_multiplier == right.dsd_rate_multiplier;
}

void PushUniqueFormat(std::vector<FormatDescriptor>& formats,
                      const FormatDescriptor& candidate) {
  const auto found = std::find_if(
      formats.begin(),
      formats.end(),
      [&](const FormatDescriptor& existing) {
        return SameFormatDescriptor(existing, candidate);
      });
  if (found == formats.end()) {
    formats.push_back(candidate);
  }
}

#endif

#if defined(_WIN32)

CapabilityProbeRecord MakeFormatProbeRecord(
    const FormatDescriptor& requested,
    HRESULT result,
    const WAVEFORMATEX* closest_match,
    bool exclusive_mode,
    std::string candidate_label) {
  CapabilityProbeRecord probe;
  probe.kind = CapabilityProbeKind::kFormatSupport;
  probe.requested_format = requested;
  probe.backend_fields["share_mode"] = exclusive_mode ? "exclusive" : "shared";
  probe.backend_fields["api"] = "IAudioClient::IsFormatSupported";
  probe.backend_fields["candidate_label"] = std::move(candidate_label);

  const auto classified = ClassifyWasapiFormatSupportResult(
      HResultCode(result),
      closest_match != nullptr);
  if (classified == CapabilitySupport::kSupported) {
    probe.support = CapabilitySupport::kSupported;
    probe.accepted_format = requested;
    probe.detail = "format supported";
    return probe;
  }

  if (classified == CapabilitySupport::kClosestMatch &&
      closest_match != nullptr) {
    probe.support = CapabilitySupport::kClosestMatch;
    probe.accepted_format = FormatDescriptorFromWaveFormat(*closest_match);
    probe.detail = "closest shared-mode match returned";
    return probe;
  }

  if (classified == CapabilitySupport::kUnsupported) {
    probe.support = CapabilitySupport::kUnsupported;
    probe.detail = "format unsupported";
    return probe;
  }

  probe.support = CapabilitySupport::kFailed;
  probe.error = MakeApiError(
      HResultCode(result),
      "IAudioClient::IsFormatSupported",
      "format support query failed");
  probe.detail = "format support query failed";
  return probe;
}

void AddDevicePeriodProbe(IAudioClient* audio_client,
                          DeviceCapabilitySnapshot& snapshot) {
  REFERENCE_TIME default_period = 0;
  REFERENCE_TIME minimum_period = 0;
  const HRESULT result =
      audio_client->GetDevicePeriod(&default_period, &minimum_period);

  CapabilityProbeRecord probe;
  probe.kind = CapabilityProbeKind::kLatency;
  probe.backend_fields["api"] = "IAudioClient::GetDevicePeriod";
  if (SUCCEEDED(result)) {
    probe.support = CapabilitySupport::kSupported;
    probe.detail = "device period queried";
    probe.backend_fields["default_period_100ns"] = std::to_string(default_period);
    probe.backend_fields["minimum_period_100ns"] = std::to_string(minimum_period);
  } else {
    probe.support = CapabilitySupport::kFailed;
    probe.detail = "device period query failed";
    probe.error = MakeApiError(
        HResultCode(result),
        "IAudioClient::GetDevicePeriod",
        "failed to query device period");
  }
  snapshot.probes.push_back(std::move(probe));
}

void AddMixFormatProbe(IAudioClient* audio_client,
                       DeviceCapabilitySnapshot& snapshot,
                       std::optional<FormatDescriptor>& mix_format) {
  WAVEFORMATEX* raw_mix_format = nullptr;
  const HRESULT result = audio_client->GetMixFormat(&raw_mix_format);

  CapabilityProbeRecord probe;
  probe.kind = CapabilityProbeKind::kFormatSupport;
  probe.backend_fields["api"] = "IAudioClient::GetMixFormat";
  if (SUCCEEDED(result) && raw_mix_format != nullptr) {
    mix_format = FormatDescriptorFromWaveFormat(*raw_mix_format);
    probe.support = CapabilitySupport::kSupported;
    probe.accepted_format = mix_format;
    probe.detail = "mix format reference; not automatically selected";
    probe.backend_fields["candidate_label"] = "mix_format_reference";
    CoTaskMemFree(raw_mix_format);
  } else {
    probe.support = CapabilitySupport::kFailed;
    probe.detail = "mix format query failed";
    probe.error = MakeApiError(
        HResultCode(result),
        "IAudioClient::GetMixFormat",
        "failed to query mix format");
  }
  snapshot.probes.push_back(std::move(probe));
}

#endif

DeviceCapabilitySnapshot MakeProbeFailure(BackendApi backend,
                                          std::string endpoint_id,
                                          ApiErrorInfo error) {
  DeviceInfo device;
  device.backend = backend;
  device.display_name = "Selected WASAPI endpoint";
  device.stable_id = std::move(endpoint_id);
  device.driver_name = "WASAPI/Core Audio";

  auto snapshot = MakeDiscoveredDeviceSnapshot(
      std::move(device),
      DeviceAvailability::kUnavailable);
  snapshot.discovery_error = std::move(error);
  return snapshot;
}

}  // namespace

CapabilitySupport ClassifyWasapiFormatSupportResult(
    std::int32_t raw_hresult,
    bool closest_match_returned) {
  const auto code = static_cast<std::uint32_t>(raw_hresult);
  if (raw_hresult == 0) {
    return CapabilitySupport::kSupported;
  }
  if (raw_hresult == 1 && closest_match_returned) {
    return CapabilitySupport::kClosestMatch;
  }
  if (code == 0x88890008u) {
    return CapabilitySupport::kUnsupported;
  }
  if (raw_hresult < 0) {
    return CapabilitySupport::kFailed;
  }
  return CapabilitySupport::kUnsupported;
}

FormatDescriptor WasapiDefaultPcmCandidate() {
  return WasapiPcmCandidate(48000, SampleType::kFloat32, 2);
}

std::vector<FormatDescriptor> WasapiDefaultPcmCandidates() {
  constexpr std::uint32_t kSampleRatesHz[] = {
      44100, 48000, 88200, 96000, 176400, 192000, 352800,
      384000, 705600, 768000, 786000, 1411200, 1536000};
  constexpr SampleType kSampleTypes[] = {
      SampleType::kInt16,
      SampleType::kInt24Packed,
      SampleType::kInt24In32,
      SampleType::kFloat32,
      SampleType::kInt32,
  };

  std::vector<FormatDescriptor> candidates;
  candidates.reserve(std::size(kSampleRatesHz) * std::size(kSampleTypes));
  for (const auto sample_rate_hz : kSampleRatesHz) {
    for (const auto sample_type : kSampleTypes) {
      candidates.push_back(WasapiPcmCandidate(sample_rate_hz, sample_type, 2));
    }
  }
  return candidates;
}

FormatDescriptor WasapiPcmCandidate(std::uint32_t sample_rate_hz,
                                    SampleType sample_type,
                                    std::uint32_t channel_count) {
  FormatDescriptor descriptor;
  descriptor.signal_mode = SignalMode::kPcm;
  descriptor.sample_rate_hz = sample_rate_hz;
  descriptor.sample_type = sample_type;
  descriptor.channels.channel_count = channel_count;
  descriptor.channels.channel_mask = channel_count == 2 ? 3 : 0;
  switch (sample_type) {
    case SampleType::kInt16:
      descriptor.valid_bits_per_sample = 16;
      descriptor.container_bits_per_sample = 16;
      break;
    case SampleType::kInt24Packed:
      descriptor.valid_bits_per_sample = 24;
      descriptor.container_bits_per_sample = 24;
      break;
    case SampleType::kInt24In32:
      descriptor.valid_bits_per_sample = 24;
      descriptor.container_bits_per_sample = 32;
      break;
    case SampleType::kInt32:
      descriptor.valid_bits_per_sample = 32;
      descriptor.container_bits_per_sample = 32;
      break;
    case SampleType::kFloat64:
      descriptor.valid_bits_per_sample = 64;
      descriptor.container_bits_per_sample = 64;
      break;
    default:
      descriptor.sample_type = SampleType::kFloat32;
      descriptor.valid_bits_per_sample = 32;
      descriptor.container_bits_per_sample = 32;
      break;
  }
  return descriptor;
}

FormatDescriptor WasapiDopCarrierCandidate(std::uint32_t dsd_rate_multiplier,
                                           std::uint32_t channel_count) {
  DopConfig config;
  config.dsd_rate_multiplier = dsd_rate_multiplier;
  config.channel_count = channel_count;
  return MakeDopCarrierFormat(config);
}

std::vector<FormatDescriptor> WasapiDefaultDopCarrierCandidates(
    std::uint32_t channel_count) {
  constexpr std::uint32_t kDsdRates[] = {64, 128, 256, 512, 1024, 2048};
  std::vector<FormatDescriptor> candidates;
  candidates.reserve(std::size(kDsdRates));
  for (const auto dsd_rate_multiplier : kDsdRates) {
    candidates.push_back(
        WasapiDopCarrierCandidate(dsd_rate_multiplier, channel_count));
  }
  return candidates;
}

DeviceCapabilitySnapshot ProbeWasapiEndpointFormats(
    const WasapiFormatProbeRequest& request) {
  if (request.endpoint_id.empty()) {
    return MakeProbeFailure(
        request.exclusive_mode
            ? BackendApi::kWasapiExclusive
            : BackendApi::kWasapiShared,
        {},
        MakeApiError(
            kHresultInvalidArg,
            "E_INVALIDARG",
            "WASAPI format probe requires an explicit endpoint ID"));
  }

#if !defined(_WIN32)
  return MakeProbeFailure(
      request.exclusive_mode
          ? BackendApi::kWasapiExclusive
          : BackendApi::kWasapiShared,
      request.endpoint_id,
      MakeApiError(
          kHresultNotImplemented,
          "E_NOTIMPL",
          "WASAPI format probe is available only on Windows"));
#else
  ScopedCom com;
  if (FAILED(com.result())) {
    return MakeProbeFailure(
        request.exclusive_mode
            ? BackendApi::kWasapiExclusive
            : BackendApi::kWasapiShared,
        request.endpoint_id,
        MakeApiError(
            HResultCode(com.result()),
            "CoInitializeEx",
            "failed to initialize COM for WASAPI format probe"));
  }

  ComPtr<IMMDeviceEnumerator> enumerator;
  HRESULT result = CoCreateInstance(
      kClsidMmDeviceEnumerator,
      nullptr,
      CLSCTX_ALL,
      IID_PPV_ARGS(enumerator.put()));
  if (FAILED(result) || !enumerator) {
    return MakeProbeFailure(
        request.exclusive_mode
            ? BackendApi::kWasapiExclusive
            : BackendApi::kWasapiShared,
        request.endpoint_id,
        MakeApiError(
            HResultCode(result),
            "CoCreateInstance(CLSID_MMDeviceEnumerator)",
            "failed to create MMDeviceEnumerator"));
  }

  const std::wstring endpoint_id = Utf8ToWide(request.endpoint_id);
  ComPtr<IMMDevice> endpoint;
  result = enumerator->GetDevice(endpoint_id.c_str(), endpoint.put());
  if (FAILED(result) || !endpoint) {
    return MakeProbeFailure(
        request.exclusive_mode
            ? BackendApi::kWasapiExclusive
            : BackendApi::kWasapiShared,
        request.endpoint_id,
        MakeApiError(
            HResultCode(result),
            "IMMDeviceEnumerator::GetDevice",
            "failed to open selected endpoint by ID"));
  }

  ComPtr<IAudioClient> audio_client;
  result = endpoint->Activate(
      __uuidof(IAudioClient),
      CLSCTX_ALL,
      nullptr,
      reinterpret_cast<void**>(audio_client.put()));
  if (FAILED(result) || !audio_client) {
    return MakeProbeFailure(
        request.exclusive_mode
            ? BackendApi::kWasapiExclusive
            : BackendApi::kWasapiShared,
        request.endpoint_id,
        MakeApiError(
            HResultCode(result),
            "IMMDevice::Activate(IAudioClient)",
            "failed to activate audio client for format probe"));
  }

  DeviceInfo device;
  device.backend = request.exclusive_mode
      ? BackendApi::kWasapiExclusive
      : BackendApi::kWasapiShared;
  device.display_name = "Selected WASAPI endpoint";
  device.stable_id = request.endpoint_id;
  device.driver_name = "WASAPI/Core Audio";

  auto snapshot = MakeDiscoveredDeviceSnapshot(
      std::move(device),
      DeviceAvailability::kActive);
  snapshot.backend_fields["api_device_id"] = request.endpoint_id;
  snapshot.backend_fields["id_kind"] = "wasapi-endpoint-id";
  snapshot.backend_fields["share_mode"] =
      request.exclusive_mode ? "exclusive" : "shared";
  snapshot.backend_fields["opens_audio_device"] = "false";
  snapshot.backend_fields["initializes_stream"] = "false";
  snapshot.backend_fields["starts_stream"] = "false";

  AddDevicePeriodProbe(audio_client.get(), snapshot);

  std::optional<FormatDescriptor> mix_format;
  AddMixFormatProbe(audio_client.get(), snapshot, mix_format);

  std::vector<std::pair<FormatDescriptor, std::string>> candidates;
  candidates.reserve(kMaxCandidateFormats);
  for (const auto& candidate : request.candidate_formats) {
    if (candidates.size() >= kMaxCandidateFormats) {
      break;
    }
    const bool duplicate = std::find_if(
        candidates.begin(),
        candidates.end(),
        [&](const auto& existing) {
          return SameFormatDescriptor(existing.first, candidate);
        }) != candidates.end();
    if (!duplicate) {
      candidates.emplace_back(
          candidate,
          candidates.empty() ? "gui_requested" : "fallback_candidate");
    }
  }
  if (mix_format && request.include_mix_format_candidate &&
      candidates.size() < kMaxCandidateFormats) {
    const bool duplicate = std::find_if(
        candidates.begin(),
        candidates.end(),
        [&](const auto& existing) {
          return SameFormatDescriptor(existing.first, *mix_format);
        }) != candidates.end();
    if (!duplicate) {
      candidates.emplace_back(*mix_format, "mix_format_fallback");
    }
  }
  if (candidates.empty()) {
    for (const auto& candidate : WasapiDefaultPcmCandidates()) {
      if (candidates.size() >= kMaxCandidateFormats) {
        break;
      }
      candidates.emplace_back(candidate, "default_fallback");
    }
  }

  for (const auto& [candidate, candidate_label] : candidates) {
    WAVEFORMATEXTENSIBLE wave = WaveFormatFromDescriptor(candidate);
    WAVEFORMATEX* closest_match = nullptr;
    WAVEFORMATEX** closest_match_out =
        request.exclusive_mode ? nullptr : &closest_match;
    result = audio_client->IsFormatSupported(
        request.exclusive_mode ? AUDCLNT_SHAREMODE_EXCLUSIVE
                               : AUDCLNT_SHAREMODE_SHARED,
        &wave.Format,
        closest_match_out);

    snapshot.probes.push_back(
        MakeFormatProbeRecord(candidate,
                              result,
                              closest_match,
                              request.exclusive_mode,
                              candidate_label));
    if (closest_match != nullptr) {
      CoTaskMemFree(closest_match);
    }
  }

  return snapshot;
#endif
}

}  // namespace aptt
