#include "backends/wasapi_endpoint_discovery.h"

#include <array>
#include <cstdint>
#include <sstream>
#include <utility>

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

namespace aptt {
namespace {

constexpr std::int32_t kHresultNotImplemented =
    static_cast<std::int32_t>(0x80004001u);

std::string JoinRoles(const std::vector<std::string>& roles) {
  std::string joined;
  for (const auto& role : roles) {
    if (!joined.empty()) {
      joined += ",";
    }
    joined += role;
  }
  return joined;
}

#if defined(_WIN32)

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

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 ScopedPropVariant {
 public:
  ScopedPropVariant() { PropVariantInit(&value_); }
  ScopedPropVariant(const ScopedPropVariant&) = delete;
  ScopedPropVariant& operator=(const ScopedPropVariant&) = delete;
  ~ScopedPropVariant() { PropVariantClear(&value_); }

  PROPVARIANT* put() { return &value_; }
  const PROPVARIANT& get() const { return value_; }

 private:
  PROPVARIANT value_{};
};

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::string WideToUtf8(const wchar_t* value) {
  if (value == nullptr || value[0] == L'\0') {
    return {};
  }

  const int required =
      WideCharToMultiByte(CP_UTF8, 0, value, -1, nullptr, 0, nullptr, nullptr);
  if (required <= 1) {
    return {};
  }

  std::string converted(static_cast<std::size_t>(required - 1), '\0');
  WideCharToMultiByte(
      CP_UTF8, 0, value, -1, converted.data(), required, nullptr, nullptr);
  return converted;
}

std::string DeviceId(IMMDevice* device) {
  if (device == nullptr) {
    return {};
  }

  wchar_t* raw_id = nullptr;
  const HRESULT result = device->GetId(&raw_id);
  if (FAILED(result) || raw_id == nullptr) {
    return {};
  }

  const std::string id = WideToUtf8(raw_id);
  CoTaskMemFree(raw_id);
  return id;
}

std::string PropertyString(IPropertyStore* properties, const PROPERTYKEY& key) {
  if (properties == nullptr) {
    return {};
  }

  ScopedPropVariant property;
  const HRESULT result = properties->GetValue(key, property.put());
  if (FAILED(result)) {
    return {};
  }

  const auto& value = property.get();
  if (value.vt == VT_LPWSTR && value.pwszVal != nullptr) {
    return WideToUtf8(value.pwszVal);
  }
  if (value.vt == VT_BSTR && value.bstrVal != nullptr) {
    return WideToUtf8(value.bstrVal);
  }
  return {};
}

std::uint32_t PropertyUInt32(IPropertyStore* properties, const PROPERTYKEY& key) {
  if (properties == nullptr) {
    return 0;
  }

  ScopedPropVariant property;
  const HRESULT result = properties->GetValue(key, property.put());
  if (FAILED(result)) {
    return 0;
  }

  const auto& value = property.get();
  if (value.vt == VT_UI4) {
    return value.ulVal;
  }
  if (value.vt == VT_I4) {
    return static_cast<std::uint32_t>(value.lVal);
  }
  return 0;
}

std::vector<std::string> DefaultRenderRoles(
    IMMDeviceEnumerator* enumerator,
    const std::string& endpoint_id) {
  if (enumerator == nullptr || endpoint_id.empty()) {
    return {};
  }

  struct RoleName {
    ERole role;
    const char* name;
  };

  constexpr std::array<RoleName, 3> kRoles = {{
      {eConsole, "console"},
      {eMultimedia, "multimedia"},
      {eCommunications, "communications"},
  }};

  std::vector<std::string> roles;
  for (const auto& role : kRoles) {
    ComPtr<IMMDevice> default_device;
    const HRESULT result =
        enumerator->GetDefaultAudioEndpoint(eRender, role.role, default_device.put());
    if (FAILED(result) || !default_device) {
      continue;
    }
    if (DeviceId(default_device.get()) == endpoint_id) {
      roles.push_back(role.name);
    }
  }
  return roles;
}

DeviceCapabilitySnapshot ReadEndpoint(
    IMMDeviceEnumerator* enumerator,
    IMMDevice* endpoint,
    std::uint32_t index) {
  DWORD state_mask = 0;
  const HRESULT state_result = endpoint->GetState(&state_mask);
  const std::string endpoint_id = DeviceId(endpoint);

  std::uint32_t data_flow = 0;
  ComPtr<IMMEndpoint> endpoint_meta;
  if (SUCCEEDED(endpoint->QueryInterface(IID_PPV_ARGS(endpoint_meta.put()))) &&
      endpoint_meta) {
    EDataFlow flow = eRender;
    if (SUCCEEDED(endpoint_meta->GetDataFlow(&flow))) {
      data_flow = static_cast<std::uint32_t>(flow);
    }
  }

  std::string friendly_name;
  std::uint32_t form_factor = 10;
  ComPtr<IPropertyStore> properties;
  if (SUCCEEDED(endpoint->OpenPropertyStore(STGM_READ, properties.put())) &&
      properties) {
    friendly_name = PropertyString(properties.get(), PKEY_Device_FriendlyName);
    form_factor = PropertyUInt32(properties.get(), PKEY_AudioEndpoint_FormFactor);
  }
  if (friendly_name.empty()) {
    std::ostringstream fallback;
    fallback << "WASAPI endpoint " << index;
    friendly_name = fallback.str();
  }

  auto snapshot = MakeWasapiEndpointSnapshot(
      endpoint_id,
      friendly_name,
      static_cast<std::uint32_t>(state_mask),
      data_flow,
      form_factor,
      DefaultRenderRoles(enumerator, endpoint_id));

  if (FAILED(state_result)) {
    snapshot.discovery_error = MakeApiError(
        HResultCode(state_result),
        "IMMDevice::GetState",
        "failed to read endpoint state");
  }
  return snapshot;
}

#endif

}  // namespace

DeviceAvailability WasapiAvailabilityFromStateMask(std::uint32_t state_mask) {
  if ((state_mask & 0x00000001u) != 0) {
    return DeviceAvailability::kActive;
  }
  if ((state_mask & 0x00000002u) != 0) {
    return DeviceAvailability::kDisabled;
  }
  if ((state_mask & 0x00000004u) != 0) {
    return DeviceAvailability::kNotPresent;
  }
  if ((state_mask & 0x00000008u) != 0) {
    return DeviceAvailability::kUnplugged;
  }
  return DeviceAvailability::kUnknown;
}

std::string WasapiStateMaskLabel(std::uint32_t state_mask) {
  struct StateLabel {
    std::uint32_t mask;
    const char* label;
  };

  constexpr std::array<StateLabel, 4> kStates = {{
      {0x00000001u, "active"},
      {0x00000002u, "disabled"},
      {0x00000004u, "not-present"},
      {0x00000008u, "unplugged"},
  }};

  std::string label;
  for (const auto& state : kStates) {
    if ((state_mask & state.mask) != 0) {
      if (!label.empty()) {
        label += "|";
      }
      label += state.label;
    }
  }

  if ((state_mask & ~0x0000000Fu) != 0) {
    if (!label.empty()) {
      label += "|";
    }
    label += "unknown";
  }
  return label.empty() ? "unknown" : label;
}

std::string WasapiDataFlowLabel(std::uint32_t data_flow) {
  switch (data_flow) {
    case 0:
      return "render";
    case 1:
      return "capture";
    case 2:
      return "all";
    default:
      return "unknown";
  }
}

std::string WasapiFormFactorLabel(std::uint32_t form_factor) {
  switch (form_factor) {
    case 0:
      return "remote-network-device";
    case 1:
      return "speakers";
    case 2:
      return "line-level";
    case 3:
      return "headphones";
    case 4:
      return "microphone";
    case 5:
      return "headset";
    case 6:
      return "handset";
    case 7:
      return "unknown-digital-passthrough";
    case 8:
      return "spdif";
    case 9:
      return "digital-audio-display-device";
    case 10:
      return "unknown";
    default:
      return "unknown";
  }
}

DeviceCapabilitySnapshot MakeWasapiEndpointSnapshot(
    std::string endpoint_id,
    std::string friendly_name,
    std::uint32_t state_mask,
    std::uint32_t data_flow,
    std::uint32_t form_factor,
    std::vector<std::string> default_roles) {
  DeviceInfo device;
  device.backend = BackendApi::kWasapiShared;
  device.display_name = std::move(friendly_name);
  device.stable_id = std::move(endpoint_id);
  device.driver_name = "WASAPI/Core Audio";
  device.is_default = !default_roles.empty();
  device.is_known_problem_device = false;

  auto snapshot = MakeDiscoveredDeviceSnapshot(
      std::move(device),
      WasapiAvailabilityFromStateMask(state_mask));
  snapshot.default_roles = std::move(default_roles);
  snapshot.backend_fields["api_device_id"] = snapshot.device.stable_id;
  snapshot.backend_fields["id_kind"] = "wasapi-endpoint-id";
  snapshot.backend_fields["state_mask"] = HexCode(state_mask);
  snapshot.backend_fields["state_label"] = WasapiStateMaskLabel(state_mask);
  snapshot.backend_fields["data_flow"] = WasapiDataFlowLabel(data_flow);
  snapshot.backend_fields["form_factor"] = WasapiFormFactorLabel(form_factor);
  snapshot.backend_fields["default_roles"] = JoinRoles(snapshot.default_roles);

  CapabilityProbeRecord probe;
  probe.kind = CapabilityProbeKind::kDiscoveryOnly;
  probe.support = CapabilitySupport::kSupported;
  probe.detail = "WASAPI endpoint enumerated without stream activation";
  probe.backend_fields["opens_audio_device"] = "false";
  probe.backend_fields["initializes_stream"] = "false";
  probe.backend_fields["starts_stream"] = "false";
  snapshot.probes.push_back(std::move(probe));
  return snapshot;
}

DiscoverySnapshot EnumerateWasapiRenderEndpoints() {
  DiscoverySnapshot snapshot;
  snapshot.backend = BackendApi::kWasapiShared;
  snapshot.opens_audio_device = false;
  snapshot.initializes_stream = false;
  snapshot.starts_stream = false;

#if !defined(_WIN32)
  snapshot.enumeration_error = MakeApiError(
      kHresultNotImplemented,
      "E_NOTIMPL",
      "WASAPI endpoint enumeration is available only on Windows");
  return snapshot;
#else
  ScopedCom com;
  if (FAILED(com.result())) {
    snapshot.enumeration_error = MakeApiError(
        HResultCode(com.result()),
        "CoInitializeEx",
        "failed to initialize COM for endpoint enumeration");
    return snapshot;
  }

  ComPtr<IMMDeviceEnumerator> enumerator;
  HRESULT result = CoCreateInstance(
      kClsidMmDeviceEnumerator,
      nullptr,
      CLSCTX_ALL,
      IID_PPV_ARGS(enumerator.put()));
  if (FAILED(result) || !enumerator) {
    snapshot.enumeration_error = MakeApiError(
        HResultCode(result),
        "CoCreateInstance(CLSID_MMDeviceEnumerator)",
        "failed to create MMDeviceEnumerator");
    return snapshot;
  }

  ComPtr<IMMDeviceCollection> collection;
  result = enumerator->EnumAudioEndpoints(
      eRender,
      DEVICE_STATEMASK_ALL,
      collection.put());
  if (FAILED(result) || !collection) {
    snapshot.enumeration_error = MakeApiError(
        HResultCode(result),
        "IMMDeviceEnumerator::EnumAudioEndpoints",
        "failed to enumerate render endpoints");
    return snapshot;
  }

  UINT count = 0;
  result = collection->GetCount(&count);
  if (FAILED(result)) {
    snapshot.enumeration_error = MakeApiError(
        HResultCode(result),
        "IMMDeviceCollection::GetCount",
        "failed to count render endpoints");
    return snapshot;
  }

  snapshot.devices.reserve(count);
  for (UINT index = 0; index < count; ++index) {
    ComPtr<IMMDevice> endpoint;
    result = collection->Item(index, endpoint.put());
    if (FAILED(result) || !endpoint) {
      DeviceInfo device;
      device.backend = BackendApi::kWasapiShared;
      device.display_name = "WASAPI endpoint unavailable";
      device.stable_id = "wasapi:item-error:" + std::to_string(index);
      device.driver_name = "WASAPI/Core Audio";
      auto failed = MakeDiscoveredDeviceSnapshot(
          std::move(device),
          DeviceAvailability::kUnavailable);
      failed.discovery_error = MakeApiError(
          HResultCode(result),
          "IMMDeviceCollection::Item",
          "failed to read endpoint from collection");
      snapshot.devices.push_back(std::move(failed));
      continue;
    }

    snapshot.devices.push_back(
        ReadEndpoint(enumerator.get(), endpoint.get(), index));
  }

  return snapshot;
#endif
}

}  // namespace aptt
