/** @file

;******************************************************************************
;* Copyright (c) 2017, 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 "HddSpinDownDxe.h"

EFI_SMM_CONTROL2_PROTOCOL                    *mSmmControl = NULL;
#if (FixedPcdGetBool (PcdL05PchResetSupported))
PCH_RESET_PROTOCOL                           *mPchResetProtocol = NULL;
PCH_RESET                                    mOriginalPchReset;
#endif
EFI_RESET_SYSTEM                             mOriginalResetSystem;
EFI_VARIABLE_DEFAULT_UPDATE_PROTOCOL         *mVariableDefaultUpdateProtocol = NULL;
EFI_VARIABLE_DEFAULT_UPDATE_FACTORY_SETTING  mOriginalUpdateFactorySetting;
BOOLEAN                                      mPostponeReset = FALSE;
BOOLEAN                                      mHddSpinDownSwSmiIsReady = FALSE;

/**

  L05 Trigger HDD spin down.

  @param  None

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval Orther                        An unexpected error occurred.
**/
EFI_STATUS
L05TriggerHddSpinDown (
  )
{
  EFI_STATUS                            Status;
  UINT8                                 SmiDataValue;

  if (!mHddSpinDownSwSmiIsReady) {
    return EFI_UNSUPPORTED;
  }

  SmiDataValue = EFI_L05_HDD_SPIN_DOWN_CALLBACK;

  Status = mSmmControl->Trigger (
                          mSmmControl,
                          &SmiDataValue,
                          NULL,
                          0,
                          0
                          );

  return Status;
}

/**
  Hook UpdateFactorySetting() to set Postpone Reset.

  @param  This                          Pointer to EFI_VARIABLE_DEFAULT_UPDATE_PROTOCOL instance.
  @param  RestoreType                   Restore type to update for the variable store.

  @retval EFI_INVALID_PARAMETER         Input parameter is invalid.
  @retval EFI_SUCCESS                   Update system setting to factory default successful.
  @return Other                         Other error cause update system to factory default failed.
**/
EFI_STATUS
EFIAPI
HookUpdateFactorySetting (
  IN EFI_VARIABLE_DEFAULT_UPDATE_PROTOCOL  *This,
  IN UINT32                                RestoreType
  )
{
  EFI_STATUS                            Status;
  
  Status = mOriginalUpdateFactorySetting (This, RestoreType);

  if (!EFI_ERROR (Status)) {
    mPostponeReset = TRUE;
  }

  return Status;
}

/**
  Restore UpdateFactorySetting() from HookUpdateFactorySetting() to original UpdateFactorySetting().
  Set ResetType to PcdL05PostponeResetType for Postpone Reset.
  Clean Postpone Reset flag.

  @param  None

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval Orther                        An unexpected error occurred.
**/
EFI_STATUS
RestoreUpdateFactorySetting (
  IN EFI_RESET_TYPE                     ResetType
  )
{
  if (!mPostponeReset) {
    return EFI_UNSUPPORTED;
  }

  PcdSet8S (PcdL05PostponeResetType, ((UINT8) ResetType));
  mVariableDefaultUpdateProtocol->UpdateFactorySetting = mOriginalUpdateFactorySetting;
  mPostponeReset = FALSE;
  
  return EFI_SUCCESS;
}

/**
  Notification function of gEfiVariableDefaultUpdateProtocolGuid.

  This is a notification function registered on gEfiVariableDefaultUpdateProtocolGuid.
  It Hook UpdateFactorySetting() to Postpone Reset.

  @param  Event                         Event whose notification function is being invoked.
  @param  Context                       Pointer to the notification function's context.
**/
VOID
EFIAPI
VariableDefaultUpdateNotify (
  IN EFI_EVENT                          Event,
  IN VOID                               *Context
  )
{
  EFI_STATUS                            Status;

  Status = gBS->LocateProtocol (
                  &gEfiVariableDefaultUpdateProtocolGuid,
                  NULL,
                  (VOID **)&mVariableDefaultUpdateProtocol
                  );
  
  if (EFI_ERROR (Status)) {
    return;
  }

  mOriginalUpdateFactorySetting = mVariableDefaultUpdateProtocol->UpdateFactorySetting;
  mVariableDefaultUpdateProtocol->UpdateFactorySetting = HookUpdateFactorySetting;
}

/**
  Postpone Reset.
  Check PcdL05PostponeResetType to do Postpone Reset.

  @param  None

  @retval None
**/
VOID
PostponeReset (
  )
{
  if (PcdGet8 (PcdL05PostponeResetType) != 0xFF) {
    gRT->ResetSystem (((EFI_RESET_TYPE) PcdGet8 (PcdL05PostponeResetType)), EFI_SUCCESS, 0, NULL);
  }
}

/**
  Hdd Spin Down register DisplayBeforeCp to PostponeReset().

 @param[in] Event                       The Event this notify function registered to.
 @param[in] Handle                      The handle associated with a previously registered checkpoint handler.
**/
STATIC
VOID
EFIAPI
HddSpinDownDisplayBeforeCp (
  IN EFI_EVENT                          Event,
  IN H2O_CP_HANDLE                      Handle
  )
{
  PostponeReset ();
}

#if (FixedPcdGetBool (PcdL05PchResetSupported))
/**
  Trigge HDD spin down before PCH reset.

  @param  None

  @retval None
**/
EFI_STATUS
EFIAPI
L05HddSpinDownPchReset (
  IN PCH_RESET_PROTOCOL                 *This,
  IN EFI_RESET_TYPE                     ResetType,
  IN UINTN                              DataSize,
  IN VOID                               *ResetData OPTIONAL
  )
{
  EFI_STATUS                            Status;

  Status = RestoreUpdateFactorySetting (ResetType);

  if (!EFI_ERROR (Status)) {
    return Status;
  }

  if (ResetType != EfiResetWarm) {
    L05TriggerHddSpinDown ();
  }

  Status = mOriginalPchReset (This, ResetType, DataSize, ResetData);

  return Status;
}

/**
  L05 HDD spin down hook PCH reset.

  @param  None

  @retval None
**/
VOID
L05HddSpinDownHookPchReset (
  )
{
  mOriginalPchReset = mPchResetProtocol->Reset;
  mPchResetProtocol->Reset = L05HddSpinDownPchReset;
}
#endif

/**
  Trigge HDD spin down before gRT->ResetSystem.

  @param  None

  @retval None
**/
VOID
L05HddSpinDownResetSystem (
  IN EFI_RESET_TYPE                     ResetType,
  IN EFI_STATUS                         ResetStatus,
  IN UINTN                              DataSize,
  IN CHAR16                             *ResetData OPTIONAL
  )
{
  EFI_STATUS                            Status;

  Status = RestoreUpdateFactorySetting (ResetType);

  if (!EFI_ERROR (Status)) {
    return;
  }

  if (ResetType != EfiResetWarm) {
    L05TriggerHddSpinDown ();
  }

  mOriginalResetSystem (ResetType, ResetStatus, DataSize, ResetData);
}

/**
  L05 HDD spin down hook gRT->ResetSystem.

  @param  None

  @retval None
**/
VOID
L05HddSpinDownHookResetSystem (
  )
{
  mOriginalResetSystem = gRT->ResetSystem;
  gRT->ResetSystem = L05HddSpinDownResetSystem;
}

/**
  Notification function of EVT_SIGNAL_VIRTUAL_ADDRESS_CHANGE.

  This is a notification function registered on EVT_SIGNAL_VIRTUAL_ADDRESS_CHANGE event.
  It converts pointer to new virtual address.

  @param  Event                         Event whose notification function is being invoked.
  @param  Context                       Pointer to the notification function's context.
**/
VOID
EFIAPI
VirtualAddressChangeNotify (
  IN EFI_EVENT                          Event,
  IN VOID                               *Context
  )
{
  EfiConvertPointer (0x0, (VOID **) &mSmmControl);
#if (FixedPcdGetBool (PcdL05PchResetSupported))
  EfiConvertPointer (0x0, (VOID **) &mOriginalPchReset);
#endif
  EfiConvertPointer (0x0, (VOID **) &mOriginalResetSystem);
  EfiConvertPointer (0x0, (VOID **) &mHddSpinDownSwSmiIsReady);
}

VOID
EFIAPI
HddSpinDownSwSmiReadyProtocolNotify (
  IN  EFI_EVENT                         Event,
  IN  VOID                              *Context
  )
{
  gBS->CloseEvent (Event);

  mHddSpinDownSwSmiIsReady = TRUE;
}

/**
  HDD spin down 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
L05HddSpinDownDxeEntry (
  IN EFI_HANDLE                         ImageHandle,
  IN EFI_SYSTEM_TABLE                   *SystemTable
  )
{
  EFI_STATUS                            Status;
  EFI_EVENT                             VirtualAddressChangeEvent;
  H2O_CP_HANDLE                         CpHandle;
  EFI_VARIABLE_DEFAULT_UPDATE_PROTOCOL  *VariableDefaultUpdateProtocol;
  EFI_EVENT                             Event;
  VOID                                  *Registration;

  Status = EFI_SUCCESS;
  VirtualAddressChangeEvent = NULL;

  Status = gBS->LocateProtocol (
                  &gEfiSmmControl2ProtocolGuid,
                  NULL,
                  (VOID **) &mSmmControl
                  );

  if (EFI_ERROR (Status)) {
    return Status;
  }

#if (FixedPcdGetBool (PcdL05PchResetSupported))

  //
  // Hook PCH Reset
  //
  Status = gBS->LocateProtocol (
                  &gPchResetProtocolGuid,
                  NULL,
                  (VOID **) &mPchResetProtocol
                  );

  if (!EFI_ERROR (Status)) {
    L05HddSpinDownHookPchReset ();
  }
#endif

  //
  // Hook gRT->ResetSystem
  //
  L05HddSpinDownHookResetSystem ();

  //
  // Register the event to convert the pointer for runtime.
  //
  Status = gBS->CreateEventEx (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  VirtualAddressChangeNotify,
                  NULL,
                  &gEfiEventVirtualAddressChangeGuid,
                  &VirtualAddressChangeEvent
                  );



  Status = gBS->LocateProtocol (
                  &gEfiVariableDefaultUpdateProtocolGuid,
                  NULL,
                  (VOID **)&VariableDefaultUpdateProtocol
                  );
  
  if (!EFI_ERROR (Status)) {
    VariableDefaultUpdateNotify (NULL, NULL);
    
  } else {
    //
    // Registers the gEfiVariableDefaultUpdateProtocolGuid protocol notification for Hook UpdateFactorySetting()
    //
    Event        = NULL;
    Registration = NULL;
  
    Status = gBS->CreateEvent (
                    EVT_NOTIFY_SIGNAL,
                    TPL_NOTIFY,
                    VariableDefaultUpdateNotify,
                    NULL,
                    &Event
                    );
    
    if (!EFI_ERROR (Status)) {
      Status = gBS->RegisterProtocolNotify (
                      &gEfiVariableDefaultUpdateProtocolGuid,
                      Event,
                      &Registration
                      );
    }
  }

  //
  // Register notification on H2O_CP_LOW of gH2OBdsCpDisplayBeforeGuid event.
  //
  Status = H2OCpRegisterHandler (
             &gH2OBdsCpDisplayBeforeGuid,
             HddSpinDownDisplayBeforeCp,
             H2O_CP_LOW,
             &CpHandle
             );
  
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Checkpoint Register Fail: %g (%r)\n", &gH2OBdsCpDisplayBeforeGuid, Status));
    return Status;
  }

  //
  // Register the event to check HddSpinDown SW SMI is ready.
  //
  Event        = NULL;
  Registration = NULL;

  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  HddSpinDownSwSmiReadyProtocolNotify,
                  NULL,
                  &Event
                  );
  
  if (!EFI_ERROR (Status)) {
    Status = gBS->RegisterProtocolNotify (
                    &gEfiL05HddSpindownSwSmiReadyProtocolGuid,
                    Event,
                    &Registration
                    );
  }

  return Status;
}