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


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

NTSTATUS ViGEm::Bus::Core::EmulationTargetPDO::PdoCreateDevice(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;
	PEMULATION_TARGET_PDO_CONTEXT	pPdoContext;

	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->PdoPrepareDevice(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

		//
		// Bind this object and device context together
		// 
		pPdoContext = EmulationTargetPdoGetContext(this->PdoDevice);
		pPdoContext->Target = this;
		
		status = this->PdoInitContext();

		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::UsbInterfaceIsDeviceHighSpeed(IN PVOID BusContext)
{
	UNREFERENCED_PARAMETER(BusContext);

	return TRUE;
}

NTSTATUS USB_BUSIFFN ViGEm::Bus::Core::EmulationTargetPDO::UsbInterfaceQueryBusInformation(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::UsbInterfaceSubmitIsoOutUrb(IN PVOID BusContext, IN PURB Urb)
{
	UNREFERENCED_PARAMETER(BusContext);
	UNREFERENCED_PARAMETER(Urb);

	return STATUS_UNSUCCESSFUL;
}

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

	return STATUS_UNSUCCESSFUL;
}

VOID USB_BUSIFFN ViGEm::Bus::Core::EmulationTargetPDO::UsbInterfaceGetUSBDIVersion(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

void ViGEm::Bus::Core::EmulationTargetPDO::UsbAbortPipe()
{
	this->AbortPipe();

	// Higher driver shutting down, emptying PDOs queues
	WdfIoQueuePurge(this->PendingUsbInRequests, NULL, NULL);
	WdfIoQueuePurge(this->PendingNotificationRequests, NULL, NULL);
}

NTSTATUS ViGEm::Bus::Core::EmulationTargetPDO::UsbGetConfigurationDescriptorType(PURB Urb)
{
	const PUCHAR Buffer = (PUCHAR)Urb->UrbControlDescriptorRequest.TransferBuffer;

	// First request just gets required buffer size back
	if (Urb->UrbControlDescriptorRequest.TransferBufferLength == sizeof(USB_CONFIGURATION_DESCRIPTOR))
	{
		const ULONG length = sizeof(USB_CONFIGURATION_DESCRIPTOR);

		this->GetConfigurationDescriptorType(Buffer, length);
		return STATUS_SUCCESS;
	}

	const ULONG length = Urb->UrbControlDescriptorRequest.TransferBufferLength;

	// Second request can store the whole descriptor
	if (length >= UsbConfigurationDescriptionSize)
	{
		this->GetConfigurationDescriptorType(Buffer, UsbConfigurationDescriptionSize);
		return STATUS_SUCCESS;
	}

	return STATUS_UNSUCCESSFUL;
}

NTSTATUS ViGEm::Bus::Core::EmulationTargetPDO::UsbSelectConfiguration(PURB Urb)
{
	TraceEvents(TRACE_LEVEL_VERBOSE,
		TRACE_USBPDO,
		">> >> >> URB_FUNCTION_SELECT_CONFIGURATION: TotalLength %d",
		Urb->UrbHeader.Length);

	if (Urb->UrbHeader.Length == sizeof(struct _URB_SELECT_CONFIGURATION))
	{
		TraceEvents(TRACE_LEVEL_VERBOSE,
			TRACE_USBPDO,
			">> >> >> URB_FUNCTION_SELECT_CONFIGURATION: NULL ConfigurationDescriptor");
		return STATUS_SUCCESS;
	}

	return this->SelectConfiguration(Urb);
}

ViGEm::Bus::Core::EmulationTargetPDO::EmulationTargetPDO(USHORT VID, USHORT PID) : VendorId(VID), ProductId(PID)
{
	KeInitializeEvent(&this->PdoBootNotificationEvent, NotificationEvent, FALSE);
}

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->PdoPrepareHardware();
}

VOID ViGEm::Bus::Core::EmulationTargetPDO::EvtIoInternalDeviceControl(
	_In_ WDFQUEUE Queue,
	_In_ WDFREQUEST Request,
	_In_ size_t OutputBufferLength,
	_In_ size_t InputBufferLength,
	_In_ ULONG IoControlCode
)
{
	const auto ctx = EmulationTargetPdoGetContext(WdfIoQueueGetDevice(Queue));
		
	UNREFERENCED_PARAMETER(OutputBufferLength);
	UNREFERENCED_PARAMETER(InputBufferLength);

	NTSTATUS                status = STATUS_INVALID_PARAMETER;
	PIRP                    irp;
	PURB                    urb;
	PIO_STACK_LOCATION      irpStack;

	TraceDbg(TRACE_BUSPDO, "%!FUNC! Entry");

	// No help from the framework available from here on
	irp = WdfRequestWdmGetIrp(Request);
	irpStack = IoGetCurrentIrpStackLocation(irp);

	switch (IoControlCode)
	{
	case IOCTL_INTERNAL_USB_SUBMIT_URB:

		TraceEvents(TRACE_LEVEL_VERBOSE,
			TRACE_BUSPDO,
			">> IOCTL_INTERNAL_USB_SUBMIT_URB");

		urb = static_cast<PURB>(URB_FROM_IRP(irp));

		switch (urb->UrbHeader.Function)
		{
		case URB_FUNCTION_CONTROL_TRANSFER:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_CONTROL_TRANSFER");

			status = ctx->Target->UsbControlTransfer(urb);

			break;

		case URB_FUNCTION_CONTROL_TRANSFER_EX:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_CONTROL_TRANSFER_EX");

			status = STATUS_UNSUCCESSFUL;

			break;

		case URB_FUNCTION_BULK_OR_INTERRUPT_TRANSFER:

			TraceDbg(
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_BULK_OR_INTERRUPT_TRANSFER");

			status = ctx->Target->UsbBulkOrInterruptTransfer(&urb->UrbBulkOrInterruptTransfer, Request);

			break;

		case URB_FUNCTION_SELECT_CONFIGURATION:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_SELECT_CONFIGURATION");

			status = ctx->Target->UsbSelectConfiguration(urb);

			break;

		case URB_FUNCTION_SELECT_INTERFACE:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_SELECT_INTERFACE");

			status = ctx->Target->UsbSelectInterface(urb);

			break;

		case URB_FUNCTION_GET_DESCRIPTOR_FROM_DEVICE:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_GET_DESCRIPTOR_FROM_DEVICE");

			switch (urb->UrbControlDescriptorRequest.DescriptorType)
			{
			case USB_DEVICE_DESCRIPTOR_TYPE:

				TraceEvents(TRACE_LEVEL_VERBOSE,
					TRACE_BUSPDO,
					">> >> >> USB_DEVICE_DESCRIPTOR_TYPE");

				ctx->Target->UsbGetDeviceDescriptorType(static_cast<PUSB_DEVICE_DESCRIPTOR>(urb->UrbControlDescriptorRequest.TransferBuffer));

				break;

			case USB_CONFIGURATION_DESCRIPTOR_TYPE:

				TraceEvents(TRACE_LEVEL_VERBOSE,
					TRACE_BUSPDO,
					">> >> >> USB_CONFIGURATION_DESCRIPTOR_TYPE");

				status = ctx->Target->UsbGetConfigurationDescriptorType(urb);

				break;

			case USB_STRING_DESCRIPTOR_TYPE:

				TraceEvents(TRACE_LEVEL_VERBOSE,
					TRACE_BUSPDO,
					">> >> >> USB_STRING_DESCRIPTOR_TYPE");

				status = ctx->Target->UsbGetStringDescriptorType(urb);

				break;

			default:

				TraceEvents(TRACE_LEVEL_VERBOSE,
					TRACE_BUSPDO,
					">> >> >> Unknown descriptor type");

				break;
			}

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				"<< <<");

			break;

		case URB_FUNCTION_GET_STATUS_FROM_DEVICE:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_GET_STATUS_FROM_DEVICE");

			// Defaults always succeed
			status = STATUS_SUCCESS;

			break;

		case URB_FUNCTION_ABORT_PIPE:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_ABORT_PIPE");

			ctx->Target->UsbAbortPipe();

			break;

		case URB_FUNCTION_CLASS_INTERFACE:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_CLASS_INTERFACE");

			status = ctx->Target->UsbClassInterface(urb);

			break;

		case URB_FUNCTION_GET_DESCRIPTOR_FROM_INTERFACE:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >> URB_FUNCTION_GET_DESCRIPTOR_FROM_INTERFACE");

			status = ctx->Target->UsbGetDescriptorFromInterface(urb);

			break;

		default:

			TraceEvents(TRACE_LEVEL_VERBOSE,
				TRACE_BUSPDO,
				">> >>  Unknown function: 0x%X",
				urb->UrbHeader.Function);

			break;
		}

		TraceEvents(TRACE_LEVEL_VERBOSE,
			TRACE_BUSPDO,
			"<<");

		break;

	case IOCTL_INTERNAL_USB_GET_PORT_STATUS:

		TraceEvents(TRACE_LEVEL_VERBOSE,
			TRACE_BUSPDO,
			">> IOCTL_INTERNAL_USB_GET_PORT_STATUS");

		// We report the (virtual) port as always active
		*static_cast<unsigned long*>(irpStack->Parameters.Others.Argument1) = USBD_PORT_ENABLED | USBD_PORT_CONNECTED;

		status = STATUS_SUCCESS;

		break;

	case IOCTL_INTERNAL_USB_RESET_PORT:

		TraceEvents(TRACE_LEVEL_VERBOSE,
			TRACE_BUSPDO,
			">> IOCTL_INTERNAL_USB_RESET_PORT");

		// Sure, why not ;)
		status = STATUS_SUCCESS;

		break;

	case IOCTL_INTERNAL_USB_SUBMIT_IDLE_NOTIFICATION:

		TraceEvents(TRACE_LEVEL_VERBOSE,
			TRACE_BUSPDO,
			">> IOCTL_INTERNAL_USB_SUBMIT_IDLE_NOTIFICATION");

		// TODO: implement?
		// This happens if the I/O latency is too high so HIDUSB aborts communication.
		status = STATUS_SUCCESS;

		break;

	default:

		TraceEvents(TRACE_LEVEL_VERBOSE,
			TRACE_BUSPDO,
			">> Unknown I/O control code 0x%X",
			IoControlCode);

		break;
	}

	if (status != STATUS_PENDING)
	{
		WdfRequestComplete(Request, status);
	}

	TraceDbg(TRACE_BUSPDO, "%!FUNC! Exit with status %!STATUS!", status);
}