/** @file
  The instance of trace hub post code handler

;******************************************************************************
;* Copyright (c) 2017 - 2019, 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 <Uefi.h>
#include <PiPei.h>
#include <Library/PeiServicesLib.h>
#include <Library/PeimEntryPoint.h>
#include <Library/DebugLib.h>
#include <Guid/H2OCp.h>
#include <Library/DebugLib.h>
#include <Library/PrintLib.h>
#include <Library/TraceHubHookLib.h>
#include <Library/TraceHubDebugLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/H2OCpLib.h>
#include <Library/HobLib.h>

STATIC EFI_GUID mBaseCpPostCodeTraceHubGuid = {0x8d39831e, 0xd253, 0x4f61, {0xb1, 0x45, 0x73, 0x10, 0x1e, 0x0c, 0x14, 0x20}};
/**
 Write post code value to trace hub.

 @param[in] Event                      A pointer to the Event that triggered the callback.
 @param[in] Handle                     Checkpoint handle.
**/
VOID
EFIAPI
PostCodeCallback (
  IN EFI_EVENT                         Event,
  IN H2O_CP_HANDLE                     Handle
  )
{
  EFI_STATUS                           Status;
  H2O_BASE_CP_POST_CODE_DATA           *PostCodeData;
  CHAR8                                AsciiStrBuffer[50];
  UINTN                                CharCount;

  Status = H2OCpLookup (Handle, (VOID **) &PostCodeData, NULL);
  if (EFI_ERROR (Status)) {
    DEBUG_CP ((DEBUG_ERROR, "Checkpoint Data Not Found: %x (%r)\n", Handle, Status));
    DEBUG_CP ((DEBUG_ERROR, "   %a\n", __FUNCTION__));
    return;
  }

  CharCount = AsciiSPrint (
                AsciiStrBuffer,
                sizeof (AsciiStrBuffer),
                "TraceHub : POSTCODE=<%02x>\n\r",
                PostCodeData->PostCodeValue
                );
  TraceHubDebugWrite (SeverityError, AsciiStrBuffer, CharCount);
}

/**
 Register post code checkpoint.

 @param[in] ImageHandle               A handle for the image that is initializing this driver
 @param[in] SystemTable               A pointer to the EFI system table

 @retval EFI_SUCCESS                  Module initialized successfully
 @retval Others                       Fail to register checkpoint.
**/
EFI_STATUS
EFIAPI
TraceHubPostCodeHandlerPeiEntry (
  IN       EFI_PEI_FILE_HANDLE         FileHandle,
  IN CONST EFI_PEI_SERVICES            **PeiServices
  )
{
  EFI_STATUS              Status;
  H2O_CP_HANDLE           CpHandle;
  EFI_BOOT_MODE           BootMode;
  EFI_HOB_GUID_TYPE       *GuidHob;
  H2O_CP_HANDLE           *GuidCpHandle;

  if (!FeaturePcdGet (PcdH2OBaseCpPostCodeSupported)) {
    return EFI_SUCCESS;
  }

  Status = PeiServicesRegisterForShadow (FileHandle);
  if (Status == EFI_ALREADY_STARTED) {
    Status = PeiServicesGetBootMode (&BootMode);
    ASSERT_EFI_ERROR (Status);

    if (BootMode == BOOT_IN_RECOVERY_MODE) {
      //
      // Unregister the CP handle in flash part
      //
      GuidHob  = GetFirstGuidHob (&mBaseCpPostCodeTraceHubGuid);
      if (GuidHob != NULL) {
        GuidCpHandle = (H2O_CP_HANDLE *) GET_GUID_HOB_DATA (GuidHob);
        Status = H2OCpUnregisterHandler (*GuidCpHandle);
      }
      //
      // Register new CP handle in memory.
      //
      Status = H2OCpRegisterHandler (
                 &gH2OBaseCpPostCodeGuid,
                 PostCodeCallback,
                 H2O_CP_MEDIUM,
                 &CpHandle
                 );
      if (EFI_ERROR (Status)) {
        DEBUG_CP ((DEBUG_ERROR, "Checkpoint Register Fail: %g (%r)\n", &gH2OBaseCpPostCodeGuid, Status));
        return Status;
      }
      DEBUG_CP ((DEBUG_INFO, "Checkpoint Registered: %g (%r)\n", &gH2OBaseCpPostCodeGuid, Status));
    }
  } else {
    Status = H2OCpRegisterHandler (
                    &gH2OBaseCpPostCodeGuid,
                    PostCodeCallback,
                    H2O_CP_MEDIUM,
                    &CpHandle
                    );
    if (EFI_ERROR (Status)) {
      DEBUG_CP ((DEBUG_ERROR, "Checkpoint Register Fail: %g (%r)\n", &gH2OBaseCpPostCodeGuid, Status));
      return Status;
    }
    DEBUG_CP ((DEBUG_INFO, "Checkpoint Registered: %g (%r)\n", &gH2OBaseCpPostCodeGuid, Status));
    //
    // Store the CP handle in GUID HOB
    //
    GuidCpHandle  = BuildGuidHob (&mBaseCpPostCodeTraceHubGuid, sizeof (H2O_CP_HANDLE));
    *GuidCpHandle = CpHandle;
  }
  return Status;
}