/** @file
  Maintain PEI Timer Interrupt in DXE phase.
;******************************************************************************
;* Copyright (c) 2015 - 2020, Insyde Software 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 Insyde Software Corporation.
;*
;******************************************************************************
*/
#include <ExtendPeiTimer.h>

EFI_LEGACY_8259_PROTOCOL             *mLegacy8259  = NULL;
EFI_EVENT                            mTimerEvent   = NULL;
UINT32                               mPeiService;
UINT32                               mFunctionPoint;
UINT32                               mPeiServicesPoint;
ALIGN_16BYTE_BOUNDRY IA32_DESCRIPTOR mIdtr;

/**
  Migrate IDT table and Save PEI Service Table Pointer.
  The PEI Services Table pointer is stored in the 4 bytes 
  immediately preceding the Interrupt Descriptor Table (IDT) in memory.
    
  @param  None.

**/
VOID
MigrateIdtTable (
  VOID
  )
{   
  EFI_STATUS            Status;
  X64_IDT_TABLE         *IdtTable;
  EFI_PHYSICAL_ADDRESS  PageAddress;
 
  IdtTable = NULL;
  //
  // Allocate New IDT table
  //
  AsmReadIdtr (&mIdtr);
  Status = gBS->AllocatePages (
                  AllocateAnyPages,
                  EfiBootServicesData,
                  EFI_SIZE_TO_PAGES(mIdtr.Limit + 1 + sizeof (UINT32) * 2),
                  &PageAddress
                  );
  ASSERT (Status == EFI_SUCCESS);
  if (!EFI_ERROR(Status)) {
    //
    // Copy IDT Table and save PeiServicePointer
    //
    IdtTable             = (X64_IDT_TABLE*)(UINTN)PageAddress;
    IdtTable->Reserved   = 0;    
    IdtTable->PeiService = mPeiService;
    
    CopyMem ((VOID*)&IdtTable->IdtTable, (VOID*)mIdtr.Base, mIdtr.Limit + 1);
    mIdtr.Base = (UINTN)&IdtTable->IdtTable;
    AsmWriteIdtr (&mIdtr);    
  }  
}

/**
  Legacy8259Protocol installed callback function.
  Open IRQ0 interrupt again.
    
  @param[in]  Event                 Wait Event
  @param[in]  Context               Passed parameter to event handler

**/
VOID
Legacy8259CallbackFunction  (
  IN EFI_EVENT        Event,
  IN VOID             *Context
  )
{   
  DEBUG ((DEBUG_INFO, "PeiTimer: Legacy8259 Callback\n"));
  gBS->LocateProtocol (&gEfiLegacy8259ProtocolGuid, NULL, (VOID**)&mLegacy8259); 
  //
  // Open the IRQ0 interrupt
  //
  mLegacy8259->EnableIrq (mLegacy8259, Efi8259Irq0, FALSE);
}

/**

  CpuArchProtocol installed callback function.
  Clear pending interrupt. 
  
  @param[in]  Event                 Wait Event
  @param[in]  Context               Passed parameter to event handler

**/
VOID
CpuCallbackFunction  (
  IN EFI_EVENT        Event,
  IN VOID             *Context
  )
{
  DEBUG ((DEBUG_INFO, "PeiTimer: CpuArchDxe Callback\n"));
  if (mLegacy8259 != NULL) {
    //
    // Disabled IRQ0 interrupts
    //
    mLegacy8259->DisableIrq (mLegacy8259, Efi8259Irq0);   
    //
    // Clear pending interrupt
    //
    mLegacy8259->EndOfInterrupt (mLegacy8259, Efi8259Irq0); 
  }

  MigrateIdtTable ();
}

/**

  Change cpu mode to IA-32 and execution PEI function.

  @param[in]  Event                 Wait Event
  @param[in]  Context               Passed parameter to event handler

**/
VOID
TimerCallbackFunction  (
  IN EFI_EVENT        Event,
  IN VOID             *Context
  )
{
  DEBUG ((DEBUG_INFO, "PeiTimer: Dxe Timer Event Callback\n"));

  if (mFunctionPoint != 0) {
    Thunk64To32 (mPeiServicesPoint, mFunctionPoint);
  }
  ASSERT (mFunctionPoint != 0);
}

/**

  CpuArchProtocol installed callback function.

  @param[in] Event          A pointer to the Event that triggered the callback.
  @param[in] Handle         Checkpoint handle.

**/
VOID
EnterBdsCallbackFunction  (
  IN EFI_EVENT                        Event,
  IN H2O_CP_HANDLE                    Handle
  )
{
  DEBUG ((DEBUG_INFO, "PeiTimer: EnterBdsCallback\n"));
  if (mTimerEvent != NULL) {
    gBS->CloseEvent (mTimerEvent);
  }
  ASSERT (mTimerEvent != NULL);
}

/**
  ExtendPeiTimer driver entry point function.

  @param[in] ImageHandle          Image handle for this driver image
  @param[in] SystemTable          Pointer to the EFI System Table

  @retval EFI_SUCCESS             The driver installed/initialized correctly.

**/
EFI_STATUS
ExtendPeiTimerEntry (
  IN EFI_HANDLE               ImageHandle,
  IN EFI_SYSTEM_TABLE         *SystemTable
  )
{
  EFI_STATUS                 Status;
  EFI_EVENT                  Event8259; 
  EFI_EVENT                  EventCpu; 
  VOID                       *Registration8259;
  VOID                       *RegistrationCpu;
  EFI_HOB_GUID_TYPE          *GuidHob;
  H2O_PEI_TIMER_DATA_HOB     *DataInHob;
  H2O_CP_HANDLE              CpHandle;
  
  //
  // Get Thunk Entry point and Pei Services Point.
  //
  GuidHob            = NULL;
  mFunctionPoint     = 0;
  mPeiServicesPoint  = 0;
  GuidHob = GetFirstGuidHob (&gH2OPeiTimerDataHobGuid);
  ASSERT (GuidHob != NULL);
  
  if (GuidHob != NULL) {
    DataInHob         = GET_GUID_HOB_DATA (GuidHob);
    mFunctionPoint    = DataInHob->CallBackFunction;
    mPeiServicesPoint = DataInHob->PeiServicesPoint;
    DEBUG ((DEBUG_INFO, "PeiTimer: Get H2OPeiTimerDataHob\n"));
    DEBUG ((DEBUG_INFO, "  FunctionPoint    = 0x%X\n",mFunctionPoint));
    DEBUG ((DEBUG_INFO, "  PeiServicesPoint = 0x%X\n",mPeiServicesPoint));
  }
  
  //
  // Save the PeiServicePointer
  //
  AsmReadIdtr (&mIdtr);
  mPeiService = (UINT32)*(UINT32*)(mIdtr.Base - PEI_POINTER_SIZE);
 
  //
  // Register Event for Monitor gEfiLegacy8259ProtocolGuid is installed.
  //
  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  Legacy8259CallbackFunction,
                  NULL,
                  &Event8259
                  );
  if(!EFI_ERROR(Status)) {
    Status = gBS->RegisterProtocolNotify (
                    &gEfiLegacy8259ProtocolGuid,
                    Event8259,
                    &Registration8259
                    );
  }
  
  //
  // Register Event for Monitor gEfiCpuArchProtocolGuid is installed.
  //
  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  CpuCallbackFunction,
                  NULL,
                  &EventCpu
                  );
  if(!EFI_ERROR(Status)) {
    Status = gBS->RegisterProtocolNotify (
                    &gEfiCpuArchProtocolGuid,
                    EventCpu,
                    &RegistrationCpu
                    );
  }
  if (FeaturePcdGet (PcdH2OBdsCpInitSupported)) {
    Status = H2OCpRegisterHandler (
               &gH2OBdsCpInitGuid,
               EnterBdsCallbackFunction,
               H2O_CP_MEDIUM,
               &CpHandle
               );
    if (EFI_ERROR (Status)) {
      DEBUG_CP ((DEBUG_ERROR, "Checkpoint Register Fail: %g (%r)\n", &gH2OBdsCpInitGuid, Status));
      return Status;
    }
    DEBUG_CP ((DEBUG_INFO, "Checkpoint Registered: %g (%r)\n", &gH2OBdsCpInitGuid, Status));
  }
  //
  // Register Timer Event for do thunk.
  //
  Status = gBS->CreateEvent (
                EVT_TIMER | EVT_NOTIFY_SIGNAL,
                TPL_CALLBACK,
                TimerCallbackFunction,
                NULL,
                &mTimerEvent
                );
  if(!EFI_ERROR(Status)) {
    Status = gBS->SetTimer (
                    mTimerEvent,
                    TimerPeriodic,
                    DEFAULT_TIMER_TICK_DURATION
                    ); 
  }          
  return Status;
}