/** @file
  Main Funcion file.

;******************************************************************************
;* Copyright (c) 2016, Lenovo Corporation. All Rights Reserved.
;*
;* You may not reproduce, distribute, publish, display, perform, modify, adapt,
;* transmit, broadcast, present, recite, release, license or otherwise exploit
;* any part of this publication in any form, by any means, without the prior
;* written permission of Lenovo Corporation.
;*
;******************************************************************************
*/
#include <Uefi.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Protocol/GraphicsOutput.h>

#include <Protocol/UsbIo.h>
#include <Protocol/DevicePath.h>
#include <Library/DevicePathLib.h>
#include <Protocol/DevicePathToText.h>
#include <Library/BaseMemoryLib.h>
#include <Library/BaseLib.h>
#include <LfcCmos.h>


#include "OneKeyBatteryLib.h"

//define if this project support OneKeyBattery through AOU bootup.
#define AOU_BOOTUP_SUPPORT 0
#if AOU_BOOTUP_SUPPORT
//Please ODM define the AOU device path as HW design.
#define AOU_Port L"PciRoot(0x0)/Pci(0x14,0x0)/USB(0x1,0x0)"
#define USB_CLASS_HID              3
#endif

//for bootup state
typedef enum {
  Normal_Bootup_State = 0,
  Keyboard_Bootup_State = 1,
  Attach_AC_Bootup_State = 2,
  Attach_AOU_Bootup_State = 3
} BOOTUP_STATE_FOR_SHOW_ONEKEYBATTERY;
BOOTUP_STATE_FOR_SHOW_ONEKEYBATTERY BootupState = Normal_Bootup_State;

//for loop run
#define SECONDS_OF_DISPLAY                3

#define GAMINGY900 1
#if GAMINGY900
#include <Library/IoLib.h>
#include <Library/TimerLib.h>

#define KBC_TIME_OUT                      0x10000
#define KEY_OBF                           1
#define KEY_IBF                           2

#define EC_DATA                           0x62
#define EC_CMD_STATE                      0x66

#define EC_READ_ECRAM_CMD                 0x80
//#[-start-1701124-JESSE0143-Add]#
#define EC_WRITE_ECRAM_CMD                0x81
//#[-end-1701124-JESSE0143-Add]#
#define BATT_REMAINING_CAP_LOW            0xC2
#define BATT_REMAINING_CAP_HIGH           0xC3
#define BATT_FULL_CHARGE_CAP_LOW          0xCC
#define BATT_FULL_CHARGE_CAP_HIGH         0xCD

#define ADAPTER_STATE                     0xA3
#define ADAPTER_STATE_ONLINE              0X80

#define EC_BOOTUP_STATE_CMD                    0x49

/* Yoga.
#define BATT_REMAINING_CAP_LOW            0x6A
#define BATT_REMAINING_CAP_HIGH           0x6B
#define BATT_FULL_CHARGE_CAP_LOW          0x6C
#define BATT_FULL_CHARGE_CAP_HIGH         0x6D

#define ADAPTER_STATE                     0x40
#define ADAPTER_STATE_ONLINE              0X01

#define EC_BOOTUP_STATE_CMD                    0xBC
*/

EFI_STATUS
WaitIBE (
  IN UINT16  CommandPort
  )
{
  UINT8                 CmdPort = 0;
  UINTN                 Index;

  for (Index = 0; Index < KBC_TIME_OUT; Index++) {
    CmdPort = IoRead8 (CommandPort);
    if (!(CmdPort & KEY_IBF)) {
      return EFI_SUCCESS;
    } else {
      MicroSecondDelay(15);
    }
  }

  return EFI_DEVICE_ERROR;
}

EFI_STATUS
WaitOBF (
  IN UINT16  CommandPort
)
{
  UINT8                 CmdPort = 0;
  UINTN                 Index;

  for (Index = 0; Index < KBC_TIME_OUT; Index++) {
    CmdPort = IoRead8 (CommandPort);
    if (CmdPort & KEY_OBF) {
      return EFI_SUCCESS;
    } else {
      MicroSecondDelay(15);
    }
  }

  return EFI_DEVICE_ERROR;
}

EFI_STATUS
WaitOBE (
  IN UINT16  CommandPort
)
{
  UINT8                 CmdPort = 0;
  UINTN                 Index;

  for (Index = 0; Index < KBC_TIME_OUT; Index++) {
    CmdPort = IoRead8 (CommandPort);
    if (CmdPort & KEY_OBF) {
      //Read data out from data port
      IoRead8 (CommandPort - 4);
    } else {
      return EFI_SUCCESS;
    }
  }

  return EFI_DEVICE_ERROR;
}

EFI_STATUS
EcCommand (
  IN UINT16  CommandPort,
  IN UINT16  DataPort,
  IN UINTN   NumOfReturnData,
  IN UINT8   *ReturnData,
  IN UINT8   Command,
  IN UINTN   NumOfArgs,
  ...
  )
{
  EFI_STATUS               Status;
  VA_LIST                  Marker;
  UINT8                    Index;
  UINT8                    *Data;

  Status = WaitIBE (CommandPort);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  Status = WaitOBE (CommandPort);
  if (EFI_ERROR (Status)) {
    return Status;
  }
  //Send command
  IoWrite8 (CommandPort, Command);

  Status = WaitIBE (CommandPort);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  if (NumOfArgs > 0){
    VA_START (Marker, NumOfArgs);
    for(Index = 0; Index < NumOfArgs; Index++){
      Data = VA_ARG (Marker, UINT8*);
      //Send data
      IoWrite8 (DataPort, *Data);

      Status = WaitIBE (CommandPort);
      if (EFI_ERROR(Status)) {
        return Status;
      }
    }
    VA_END (Marker);
    //Get return data
  }

  if (NumOfReturnData > 0) {
    for (Index = 0; Index < NumOfReturnData; Index++) {
      Status = WaitOBF (CommandPort);
      if (EFI_ERROR(Status)) {
        return Status;
      }
      *(ReturnData + Index) = IoRead8 (DataPort);
    }
  }
  return Status;
}


EFI_STATUS
EcRamRead(
  IN  UINT8  Index,
  OUT UINT8  *Data
  )
{
  EFI_STATUS           Status;
  UINT8                Data8;

  Status = EcCommand (
             EC_CMD_STATE,
             EC_DATA,
             1,
             &Data8,
             EC_READ_ECRAM_CMD,
             1,
             &Index);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  *Data = Data8;
  return  Status;
}

//#[-start-1701124-JESSE0143-Add]#
EFI_STATUS
EcRamWrite(
  IN UINT8                           Index,
  IN UINT8                           Value
  )
{
  EFI_STATUS           Status;

  Status = EcCommand (
             EC_CMD_STATE,
             EC_DATA,
             0,
             NULL,
             EC_WRITE_ECRAM_CMD,
             2,
             &Index,
             &Value);
  return Status;
}
//#[-end-1701124-JESSE0143-Add]#

/**
  Get Battery Percentage.

  @param[OUT] BatteryPercentage      The Battery Percentage.

  @retval     EFI_SUCCESS            Executed successfully.
  @retval     other                  Error occurs.

**/
EFI_STATUS
GetBatteryPercentage (
  IN OUT  UINT8  *BatteryPercentage
  )
{
  UINT16              RemainingCap, FullChargeCap;
  EFI_STATUS          Status;
  UINT8               Data1 = 1, Data2 = 1, Data3 = 1, Data4 = 1;


  Status = EcRamRead (BATT_REMAINING_CAP_HIGH, &Data1);
  if (EFI_ERROR(Status)) {
    return Status;
  }
  Status = EcRamRead (BATT_REMAINING_CAP_LOW, &Data2);
  if (EFI_ERROR(Status)) {
    return Status;
  }
  Status = EcRamRead (BATT_FULL_CHARGE_CAP_HIGH, &Data3);
  if (EFI_ERROR(Status)) {
    return Status;
  }
  Status = EcRamRead (BATT_FULL_CHARGE_CAP_LOW, &Data4);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  RemainingCap  = (((UINT16)(Data1))<<8) + (UINT16)Data2;
  FullChargeCap = (((UINT16)(Data3))<<8) + (UINT16)Data4;

  //If battery not exist or battery error
  if (FullChargeCap == 0 || 100*RemainingCap/FullChargeCap > 100) {
    //If error
    *BatteryPercentage = 15;
    return EFI_SUCCESS;
  }

  *BatteryPercentage = (UINT8)(100*RemainingCap/FullChargeCap);
  //if (*BatteryPercentage >=20 && *BatteryPercentage <= 30) {
  //  *BatteryPercentage = 15;  //Make it under 20 to meet flashit's error condition
  //}
  return EFI_SUCCESS;
}

/**
  Get AC state.

  @param[OUT] ACStatus      The AC is pluged in or not. (TRUE or FALSE)

  @retval EFI_SUCCESS       Executed successfully.
  @retval other             Error occurs.

**/
EFI_STATUS
GetACStatus (
  OUT BOOLEAN  * ACStatus
  )
{
  EFI_STATUS        Status;
  UINT8             Data = 0;

  //Get AC Status.
  Status = EcRamRead (ADAPTER_STATE, &Data );
  if (EFI_ERROR(Status)) {
    return Status;
  }

  if((Data &ADAPTER_STATE_ONLINE) == 0)
    *ACStatus = FALSE;
  else
    *ACStatus = TRUE;

  return EFI_SUCCESS;
}

#if AOU_BOOTUP_SUPPORT
/**
  Get USB state.

  @param[OUT] AOUStatus     The USB device is pluged in AOU port or not. (TRUE or FALSE)

  @retval EFI_SUCCESS       Executed successfully.
  @retval other             Error occurs.

**/
EFI_STATUS
GetAOUStatus (
  OUT BOOLEAN  * AOUStatus
  )
{
	//Get AOC Status.
	//The AOC status requires ODM implementation., we set FALSE for test.
  *AOUStatus = FALSE;

  return EFI_SUCCESS;
}
#endif

/**
  Get Boot up state.

  @param[OUT] Bootup_State     For return Bootup state(KB,AC,AOU,Normal)

  @retval EFI_SUCCESS          Executed successfully.
  @retval other                Error occurs.

**/
EFI_STATUS
EcBootupState (
  OUT BOOTUP_STATE_FOR_SHOW_ONEKEYBATTERY *Bootup_State
  )
{
  EFI_STATUS        Status = EFI_SUCCESS;
  UINT8             BootupStatus;

  //
  // Get BootUp state
  //
  Status = EcCommand (
             EC_CMD_STATE,
             EC_DATA,
             1,
             &BootupStatus,
             EC_BOOTUP_STATE_CMD,
             0);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  // Check if bootup state is for show OneKeyBattery
  if ((BootupStatus  == Normal_Bootup_State) ||
    (BootupStatus  == Keyboard_Bootup_State) ||
    (BootupStatus  == Attach_AC_Bootup_State) ||
    (BootupStatus  == Attach_AOU_Bootup_State)) {
    *Bootup_State = BootupStatus;
    return Status;
  } else{
    return EFI_NOT_FOUND;
  }
}
#else
/**
  Get Battery Percentage.

  @param[OUT] BatteryPercentage      The Battery Percentage.

  @retval     EFI_SUCCESS            Executed successfully.
  @retval     other                  Error occurs.

**/
EFI_STATUS
GetBatteryPercentage (
  OUT  UINT8  *BatteryPercentage
  )
{
  /*++
      Todo:
        Add project specific code in here.
    --*/

    return EFI_UNSUPPORTED;
}

/**
  Get AC state.

  @param[OUT] ACStatus      The AC in pluged in or not. (TRUE or FALSE)

  @retval EFI_SUCCESS       Executed successfully.
  @retval other             Error occurs.

**/
EFI_STATUS
GetACStatus (
  OUT BOOLEAN  *ACStatus
  )
{
  /*++
      Todo:
        Add project specific code in here.
    --*/

    return EFI_UNSUPPORTED;
}

#if AOU_BOOTUP_SUPPORT
/**
  Get AOU state.

  @param[OUT] AOUStatus     The USB device is pluged in AOU port or not. (TRUE or FALSE)

  @retval EFI_SUCCESS       Executed successfully.
  @retval other             Error occurs.

**/
EFI_STATUS
GetAOUStatus (
  OUT BOOLEAN  * AOUStatus
  )
{
  /*++
      Todo:
        Add project specific code in here.
    --*/

    return EFI_UNSUPPORTED;
}
#endif

/**
  Get Boot up state.

  @param[OUT] Bootup_State     For return Bootup state(KB,AC,AOU,Normal)

  @retval EFI_SUCCESS          Executed successfully.
  @retval other                Error occurs.

**/
EFI_STATUS
EcBootupState (
  OUT BOOTUP_STATE_FOR_SHOW_ONEKEYBATTERY *Bootup_State
  )
{
  /*++
      Todo:
        Add project specific code in here.
    --*/

    return EFI_UNSUPPORTED;
}

#endif

#if AOU_BOOTUP_SUPPORT
BOOLEAN
CheckUsbHidDevice(
  VOID
  )
{
  EFI_STATUS                        Status;
  UINTN                             UsbIoHandleCount;
  EFI_HANDLE                        *UsbIoHandleBuffer;
  EFI_USB_IO_PROTOCOL               *UsbIo;
  EFI_USB_INTERFACE_DESCRIPTOR      IfDesc;

  EFI_DEVICE_PATH_TO_TEXT_PROTOCOL  *DpathToText;
  EFI_DEVICE_PATH_PROTOCOL          *DevPath;
  CHAR16                            *TextStr;
  UINTN                             Index;

  //
  // Get all UsbIo Handles.
  //
  UsbIoHandleCount = 0;
  UsbIoHandleBuffer = NULL;
  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiUsbIoProtocolGuid,
                  NULL,
                  &UsbIoHandleCount,
                  &UsbIoHandleBuffer
                  );
  if (EFI_ERROR (Status) || (UsbIoHandleCount == 0) || (UsbIoHandleBuffer == NULL)) {
    return FALSE;
  }

  //get function which converts device path to text
  DpathToText = NULL;
  Status = gBS->LocateProtocol (
                  &gEfiDevicePathToTextProtocolGuid,
                  NULL,
                  (VOID **)&DpathToText
                  );
  if (EFI_ERROR (Status) || (DpathToText == NULL)) {
    return FALSE;
  }

  //polling all the usb devices
  for (Index = 0; Index < UsbIoHandleCount; Index++) {
    Status = gBS->HandleProtocol (
                    UsbIoHandleBuffer[Index],
                    &gEfiDevicePathProtocolGuid,
                    (VOID **)&DevPath
                    );

    //converts device path to text
    TextStr = DpathToText->ConvertDevicePathToText (DevPath, TRUE, TRUE);

    //compare to the default usb charge ic port string.
    if (StrCmp (AOU_Port, TextStr) == 0) {
      Status = gBS->HandleProtocol(
                    UsbIoHandleBuffer[Index],
                    &gEfiUsbIoProtocolGuid,
                    (VOID **) &UsbIo
                    );
      if (EFI_ERROR (Status)) {
        return FALSE;
      }
      Status = UsbIo->UsbGetInterfaceDescriptor (UsbIo, &IfDesc);
      if (EFI_ERROR (Status)) {
        return FALSE;
      }

      gBS->FreePool (TextStr);

      //check if the device is usb hid devices.
      if(IfDesc.InterfaceClass == USB_CLASS_HID){
        return TRUE;
      }else{
        return FALSE;
      }
    }else {
      gBS->FreePool (TextStr);
    }
  }

  return FALSE;
}
#endif

EFI_STATUS
CheckGraphicsOutputProtocol(
  VOID
  )
{
  EFI_HANDLE                    *HandleBuffer;
  UINTN                         NumberOfHandles;
  UINTN                         Index;
  EFI_DEVICE_PATH_PROTOCOL      *DevicePath;
  EFI_STATUS                    Status;
  EFI_GRAPHICS_OUTPUT_PROTOCOL  *GraphicsOutput;

  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiGraphicsOutputProtocolGuid,
                  NULL,
                  &NumberOfHandles,
                  &HandleBuffer
                  );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  for (Index = 0; Index < NumberOfHandles; Index++) {
    Status = gBS->HandleProtocol (
                    HandleBuffer[Index],
                    &gEfiDevicePathProtocolGuid,
                    (VOID*)&DevicePath
                    );
    if (EFI_ERROR (Status)) {
      continue;
    }

    Status = gBS->HandleProtocol (
                    HandleBuffer[Index],
                    &gEfiGraphicsOutputProtocolGuid,
                    (VOID **)&GraphicsOutput
                    );
    if (!EFI_ERROR (Status)) {
      break;
    }
  }

  gBS->FreePool (HandleBuffer);

  if (Index >= NumberOfHandles) {
    return EFI_NOT_FOUND;
  }

  return EFI_SUCCESS;
}
//#[-start-1701124-JESSE0143-Add]#
VOID
DisableEcPower (
  VOID
  )
{
  UINT8    Data;
  //Disable EC power
  EcRamRead(0xA1, &Data);
  Data &= (UINT8)~(BIT6|BIT7);  // Clear BIT6 and BIT7
  EcRamWrite(0xA1,Data);
}
//#[-end-1701124-JESSE0143-Add]#
EFI_STATUS
OneKeyBatterySetupCallBack (
  IN EFI_EVENT  Event,
  IN VOID       *Context
  )
{
  EFI_STATUS         Status = EFI_SUCCESS;

  UINTN              DisplayTime;
  UINT8              BatteryPercentage;
  BOOLEAN            ACStatus         = FALSE;
#if AOU_BOOTUP_SUPPORT
  BOOLEAN            AOUStatus        = FALSE;
#else
	//if attach usb to boot up, and AOU_BOOTUP_SUPPORT==0, then don't show OneKeyBattery.
  if(BootupState == Attach_AOU_Bootup_State){
    goto OneKeyBatteryExit;
  }
#endif

  //check if GraphicsOutputProtocol is avaiable.
  Status = CheckGraphicsOutputProtocol();
  if (EFI_ERROR (Status)) {
    return Status;
  }

#if AOU_BOOTUP_SUPPORT
  //if attach usb to boot up, and it's usb hid device, then don't show OneKeyBattery.
  if(BootupState == Attach_AOU_Bootup_State){
    if(CheckUsbHidDevice()){
      goto OneKeyBatteryExit;
    }
  }
#endif

  // Initializ graphics console
  //  gST->ConOut->ClearScreen (gST->ConOut);
  Status = InitializeGUI();
  if (EFI_ERROR (Status)) {
    goto OneKeyBatteryExit;
  }

  Status = GetACStatus(&ACStatus);
  if (EFI_ERROR(Status)) {
    goto OneKeyBatteryExit;
  }
#if AOU_BOOTUP_SUPPORT
  Status = GetAOUStatus(&AOUStatus);
  if (EFI_ERROR(Status)) {
    goto OneKeyBatteryExit;
  }
#endif

// Get Battery Percentage
  Status = GetBatteryPercentage(&BatteryPercentage);
  if (EFI_ERROR(Status)) {
     goto OneKeyBatteryExit;
  }
  if(BatteryPercentage > 100)
    BatteryPercentage = 100;

  // Draw Battery
  Status = DrawBattery(BatteryPercentage,ACStatus);
  if (EFI_ERROR (Status)) {
    goto OneKeyBatteryExit;
  }

  switch(BootupState){
      case Keyboard_Bootup_State://in this case,show static AC and static USB
      {
        //Draw Charger
        if (ACStatus == TRUE) {
          Status = DrawCharger(FALSE);
          if (EFI_ERROR (Status)) {
            goto OneKeyBatteryExit;
          }
        }
#if AOU_BOOTUP_SUPPORT
        //Draw USB device
        if (AOUStatus == TRUE) {
          Status = DrawAOUDevice(FALSE);
          if (EFI_ERROR (Status)) {
            goto OneKeyBatteryExit;
          }
        }
        break;
      }
      case Attach_AC_Bootup_State://in this case,show static USB, show dynamic AC in while loop
      {
        //Draw USB device
        if (AOUStatus == TRUE) {
          Status = DrawAOUDevice(FALSE);
          if (EFI_ERROR (Status)) {
            goto OneKeyBatteryExit;
          }
        }
        break;
      }
      case Attach_AOU_Bootup_State://in this case,show static AC, and show dynamic USB in while loop
      {
        //Draw Charger
        if (ACStatus == TRUE) {
          Status = DrawCharger(FALSE);
          if (EFI_ERROR (Status)) {
            goto OneKeyBatteryExit;
          }
        }
#endif
        break;
      }
  }

  /*
   *  When attach AC, the AC picture should be dynamic.
   *  When attach USB, the USB picture should be dynamic.
   */
  DisplayTime = 0;
  while(DisplayTime < SECONDS_OF_DISPLAY*3){
    switch(BootupState){
      case Attach_AC_Bootup_State:
      {
        //Draw Charger. in this case, do NOT check AC status.
//        if (ACStatus == TRUE) {
          Status = DrawCharger(TRUE);
          if (EFI_ERROR (Status)) {
            goto OneKeyBatteryExit;
          }
//        }
        break;
      }
#if AOU_BOOTUP_SUPPORT
      case Attach_AOU_Bootup_State:
      {
        //Draw USB device. in this case, do NOT check USB device status.
//        if (USBStatus == TRUE) {
          Status = DrawAOUDevice(TRUE);
          if (EFI_ERROR (Status)) {
            goto OneKeyBatteryExit;
          }
//        }
        break;
      }
#endif
    }

    //stall for 300ms
    gBS->Stall(1000*300);
    DisplayTime++;
  }
  gBS->Stall(1000*100*SECONDS_OF_DISPLAY);

OneKeyBatteryExit:
//#[-start-1701124-JESSE0143-Add]#
  DisableEcPower();
//#[-end-1701124-JESSE0143-Add]#
//#[-start-1701219-JESSE0146-Add]#
  IoWrite8 (LFC_CMOS_INDEX, LFC_WAKE_SRC_INDEX);
  IoWrite8 (LFC_CMOS_DATA, 0x01);
//#[-end-1701219-JESSE0146-Add]#
  gRT->ResetSystem(EfiResetShutdown, EFI_SUCCESS, 0, NULL);
  return Status;
}

/**
  The user Entry Point for Dxe Driver. The user code starts with this function
  as the real entry point for the Dxe Driver.

  @param[in] ImageHandle    The firmware allocated handle for the EFI image.
  @param[in] SystemTable    A pointer to the EFI System Table.

  @retval EFI_SUCCESS       The entry point is executed successfully.
  @retval other             Some error occurs when executing this entry point.

**/
EFI_STATUS
EFIAPI
OneKeyBatterySetupEntry (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS         Status;
  EFI_EVENT          Event;
  VOID               *Registration;

  Status = EcBootupState(&BootupState);
  if (EFI_ERROR(Status)) {
    return Status;
  }

  PcdSet8 (NotifyECBootUp, 0);

  if(BootupState != Normal_Bootup_State){
    PcdSet8 (NotifyECBootUp, 1);
    // Install notify to callback
    Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  OneKeyBatterySetupCallBack,
                  NULL,
                  &Event
                  );
    if (!EFI_ERROR (Status)) {
    	gBS->RegisterProtocolNotify (
                  &gEfiGraphicsOutputProtocolGuid,
                  Event,
                  &Registration
                  );
    }
  }

  return Status;
}