#include "EmulationTargetPDO.hpp"
#include "CRTCPP.hpp"
#include "trace.h"
#include "EmulationTargetPDO.tmh"
#define NTSTRSAFE_LIB
#include <ntstrsafe.h>
#include <usbiodef.h>


PCWSTR ViGEm::Bus::Core::EmulationTargetPDO::_deviceLocation = L"Virtual Gamepad Emulation Bus";

NTSTATUS ViGEm::Bus::Core::EmulationTargetPDO::CreateDevice(WDFDEVICE Device, PWDFDEVICE_INIT DeviceInit, PPDO_IDENTIFICATION_DESCRIPTION Description)
{
	NTSTATUS                        status = STATUS_UNSUCCESSFUL;
	WDF_DEVICE_PNP_CAPABILITIES     pnpCaps;
	WDF_DEVICE_POWER_CAPABILITIES   powerCaps;
	WDF_PNPPOWER_EVENT_CALLBACKS    pnpPowerCallbacks;
	WDF_OBJECT_ATTRIBUTES           pdoAttributes;
	WDF_IO_QUEUE_CONFIG             defaultPdoQueueConfig;
	WDFQUEUE                        defaultPdoQueue;
	UNICODE_STRING                  deviceDescription;
	WDF_OBJECT_ATTRIBUTES           attributes;
	WDF_IO_QUEUE_CONFIG             usbInQueueConfig;
	WDF_IO_QUEUE_CONFIG             notificationsQueueConfig;

	DECLARE_CONST_UNICODE_STRING(deviceLocation, L"Virtual Gamepad Emulation Bus");
	DECLARE_UNICODE_STRING_SIZE(buffer, MAX_INSTANCE_ID_LEN);
	// reserve space for device id
	DECLARE_UNICODE_STRING_SIZE(deviceId, MAX_INSTANCE_ID_LEN);

	UNREFERENCED_PARAMETER(Description);
	UNREFERENCED_PARAMETER(Device);

	PAGED_CODE();

	// set device type
	WdfDeviceInitSetDeviceType(DeviceInit, FILE_DEVICE_BUS_EXTENDER);
	// Bus is power policy owner
	WdfDeviceInitSetPowerPolicyOwnership(DeviceInit, FALSE);

	do
	{
#pragma region Enter RAW device mode

		status = WdfPdoInitAssignRawDevice(DeviceInit, &GUID_DEVCLASS_VIGEM_RAWPDO);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfPdoInitAssignRawDevice failed with status %!STATUS!",
				status);
			break;
		}

		WdfDeviceInitSetCharacteristics(DeviceInit, FILE_AUTOGENERATED_DEVICE_NAME, TRUE);

		status = WdfDeviceInitAssignSDDLString(DeviceInit, &SDDL_DEVOBJ_SYS_ALL_ADM_RWX_WORLD_RWX_RES_RWX);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfDeviceInitAssignSDDLString failed with status %!STATUS!",
				status);
			break;
		}

#pragma endregion

#pragma region Prepare PDO

		status = this->PrepareDevice(DeviceInit, &deviceId, &deviceDescription);

		if (!NT_SUCCESS(status))
			break;

		// set device id
		status = WdfPdoInitAssignDeviceID(DeviceInit, &deviceId);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfPdoInitAssignDeviceID failed with status %!STATUS!",
				status);
			break;
		}

		// prepare instance id
		status = RtlUnicodeStringPrintf(&buffer, L"%02d", this->SerialNo);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"RtlUnicodeStringPrintf failed with status %!STATUS!",
				status);
			break;
		}

		// set instance id
		status = WdfPdoInitAssignInstanceID(DeviceInit, &buffer);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfPdoInitAssignInstanceID failed with status %!STATUS!",
				status);
			break;
		}

		// set device description (for English operating systems)
		status = WdfPdoInitAddDeviceText(DeviceInit, &deviceDescription, &deviceLocation, 0x409);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfPdoInitAddDeviceText failed with status %!STATUS!",
				status);
			break;
		}

		// default locale is English
		// TODO: add more locales
		WdfPdoInitSetDefaultLocale(DeviceInit, 0x409);

#pragma region PNP/Power event callbacks

		WDF_PNPPOWER_EVENT_CALLBACKS_INIT(&pnpPowerCallbacks);

		pnpPowerCallbacks.EvtDevicePrepareHardware = EvtDevicePrepareHardware;

		WdfDeviceInitSetPnpPowerEventCallbacks(DeviceInit, &pnpPowerCallbacks);

#pragma endregion

		// NOTE: not utilized at the moment
		WdfPdoInitAllowForwardingRequestToParent(DeviceInit);

#pragma endregion

#pragma region Create PDO

		// Add common device data context
		WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&pdoAttributes, EMULATION_TARGET_PDO_CONTEXT);

		status = WdfDeviceCreate(&DeviceInit, &pdoAttributes, &this->PdoDevice);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfDeviceCreate failed with status %!STATUS!",
				status);
			break;
		}

		TraceEvents(TRACE_LEVEL_VERBOSE,
			TRACE_BUSPDO,
			"Created PDO 0x%p",
			this->PdoDevice);

#pragma endregion

#pragma region Expose USB Interface

		status = WdfDeviceCreateDeviceInterface(
			Device, 
			const_cast<LPGUID>(&GUID_DEVINTERFACE_USB_DEVICE), 
			NULL
		);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfDeviceCreateDeviceInterface failed with status %!STATUS!",
				status);
			break;
		}

#pragma endregion

#pragma region Set PDO contexts

		status = this->InitContext();

		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"Couldn't initialize additional contexts: %!STATUS!",
				status);
			break;
		}

#pragma endregion

#pragma region Create Queues & Locks

		WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
		attributes.ParentObject = this->PdoDevice;

		// Create and assign queue for incoming interrupt transfer
		WDF_IO_QUEUE_CONFIG_INIT(&usbInQueueConfig, WdfIoQueueDispatchManual);

		status = WdfIoQueueCreate(
			this->PdoDevice,
			&usbInQueueConfig,
			WDF_NO_OBJECT_ATTRIBUTES,
			&this->PendingUsbInRequests
		);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfIoQueueCreate (PendingUsbInRequests) failed with status %!STATUS!",
				status);
			break;
		}

		// Create and assign queue for user-land notification requests
		WDF_IO_QUEUE_CONFIG_INIT(&notificationsQueueConfig, WdfIoQueueDispatchManual);

		status = WdfIoQueueCreate(
			Device, 
			&notificationsQueueConfig, 
			WDF_NO_OBJECT_ATTRIBUTES, 
			&this->PendingNotificationRequests
		);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfIoQueueCreate (PendingNotificationRequests) failed with status %!STATUS!",
				status);
			break;
		}

#pragma endregion 

#pragma region Default I/O queue setup

		WDF_IO_QUEUE_CONFIG_INIT_DEFAULT_QUEUE(&defaultPdoQueueConfig, WdfIoQueueDispatchParallel);

		defaultPdoQueueConfig.EvtIoInternalDeviceControl = EvtIoInternalDeviceControl;

		status = WdfIoQueueCreate(
			this->PdoDevice, 
			&defaultPdoQueueConfig, 
			WDF_NO_OBJECT_ATTRIBUTES, 
			&defaultPdoQueue
		);
		if (!NT_SUCCESS(status))
		{
			TraceEvents(TRACE_LEVEL_ERROR,
				TRACE_BUSPDO,
				"WdfIoQueueCreate (Default) failed with status %!STATUS!",
				status);
			break;
		}

#pragma endregion

#pragma region PNP capabilities

		WDF_DEVICE_PNP_CAPABILITIES_INIT(&pnpCaps);

		pnpCaps.Removable = WdfTrue;
		pnpCaps.EjectSupported = WdfTrue;
		pnpCaps.SurpriseRemovalOK = WdfTrue;

		pnpCaps.Address = this->SerialNo;
		pnpCaps.UINumber = this->SerialNo;

		WdfDeviceSetPnpCapabilities(this->PdoDevice, &pnpCaps);

#pragma endregion

#pragma region Power capabilities

		WDF_DEVICE_POWER_CAPABILITIES_INIT(&powerCaps);

		powerCaps.DeviceD1 = WdfTrue;
		powerCaps.WakeFromD1 = WdfTrue;
		powerCaps.DeviceWake = PowerDeviceD1;

		powerCaps.DeviceState[PowerSystemWorking] = PowerDeviceD0;
		powerCaps.DeviceState[PowerSystemSleeping1] = PowerDeviceD1;
		powerCaps.DeviceState[PowerSystemSleeping2] = PowerDeviceD3;
		powerCaps.DeviceState[PowerSystemSleeping3] = PowerDeviceD3;
		powerCaps.DeviceState[PowerSystemHibernate] = PowerDeviceD3;
		powerCaps.DeviceState[PowerSystemShutdown] = PowerDeviceD3;

		WdfDeviceSetPowerCapabilities(this->PdoDevice, &powerCaps);

#pragma endregion

	} while (FALSE);

	TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_BUSPDO, "%!FUNC! Exit with status %!STATUS!", status);
	
	return status;
}

VOID ViGEm::Bus::Core::EmulationTargetPDO::SetSerial(ULONG Serial)
{
	this->SerialNo = Serial;
}

ULONG ViGEm::Bus::Core::EmulationTargetPDO::GetSerial() const
{
	return this->SerialNo;
}

#pragma region USB Interface Functions

BOOLEAN USB_BUSIFFN ViGEm::Bus::Core::EmulationTargetPDO::UsbIsDeviceHighSpeed(IN PVOID BusContext)
{
	UNREFERENCED_PARAMETER(BusContext);

	return TRUE;
}

NTSTATUS USB_BUSIFFN ViGEm::Bus::Core::EmulationTargetPDO::UsbQueryBusInformation(IN PVOID BusContext, IN ULONG Level,
	IN OUT PVOID BusInformationBuffer,
	IN OUT PULONG BusInformationBufferLength,
	OUT PULONG BusInformationActualLength)
{
	UNREFERENCED_PARAMETER(BusContext);
	UNREFERENCED_PARAMETER(Level);
	UNREFERENCED_PARAMETER(BusInformationBuffer);
	UNREFERENCED_PARAMETER(BusInformationBufferLength);
	UNREFERENCED_PARAMETER(BusInformationActualLength);

	return STATUS_UNSUCCESSFUL;
}

NTSTATUS USB_BUSIFFN ViGEm::Bus::Core::EmulationTargetPDO::UsbSubmitIsoOutUrb(IN PVOID BusContext, IN PURB Urb)
{
	UNREFERENCED_PARAMETER(BusContext);
	UNREFERENCED_PARAMETER(Urb);

	return STATUS_UNSUCCESSFUL;
}

NTSTATUS USB_BUSIFFN ViGEm::Bus::Core::EmulationTargetPDO::UsbQueryBusTime(IN PVOID BusContext, IN OUT PULONG CurrentUsbFrame)
{
	UNREFERENCED_PARAMETER(BusContext);
	UNREFERENCED_PARAMETER(CurrentUsbFrame);

	return STATUS_UNSUCCESSFUL;
}

VOID USB_BUSIFFN ViGEm::Bus::Core::EmulationTargetPDO::UsbGetUSBDIVersion(IN PVOID BusContext,
	IN OUT PUSBD_VERSION_INFORMATION VersionInformation,
	IN OUT PULONG HcdCapabilities)
{
	UNREFERENCED_PARAMETER(BusContext);

	if (VersionInformation != nullptr)
	{
		VersionInformation->USBDI_Version = 0x500; /* Usbport */
		VersionInformation->Supported_USB_Version = 0x200; /* USB 2.0 */
	}

	if (HcdCapabilities != nullptr)
	{
		*HcdCapabilities = 0;
	}
}

#pragma endregion

ViGEm::Bus::Core::EmulationTargetPDO::EmulationTargetPDO(USHORT VID, USHORT PID) : VendorId(VID), ProductId(PID)
{
}

NTSTATUS ViGEm::Bus::Core::EmulationTargetPDO::EvtDevicePrepareHardware(
	_In_ WDFDEVICE Device,
	_In_ WDFCMRESLIST ResourcesRaw,
	_In_ WDFCMRESLIST ResourcesTranslated
)
{
	UNREFERENCED_PARAMETER(ResourcesRaw);
	UNREFERENCED_PARAMETER(ResourcesTranslated);
	
	const auto ctx = EmulationTargetPdoGetContext(Device);

	return ctx->Target->PrepareHardware();
}

VOID ViGEm::Bus::Core::EmulationTargetPDO::EvtIoInternalDeviceControl(
	_In_ WDFQUEUE Queue,
	_In_ WDFREQUEST Request,
	_In_ size_t OutputBufferLength,
	_In_ size_t InputBufferLength,
	_In_ ULONG IoControlCode
)
{
	UNREFERENCED_PARAMETER(Queue);
	UNREFERENCED_PARAMETER(OutputBufferLength);
	UNREFERENCED_PARAMETER(InputBufferLength);
	UNREFERENCED_PARAMETER(IoControlCode);

	WdfRequestComplete(Request, STATUS_UNSUCCESSFUL);
}