#include  <Uefi.h>
#include  <Library/UefiLib.h>
#include  <Library/ShellCEntryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/BaseMemoryLib.h>
#include  <Protocol/UsbIo.h>

#include "EfiKey.h"

extern EFI_BOOT_SERVICES         *gBS;
extern EFI_HANDLE					 gImageHandle;

EFI_GUID  gEfiUsbIoProtocolGuid   = 
	{ 0x2B2F68D6, 0x0CD2, 0x44CF, 
		{ 0x8E, 0x8B, 0xBB, 0xA2, 0x0B, 0x1B, 0x5B, 0x75 }};

EFI_GUID  gEfiSimpleTextInputExProtocolGuid = 
	{0xdd9e7534, 0x7762, 0x4698, 
		{ 0x8c, 0x14, 0xf5, 0x85, 0x17, 0xa6, 0x25, 0xaa }};
		
//C:\UDK2015\MdePkg\Library\UefiUsbLib\Hid.c
/**
  Set the report descriptor of the specified USB HID interface.

  Submit a USB set HID report request for the USB device specified by UsbIo,
  Interface, ReportId, and ReportType, and set the report descriptor using the
  buffer specified by ReportLength and Report.
  If UsbIo is NULL, then ASSERT().
  If Report is NULL, then ASSERT().

  @param  UsbIo         A pointer to the USB I/O Protocol instance for the specific USB target.
  @param  Interface     The index of the report interface on the USB target.
  @param  ReportId      The identifier of the report to retrieve.
  @param  ReportType    The type of report to retrieve.
  @param  ReportLength  The size, in bytes, of Report.
  @param  Report        A pointer to the report descriptor buffer to set.

  @retval  EFI_SUCCESS       The request executed successfully.
  @retval  EFI_TIMEOUT       A timeout occurred executing the request.
  @retval  EFI_DEVICE_ERROR  The request failed due to a device error.

**/
EFI_STATUS
EFIAPI
UsbSetReportRequest (
  IN EFI_USB_IO_PROTOCOL     *UsbIo,
  IN UINT8                   Interface,
  IN UINT8                   ReportId,
  IN UINT8                   ReportType,
  IN UINT16                  ReportLen,
  IN UINT8                   *Report
  )
{
  UINT32                  Status;
  EFI_STATUS              Result;
  EFI_USB_DEVICE_REQUEST  Request;


  //
  // Fill Device request packet
  //
  Request.RequestType = USB_HID_CLASS_SET_REQ_TYPE;
  Request.Request = EFI_USB_SET_REPORT_REQUEST;
  Request.Value   = (UINT16) ((ReportType << 8) | ReportId);
  Request.Index   = Interface;
  Request.Length  = ReportLen;

  Result = UsbIo->UsbControlTransfer (
                    UsbIo,
                    &Request,
                    EfiUsbDataOut,
                    3000, //PcdGet32 (PcdUsbTransferTimeoutValue),
                    Report,
                    ReportLen,
                    &Status
                    );

  return Result;
}

VOID
RunSetKeyLED (
  IN  USB_KB_DEV    *UsbKeyboardDevice
  )
{
  LED_MAP Led;
  UINT8   ReportId;
  UINT8		i;
  
  for (i=0;i<32;i++)	
  {
	  
		  //
		  // Set each field in Led map.
		  //
		  Led.NumLock    = i      % 2;
		  Led.CapsLock   = (i>>1) % 2;
		  Led.ScrollLock = (i>>2) % 2;
		  Led.Resrvd     = 0;

		  ReportId       = 0;
		  //
		  // Call Set_Report Request to lighten the LED.
		  //
		  UsbSetReportRequest (
			UsbKeyboardDevice->UsbIo,
			UsbKeyboardDevice->InterfaceDescriptor.InterfaceNumber,
			ReportId,
			HID_OUTPUT_REPORT,
			1,
			(UINT8 *) &Led
			);
			gBS->Stall(100000);
	}
}
		
UINTN GetUSB( )
{
  EFI_STATUS  Status;
  UINTN       HandleIndex, HandleCount;
  EFI_HANDLE  *DevicePathHandleBuffer = NULL;
  EFI_USB_IO_PROTOCOL 				*USBIO;
  EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL	*SimpleEx;
  USB_KB_DEV                     	*UsbKeyboardDevice;
  EFI_USB_DEVICE_DESCRIPTOR     DeviceDescriptor;
  
  //Get all the Handles that have UsbIO Protocol
  Status = gBS->LocateHandleBuffer(
                  ByProtocol,
                  &gEfiUsbIoProtocolGuid,
                  NULL,
                  &HandleCount,
                  &DevicePathHandleBuffer);
  if (EFI_ERROR(Status)) 
  {
    Print(L"ERROR : Get USBIO count fail.\n");
    return 0;
  }   

  for (HandleIndex = 0; HandleIndex < HandleCount; HandleIndex++) 
  { 
    Status = gBS->HandleProtocol(
                      DevicePathHandleBuffer[HandleIndex],
                      &gEfiUsbIoProtocolGuid,
                      (VOID**)&USBIO);
    if (EFI_ERROR(Status))
      {
        Print(L"ERROR : Open USBIO fail.\n");
        gBS->FreePool(DevicePathHandleBuffer);  
        return 0;
      }
	
	//Check if the Handle has SimpleEx Protocol
 	Status = gBS->OpenProtocol(
		DevicePathHandleBuffer[HandleIndex],
		&gEfiSimpleTextInputExProtocolGuid,
		(VOID**)&SimpleEx,
		gImageHandle,
		NULL,
		EFI_OPEN_PROTOCOL_GET_PROTOCOL);
	//If it has the Protocol, it means it's a USB Keyboard
	if (!EFI_ERROR(Status)) {
		
		//Get USB Device Descriptor
		Status = USBIO->UsbGetDeviceDescriptor(USBIO, &DeviceDescriptor);     
		if (EFI_ERROR(Status))
		{
			Print(L"ERROR : Usb Get Device Descriptor fail.\n");
			gBS->FreePool(DevicePathHandleBuffer);  
			return 0;
		}
      
		//Show the PID and VID
		Print(L"Found a USB Keyboard. VendorID = %04X, ProductID = %04X\n", 
                              DeviceDescriptor.IdVendor, 
                              DeviceDescriptor.IdProduct);       
		
		//Get USB_KB_DEV struct by SimpleEx Protocol
		UsbKeyboardDevice = TEXT_INPUT_EX_USB_KB_DEV_FROM_THIS (SimpleEx);
		//Change LED status
		RunSetKeyLED(UsbKeyboardDevice);
	}
  }
  gBS->FreePool(DevicePathHandleBuffer);       
  return HandleCount;
}


int
EFIAPI
main (
  IN int Argc,
  IN CHAR16 **Argv
  )
{
  GetUSB( );
  return EFI_SUCCESS;
}

