#pragma region CPL License
/*
Nuclex Native Framework
Copyright (C) 2002-2013 Nuclex Development Labs

This library is free software; you can redistribute it and/or
modify it under the terms of the IBM Common Public License as
published by the IBM Corporation; either version 1.0 of the
License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
IBM Common Public License for more details.

You should have received a copy of the IBM Common Public
License along with this library
*/
#pragma endregion // CPL License

// If the library is compiled as a DLL, this ensures symbols are exported
#define NUCLEX_INPUT_SOURCE 1

#include "XInputDeviceEnumerator.h"

#include <stdexcept>
#include <cstdlib>

#if !defined(_MSC_VER)
namespace {
#endif

  // ------------------------------------------------------------------------------------------- //

  /// <summary>Initializes COM for the duration of the scope</summary>
  class ComScope {

    /// <summary>Initializes COM</summary>
    public: ComScope() {
      HRESULT resultHandle = ::CoInitialize(nullptr);
      if(FAILED(resultHandle)) {
        throw std::runtime_error("Could not initialize COM");
      }
    }

    /// <summary>Shuts down COM</summary>
    public: ~ComScope() {
      ::CoUninitialize();
    }

  };

  // ------------------------------------------------------------------------------------------- //

  /// <summary>Releases an array of device instances</summary>
  class DeviceArrayScope {

    /// <summary>Initializes a new device array releaser</summary>
    /// <param name="devices">Contains the devices that need to be released</param>
    /// <param name="count">Number of devices that need to be released</param>
    public: DeviceArrayScope(IWbemClassObject **devices, std::size_t count) :
      devices(devices), count(count) {}

    /// <summary>Releases the devices the instance was constructed with</summary>
    public: ~DeviceArrayScope() {
      for(std::size_t index = 0; index < this->count; ++index) {
        this->devices[index]->Release();
      }
    }

    /// <summary>Array containing the devices that need to be released</summary>
    private: IWbemClassObject **devices;
    /// <summary>Number of devices contained in the array</summary>
    private: std::size_t count;

  };

  // ------------------------------------------------------------------------------------------- //

  /// <summary>Number of device instances that will be queried from WMI at once</summary>
  const std::size_t BatchSize = 32;

  /// <summary>Timeout for the enumerator to look up the next batch</summary>
  const long TimeoutMilliseconds = 10000;

  // ------------------------------------------------------------------------------------------- //

#if !defined(_MSC_VER)
} // anonymous namespace
#endif

namespace Nuclex { namespace Input {

  // ------------------------------------------------------------------------------------------- //

  std::set<DWORD> XInputDeviceEnumerator::GetXInputVidPidPairs() {
    static const std::wstring vid(L"VID_");
    static const std::wstring pid(L"PID_");
    std::set<DWORD> vidPidPairs;

    // Initialize COM. Yep, this stuff requires COM :-/
    ComScope comScope;

    // Obtain an enumerator for all the device instances in the system
    IEnumWbemClassObjectPtr enumerator = createDeviceInstanceEnumerator(
      getWbemServices(), _bstr_t(L"Win32_PNPEntity")
    );

    // Enumerate all of the devices, looking for the ones which we can indentify as
    // XINPUT devices and adding the VID+PID pairs of these to a list
    IWbemClassObject *devices[BatchSize] = {0};
    for(;;) {
      ULONG returnedDeviceCount;
      HRESULT resultHandle = enumerator->Next(
        TimeoutMilliseconds, BatchSize, devices, &returnedDeviceCount
      );
      if(FAILED(resultHandle)) {
        throw std::runtime_error("Error enumerating WBEM PnP devices");
      }
      if(returnedDeviceCount == 0) {
        break;
      }

      // We have been given a list of devices. Make sure their reference counters are
      // decremented again if we exit for one reason or another
      DeviceArrayScope deviceArrayScope(devices, returnedDeviceCount);
  
      // Look for device instances of XINPUT devices
      for(std::size_t index = 0; index < returnedDeviceCount; ++index) {
        std::wstring deviceId = getDeviceId(devices[index]);
        if(isXInputDeviceId(deviceId)) {
          vidPidPairs.insert(MAKELONG(extractId(deviceId, vid), extractId(deviceId, pid)));
        }
      }

      if(returnedDeviceCount < BatchSize) {
        break;
      }
    }

    return vidPidPairs;
  }

  // ------------------------------------------------------------------------------------------- //

  IWbemServicesPtr XInputDeviceEnumerator::getWbemServices() {
    IWbemLocatorPtr locator(CLSID_WbemLocator);

    // Obtain the WBEM services interface for the local system, accessing the devices
    // namespace in which the PnP devices will be found.
    IWbemServicesPtr services;
    {
      _bstr_t servicesNamespace(L"\\\\.\\root\\cimv2");

      HRESULT resultHandle = locator->ConnectServer(
        servicesNamespace, nullptr, nullptr, nullptr, 0, nullptr, nullptr, &services
      );
      if(FAILED(resultHandle)) {
        throw std::runtime_error("Could not obtain to WBEM services interface");
      }
    }

    // Scary sounding stuff that configures the security options as which we are
    // going to access the services interface (which is supposedly a proxy)
    HRESULT resultHandle = ::CoSetProxyBlanket(
      services, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, nullptr,
      RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE, nullptr, EOAC_NONE
    );
    if(FAILED(resultHandle)) {
      throw std::runtime_error("Could not configure security level for WBEM services");
    }

    return services;
  }

  // ------------------------------------------------------------------------------------------- //

  IEnumWbemClassObjectPtr XInputDeviceEnumerator::createDeviceInstanceEnumerator(
    const IWbemServicesPtr &services, const _bstr_t &deviceClass
  ) {
    IEnumWbemClassObjectPtr enumerator;

    HRESULT resultHandle = services->CreateInstanceEnum(deviceClass, 0, nullptr, &enumerator);
    if(FAILED(resultHandle)) {
      throw std::runtime_error("Could not create enumerator for WBEM PnP entities");
    }

    return enumerator;
  }

  // ------------------------------------------------------------------------------------------- //

  std::wstring XInputDeviceEnumerator::getDeviceId(IWbemClassObject *deviceInstance) {
    static const _bstr_t deviceId(L"DeviceID");

    _variant_t deviceIdValue;
    HRESULT resultHandle = deviceInstance->Get(deviceId, 0, &deviceIdValue, nullptr, nullptr);
    if(FAILED(resultHandle)) {
      throw std::runtime_error("Could not obtain device id of PnP device");
    }

    deviceIdValue.ChangeType(VT_BSTR);

    return std::wstring(deviceIdValue.bstrVal, ::SysStringLen(deviceIdValue.bstrVal));
  }

  // ------------------------------------------------------------------------------------------- //

  bool XInputDeviceEnumerator::isXInputDeviceId(const std::wstring &deviceId) {
    static const std::wstring ig(L"IG_");
    return (deviceId.find(ig) != std::wstring::npos);
  }

  // ------------------------------------------------------------------------------------------- //

  DWORD XInputDeviceEnumerator::extractId(
    const std::wstring &deviceId, const std::wstring &idName
  ) {
    std::size_t index = deviceId.find(idName);
    if(index == std::wstring::npos) {
      return 0;
    }

    // Skip the prefix, we're only interested in the actual number
    index += idName.length();

    // Extract the VID from the device id
    std::wstring value;
    {
      std::size_t end = deviceId.find_first_of(L'&', index);
      if(end == std::wstring::npos) {
        value = deviceId.substr(index);
      } else {
        value = deviceId.substr(index, end - index);
      }
    }

    return ::wcstoul(value.c_str(), nullptr, 16);
  }

  // ------------------------------------------------------------------------------------------- //

}} // namespace Nuclex::Input