/** @file
  One Key Battery DXE Driver.

;******************************************************************************
;* Copyright (c) 2020, Insyde Software Corp. 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 Insyde Software Corporation.
;*
;******************************************************************************
*/

#include "OneKeyBatteryDxe.h"
//[-start-210702-Dongxu0008-add]// 
#ifdef LCFC_SUPPORT
#include <Library/LfcEcLib.h>
#include <LfcCmos.h>
#endif
//[-end-210702-Dongxu0008-add]// 
//[-start-211112-OWENWU0022-add]//
#if defined(S370_SUPPORT) || defined(S570_SUPPORT)
#include <Protocol/AtaPassThru.h>
#include <Protocol/ResetNotification.h>
#include <SetupVariable.h>
#include <Library/VariableLib.h>
#endif
//[-end-211112-OWENWU0022-add]// 
BOOTUP_STATE_FOR_SHOW_ONEKEYBATTERY     mBootupState = Normal_Bootup_State;
BOOLEAN                                 mAouStatus   = FALSE;

/**
  Get Battery Percentage.

  @param  BatteryPercentage             The Battery Percentage.
  
  @retval EFI_SUCCESS                   Executed successfully.
  @retval other                         Error occurs.

**/
EFI_STATUS
GetBatteryPercentage (
  IN OUT UINT8                          *BatteryPercentage
  )
{
  EFI_STATUS                            Status;

  Status = EFI_UNSUPPORTED;

  Status = OemSvcGetBatteryPercentage (BatteryPercentage);

  return Status;
}

/**
  Get AC state.

  @param  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;

  Status = EFI_UNSUPPORTED;

  //
  // If BootupState is Attach_AC_Bootup_State, directly assume that the AC is plugged in.
  //
  if (mBootupState == Attach_AC_Bootup_State) {
    *AcStatus = TRUE;
    return EFI_SUCCESS;
  }

  Status = OemSvcGetAcStatus (AcStatus);

  return Status;
}

/**
  Get USB state.

  @param  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
  )
{
  EFI_STATUS                            Status;

  Status = EFI_UNSUPPORTED;

  Status = OemSvcGetAouStatus (AouStatus);

  return Status;	
}

/**
  Get Boot up state.

  @param  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 *BootupState
  )
{
  EFI_STATUS                              Status;

  Status = EFI_UNSUPPORTED;

  Status = OemSvcGetBootUpState (BootupState);

  return Status;	
}

//[-start-211112-OWENWU0022-add]//

#if defined(S370_SUPPORT) || defined(S570_SUPPORT) 
/**
  Wrapper for EFI_ATA_PASS_THRU_PROTOCOL.PassThru().

  This function wraps the PassThru() invocation for ATA pass through function
  for an ATA device. It assembles the ATA pass through command packet for ATA
  transaction.

  @param         AtaPassThruPtr        .
  @param[in]     Port                  The port number of the ATA device to send
                                       the command.
  @param[in]     PortMultiplierPort    The port multiplier port number of the ATA
                                       device to send the command.
                                       If there is no port multiplier, then specify
                                       0xFFFF.

  @return The return status from EFI_ATA_PASS_THRU_PROTOCOL.PassThru().

**/
EFI_STATUS
EnableVMDAHCISpinDown (
  IN EFI_ATA_PASS_THRU_PROTOCOL         *AtaPassThruPtr,
  IN UINT16                              Port,
  IN UINT16                              PortMultiplierPort
  )
{
  EFI_STATUS                         Status;
  EFI_ATA_STATUS_BLOCK               Asb;
  EFI_ATA_COMMAND_BLOCK              Acb;
  EFI_ATA_PASS_THRU_COMMAND_PACKET   AtaPassThruCmdPacket;

  Status = EFI_SUCCESS;

  //
  // Initial Command packet
  //
  ZeroMem (&AtaPassThruCmdPacket, sizeof (EFI_ATA_PASS_THRU_COMMAND_PACKET));
  ZeroMem (&Asb, sizeof (EFI_ATA_STATUS_BLOCK));
  ZeroMem (&Acb, sizeof (EFI_ATA_COMMAND_BLOCK));
  AtaPassThruCmdPacket.Asb = &Asb;
  AtaPassThruCmdPacket.Acb = &Acb;

  //
  // Fill in command packet
  //  
  AtaPassThruCmdPacket.Acb->AtaCommand       = 0xE0; // STANDBY_IMMEDIATE_CMD_ALIAS
  AtaPassThruCmdPacket.Timeout               = 100000000;
  AtaPassThruCmdPacket.Protocol              = EFI_ATA_PASS_THRU_PROTOCOL_ATA_NON_DATA;
  AtaPassThruCmdPacket.Length                = EFI_ATA_PASS_THRU_LENGTH_BYTES;

  //
  // Processing AtaPassThru command
  // 
  Status = AtaPassThruPtr->PassThru (
                          AtaPassThruPtr,
                          Port,
                          PortMultiplierPort,
                          &AtaPassThruCmdPacket,
                          NULL
                          );
  return Status;
}

/**
  Calling this routine can make the SATA device enter standby mode correctly .

  @param[in]  ResetType         The type of reset to perform.
  @param[in]  ResetStatus       The status code for the reset.
  @param[in]  DataSize          The size, in bytes, of ResetData.
  @param[in]  ResetData         For a ResetType of EfiResetCold, EfiResetWarm, or
                                EfiResetShutdown the data buffer starts with a Null-terminated
                                string, optionally followed by additional binary data.
                                The string is a description that the caller may use to further
                                indicate the reason for the system reset.
                                For a ResetType of EfiResetPlatformSpecific the data buffer
                                also starts with a Null-terminated string that is followed
                                by an EFI_GUID that describes the specific type of reset to perform.
**/
VOID
EFIAPI
VMDDeviceSpinDownCallback (
  )
{
  UINTN                      HandleNum;
  EFI_HANDLE                 *AtaPassThruHandles;
  EFI_ATA_PASS_THRU_PROTOCOL *AtaPassThruPtr;
  UINTN                      Index;
  UINT16                     Port;
  UINT16                     PortMultiplierPort;
  EFI_STATUS                 Status;

  
  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiAtaPassThruProtocolGuid,
                  NULL,
                  &HandleNum,
                  &AtaPassThruHandles
                  );

  if (EFI_ERROR (Status)) {
    return ;
  }

  for (Index = 0; Index < HandleNum; Index++) {

    Status = gBS->HandleProtocol (
                    AtaPassThruHandles[Index],
                    &gEfiAtaPassThruProtocolGuid,
                    (VOID **) &AtaPassThruPtr
                    );
    if (EFI_ERROR (Status)) {
      return ;
    }

    Port = 0xFFFF;
    while (TRUE) {
      Status = AtaPassThruPtr->GetNextPort (AtaPassThruPtr, &Port);
      if (EFI_ERROR (Status)) {
        if (Status == EFI_NOT_FOUND) {
          break;
        } else {
          return ;
        }
      }

      PortMultiplierPort = 0xFFFF;
      while (TRUE) {
        Status = AtaPassThruPtr->GetNextDevice (AtaPassThruPtr, Port, &PortMultiplierPort);
        if (EFI_ERROR (Status)) {

          if (Status == EFI_NOT_FOUND) {
            break;
          } else {
            return;
          }
        }
       // ATA command STANDBY_IMMEDIATE_CMD_ALIAS
       Status = EnableVMDAHCISpinDown (AtaPassThruPtr, Port, PortMultiplierPort); 
      }
    }
  }
  return ;
}
#endif
//[-end-211112-OWENWU0022-add]//

EFI_STATUS
OneKeyBatterySetupCallBack (
  IN  EFI_EVENT                         Event,
  IN  H2O_CP_HANDLE                     Handle
  )

{ 
  EFI_STATUS                            Status;
  UINTN                                 DisplayTime;  
  UINT8                                 BatteryPercentage;
  BOOLEAN                               AcStatus;
  EFI_STATUS                            AouSupportStatus;

  Status            = EFI_SUCCESS;
  DisplayTime       = 0;
  BatteryPercentage = 0;
  AcStatus          = FALSE;
  AouSupportStatus  = EFI_UNSUPPORTED;

  H2OCpUnregisterHandler (Handle);

  //
  // If attach usb to boot up, and AOU_BOOTUP_SUPPORT==0, then don't show OneKeyBattery.
  //
  AouSupportStatus = GetAouStatus (&mAouStatus);

  if(EFI_ERROR (AouSupportStatus) && mBootupState == Attach_AOU_Bootup_State) {
    goto OneKeyBatteryExit;
  }

  Status = InitializeGUI ();

  if (EFI_ERROR (Status)) {
    goto OneKeyBatteryExit;
  }

  Status = GetAcStatus (&AcStatus);

  if (EFI_ERROR (Status)) {
    goto OneKeyBatteryExit;
  }

  //
  // 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 (mBootupState) {

  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;
      }
    }

    //
    // Draw USB device.
    //
    if (!EFI_ERROR (AouSupportStatus) && mAouStatus == TRUE) {
      Status = DrawAOUDevice (FALSE);
      if (EFI_ERROR (Status)) {
        goto OneKeyBatteryExit;
      }
    }

    break;

  //
  // In this case, show static USB, show dynamic AC in while loop.
  //
  case Attach_AC_Bootup_State:

    //
    // Draw USB device.
    //
    if (!EFI_ERROR (AouSupportStatus) && mAouStatus == TRUE) {
      Status = DrawAOUDevice (FALSE);
      if (EFI_ERROR (Status)) {
        goto OneKeyBatteryExit;
      }
    }
    break;

  //
  // In this case, show static AC, and show dynamic USB in while loop.
  //
  case Attach_AOU_Bootup_State:

    //
    // Draw Charger.
    //
    if (!EFI_ERROR (AouSupportStatus) && AcStatus == TRUE) {
      Status = DrawCharger (FALSE);
      if (EFI_ERROR (Status)) {
        goto OneKeyBatteryExit;
      }
    }

    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 (mBootupState) {

    case Attach_AC_Bootup_State:

      //
      // Draw Charger. In this case, do NOT check AC status.
      //
      Status = DrawCharger (TRUE);
      if (EFI_ERROR (Status)) {
        goto OneKeyBatteryExit;
      }

      break;

    case Attach_AOU_Bootup_State:

      if (!EFI_ERROR (AouSupportStatus)) {

        //
        // Draw USB device. In this case, do NOT check USB device status.
        //
        Status = DrawAOUDevice (TRUE);
        if (EFI_ERROR (Status)) {
          goto OneKeyBatteryExit;
        }
      }
      break;
    }

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

  gBS->Stall (1000 * 100 * SECONDS_OF_DISPLAY);  
  
OneKeyBatteryExit:
//[-start-210702-Dongxu0008-add]// 
#ifdef LCFC_SUPPORT
  DisableEcPower();
  IoWrite8 (LFC_CMOS_INDEX, LFC_WAKE_SRC_INDEX);
  IoWrite8 (LFC_CMOS_DATA, 0x01);
#endif
//[-start-211112-OWENWU0022-add]//
#if defined(S370_SUPPORT) || defined(S570_SUPPORT)
  VMDDeviceSpinDownCallback();
#endif

//[-end-211112-OWENWU0022-add]//
//[-end-210702-Dongxu0008-add]// 
  gRT->ResetSystem (
         EfiResetShutdown,
         EFI_SUCCESS,
         0,
         NULL
         );
  return Status;
}

/**
  This is the callback function after gEfiSetupUtilityProtocolGuid is installed.
  It will check SYSTEM_CONFIGURATION.L05OneKeyBattery to init One Key Battery.

  @param  Event                         Event whose notification function is being invoked.
  @param  Context                       Pointer to the notification function's context.
**/
//[-start-210702-Dongxu0008-remove]// 
//VOID
//EFIAPI
//OneKeyBatteryCallback (
//  IN  EFI_EVENT                         Event,
//  IN  VOID                              *Context
//  )
//{
//  EFI_STATUS                            Status;
//  VOID                                  *SetupUtility;
//  SYSTEM_CONFIGURATION                  *SetupNvData;
//  BOOLEAN                               IsOneKeyBatteryEnable;
//  H2O_CP_HANDLE                         CpHandle;
//
//  SetupUtility          = NULL;
//  SetupNvData           = NULL;
//  IsOneKeyBatteryEnable = FALSE;
//
//  Status = gBS->LocateProtocol (
//                  &gEfiSetupUtilityProtocolGuid,
//                  NULL,
//                  (VOID **) &SetupUtility
//                  );
//
//  if (EFI_ERROR (Status)) {
//    return;
//  }
//
//  gBS->CloseEvent (Event);
//
//  SetupNvData = CommonGetVariableData (SETUP_VARIABLE_NAME, &gSystemConfigurationGuid);
//
//  if (SetupNvData == NULL) {
//    return;
//  }
//
//  if ((SetupNvData->L05OneKeyBatterySupport == 0x01) && (SetupNvData->L05OneKeyBattery == 0x01)) { // 0:Disable, 1:Enable
//    IsOneKeyBatteryEnable = TRUE;
//  }
//
//  if (SetupNvData != NULL) {
//    FreePool (SetupNvData);
//  }
//
//  if (!IsOneKeyBatteryEnable) {
//    return;
//  }
//
//  Status = EcBootupState (&mBootupState);
//
//  if (EFI_ERROR (Status)) {
//    return;
//  }
//
//  if(mBootupState != Normal_Bootup_State) {
//    Status = H2OCpRegisterHandler (
//               &gH2OBdsCpConOutConnectAfterGuid,
//               OneKeyBatterySetupCallBack,
//               H2O_CP_MEDIUM,
//               &CpHandle
//               );
//  }      
//
//  return;
//}
//[-end-210702-Dongxu0008-add]// 

/**
  One Key Battery DXE Entry.

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

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
EFIAPI
OneKeyBatteryDxeEntryPoint (
  IN  EFI_HANDLE                        ImageHandle,
  IN  EFI_SYSTEM_TABLE                  *SystemTable
  )
{
//[-start-210702-Dongxu0008-modify]// 
//  VOID                                  *Registration;
//
//  Registration = NULL;
//
//  EfiCreateProtocolNotifyEvent (
//    &gEfiSetupUtilityProtocolGuid,
//    TPL_NOTIFY,
//    OneKeyBatteryCallback,
//    NULL,
//    &Registration
//    );
//
//  return EFI_SUCCESS;
  EFI_STATUS                            Status;
  H2O_CP_HANDLE                         CpHandle;

//[-start-211112-OWENWU0022-add]//
#if defined(S370_SUPPORT) || defined(S570_SUPPORT) 
  SA_SETUP                              SaSetup;
  UINTN                                 VarSize = 0;
#endif
//[-end-211112-OWENWU0022-add]//
  Status = EcBootupState (&mBootupState);

  if (EFI_ERROR (Status)) {
    return Status;
  }
  
  if(mBootupState != Normal_Bootup_State) {
//[-start-211112-OWENWU0022-add//
#if defined(S370_SUPPORT) || defined(S570_SUPPORT)
    VarSize = sizeof(SA_SETUP);
    Status = CommonGetVariable(
               SA_SETUP_VARIABLE_NAME,
               &gSaSetupVariableGuid,
               &VarSize,
               &SaSetup
               );
    Status = H2OCpRegisterHandler (
               &gH2OBdsCpDisplayBeforeGuid,
               OneKeyBatterySetupCallBack,
               H2O_CP_MEDIUM,
               &CpHandle
               );
#else
    Status = H2OCpRegisterHandler (
               &gH2OBdsCpConOutConnectAfterGuid,
               OneKeyBatterySetupCallBack,
               H2O_CP_MEDIUM,
               &CpHandle
               );
#endif
//[-end-211112-OWENWU0022-add]//
  }      
  
  return Status;
//[-end-210702-Dongxu0008-modify]//   
}