/** @file

;******************************************************************************
;* Copyright (c) 2014 - 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 <SmbiosUpdateDxe.h>
#include <MemoryConfig.h>
#include <Library/H2OCpLib.h>
#include <Guid/H2OCp.h>


EFI_SMBIOS_PROTOCOL        *Smbios;

/**
  Once DIMM settings keep in original comparing with the previous boot,
  the RC will ignore processing some SPD data into MEM_INFO_PROTOCOL during
  POST. This function used to restore SPD data from the variable service
  in order to ensure system can correctly and completely create the SMBIOS
  type 17 record.

  @param[in] *MemoryInfo     - A pointer points to the memory INFO data structure.
**/
VOID
UpdateSpdInfomation (
  IN MEM_INFO_PROTOCOL        *MemoryInfo
  )
{
  UINTN                       DataSize;
  EFI_STATUS                  Status;
  EFI_STATUS                  GetSpdVariableStatus;
  SPD_SAVE_DATA               *SpdSaveData;
  EFI_GUID                    gSpdSaveDataGuid = SPD_DATA_GUID;
  UINT8                       DimmNo;
  UINTN                       SmbusOffset;
  UINT8                       *OriginalRawData;
  UINT32                      ModuleSN = 0;
  UINT32                      TmpModuleSN = 0;
  BOOLEAN                     DimmSameAsPrevious;
  BOOLEAN                     SetVariableRequired;


  DataSize            = (sizeof (SPD_SAVE_DATA));
  SetVariableRequired = FALSE;
  DimmSameAsPrevious  = FALSE;
  SpdSaveData         = (SPD_SAVE_DATA *) AllocatePool (sizeof (SPD_SAVE_DATA));

  ASSERT (SpdSaveData != NULL);
  if (SpdSaveData == NULL){
    return;
  }

  //
  // Get the DIMM records of the previous boot.
  //
  GetSpdVariableStatus = gRT->GetVariable (L"SpdData",
                                           &gSpdSaveDataGuid,
                                           NULL,
                                           &DataSize,
                                           (VOID *)SpdSaveData
                                           );
  //
  // Add this fault-tolerance to prevent some unexpected errors occured.
  //
  if (GetSpdVariableStatus == EFI_BUFFER_TOO_SMALL) {
    if (SpdSaveData != NULL) {
      FreePool (SpdSaveData);
    }

    SpdSaveData = (SPD_SAVE_DATA *) AllocatePool (DataSize);
    ASSERT (SpdSaveData != NULL);
    if (SpdSaveData == NULL){
      return;
    }

    GetSpdVariableStatus = gRT->GetVariable (L"SpdData",
                                             &gSpdSaveDataGuid,
                                             NULL,
                                             &DataSize,
                                             (VOID *)SpdSaveData
                                             );
  }

  //
  // Process every DIMMs
  //
  for (DimmNo = 0; DimmNo < MAX_SOCKETS; DimmNo++) {
    //
    // To simply the declaration
    //
    OriginalRawData = MemoryInfo->MemInfoData.DimmsSpdData[DimmNo];
    //
    // Read the current Module Serial Number.
    //
    for (SmbusOffset = 122; SmbusOffset < 126; SmbusOffset++) {
      TmpModuleSN = OriginalRawData[SmbusOffset];
      if (SmbusOffset == 122) {
        ModuleSN = TmpModuleSN;
      } else {
        ModuleSN = (ModuleSN << 8) + (TmpModuleSN & 0xFF);
      }
    }

    //
    // Previous record found, check the DIMM in the slot is changed (update, remove or insert) or not
    //
    if (!EFI_ERROR (GetSpdVariableStatus)) {
      //
      // Compare ModuleSN with the previous boot
      //
      DimmSameAsPrevious = TRUE;
      if (ModuleSN != SpdSaveData->DimmData[DimmNo].ModuleSN) {
        DimmSameAsPrevious = FALSE;
      }
    }

    if (EFI_ERROR (GetSpdVariableStatus) || (!DimmSameAsPrevious)) {
      //
      // 1. It's the first boot (record not found) >> Create variable
      // 2. Memory module(s) is (are) changed      >> Update variable
      //
      SetVariableRequired = TRUE;
      SpdSaveData->DimmData[DimmNo].ModuleSN = ModuleSN;
      //
      // Sync with BDW 3.7 (only stores 144 Bytes)
      //
      CopyMem (SpdSaveData->DimmData[DimmNo].SmbusBufferData, OriginalRawData, 144);
    }
  }

  if (SetVariableRequired) {
    //
    // Module changed (update, remove or insert)
    //
    Status = gRT->SetVariable (L"SpdData",
                               &gSpdSaveDataGuid,
                               EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                               DataSize,
                               (VOID *) SpdSaveData
                               );
  } else {
    //
    // Module settings keep the same as the previous boot, some SPD data initializations will be skipped,
    // e.g. SPD_DDR3_MTBDD, SPD_DDR3_MTBDS, and SPD_DDR3_TCLK. So we use variable to restore them;
    // Otherwise, those fields will keep "0x00".
    //
    for (DimmNo = 0; DimmNo < MAX_SOCKETS; DimmNo++) {
      CopyMem (MemoryInfo->MemInfoData.DimmsSpdData[DimmNo],
               SpdSaveData->DimmData[DimmNo].SmbusBufferData,
               144
               );
    }
  }

  if (SpdSaveData != NULL) {
    FreePool (SpdSaveData);
  }
}

EFI_STATUS
SmbiosType20 (
  IN EFI_EVENT Event,
  IN VOID      *Context
  )
{
  //
  // SMBIOS_TABLE_TYPE20
  //
  SMBIOS_TABLE_TYPE20             *Type20Record;
  UINT8                            Type20RecordSize;

  EFI_STATUS                       Status;
  EFI_SMBIOS_HANDLE                SmbiosHandle;
  MEM_INFO_PROTOCOL               *MemInfoHob;
  SA_POLICY_PROTOCOL              *DxePlatformSaPolicy;
  UINT8                            ChannelASlotMap;
  UINT8                            ChannelBSlotMap;
  UINT8                            ControllerIndex;
  UINT8                            ChannelIndex;
  UINT8                            Dimm;
  UINT8                            DimmIndex;
  UINT8                            Index;
  UINT8                            Sockets;
  UINT64                           TotalMemorySizeInKB;
  UINT16                          *Record;
  UINTN                            HandleCount;
  UINT8                            BitIndex;
  MEMORY_DXE_CONFIG                *MemoryDxeConfig= NULL;

  Type20Record        = NULL;
  Type20RecordSize    = sizeof(SMBIOS_TABLE_TYPE20);

  MemInfoHob          = NULL;
  DxePlatformSaPolicy = NULL;
  ControllerIndex     = 0;
  ChannelIndex        = 0;
  DimmIndex           = 0;
  Index               = 0;
  Sockets             = 0;
  TotalMemorySizeInKB = 0;
  HandleCount         = 0;
  BitIndex            = 0;

  Status = gBS->LocateProtocol (&gMemInfoProtocolGuid, NULL, (VOID **) &MemInfoHob);
  Status = gBS->LocateProtocol (&gSaPolicyProtocolGuid, NULL, (VOID **) &DxePlatformSaPolicy);
  Status = GetConfigBlock ((VOID *)DxePlatformSaPolicy, &gMemoryDxeConfigGuid, (VOID *)&MemoryDxeConfig);
  ASSERT_EFI_ERROR (Status);


  ChannelASlotMap = MemoryDxeConfig->ChannelASlotMap;
  ChannelBSlotMap = MemoryDxeConfig->ChannelBSlotMap;
  
  for (ControllerIndex = 0; ControllerIndex < MEM_CFG_MAX_CONTROLLERS; ControllerIndex++) {    
    for (Dimm = 0; Dimm < (MEM_CFG_MAX_SOCKETS / MEM_CFG_MAX_CONTROLLERS); Dimm++) {
      ChannelIndex = Dimm >> 1;
      DimmIndex = Dimm & 0x1;

      Type20Record = (SMBIOS_TABLE_TYPE20 *) AllocateZeroPool (Type20RecordSize + 2);   //add double 0x00 at the end
      ASSERT (Type20Record != NULL);
      if (Type20Record == NULL) {
        return EFI_OUT_OF_RESOURCES;
      }

      if (MemInfoHob->MemInfoData.dimmSize[Index] > 0 && (MemInfoHob->MemInfoData.DimmStatus[Index] < DIMM_NOT_PRESENT)) {
        Type20Record->Hdr.Type                 = 20;
        Type20Record->Hdr.Length               = Type20RecordSize;
        Type20Record->StartingAddress          = 0;

        Type20Record->PartitionRowPosition = PcdGet8(PcdSmbiosType20PartitionRowPosition);
        GetLinkTypeHandle(EFI_SMBIOS_TYPE_MEMORY_DEVICE,&Record,&HandleCount);
        ASSERT (Record != NULL);
        if(Record == NULL) {
          return EFI_NOT_FOUND;
        }
        if (Sockets < HandleCount){
          Type20Record->MemoryDeviceHandle = Record[Sockets];
        }
        GetLinkTypeHandle(EFI_SMBIOS_TYPE_MEMORY_ARRAY_MAPPED_ADDRESS,&Record,&HandleCount);
        ASSERT (Record != NULL);
        if(Record == NULL) {
          return EFI_NOT_FOUND;
        }

        Type20Record->MemoryArrayMappedAddressHandle = Record[0]; //Record->Handle;
        SmbiosHandle = SMBIOS_HANDLE_PI_RESERVED;

        TotalMemorySizeInKB = LShiftU64 (MemInfoHob->MemInfoData.dimmSize[Index], 10);
        if (TotalMemorySizeInKB > 0xFFFFFFFF) {
          Type20Record->StartingAddress = 0xFFFFFFFF;
          Type20Record->EndingAddress = 0xFFFFFFFF;
          Type20Record->ExtendedEndingAddress = TotalMemorySizeInKB - 1;
        } else {
          Type20Record->EndingAddress = (UINT32) (TotalMemorySizeInKB - 1);
        }
        Type20Record->InterleavePosition    = (UINT8)(ChannelIndex + 1);
        Type20Record->InterleavedDataDepth  = (UINT8)(1);
        Status = Smbios->Add (Smbios, NULL, &SmbiosHandle,  (EFI_SMBIOS_TABLE_HEADER*) Type20Record);
        Sockets++;
      } else if ((BIT0 << DimmIndex) & ((ChannelIndex == 0) ? ChannelASlotMap : ChannelBSlotMap)) {
        Sockets++;
      }
      Index++;
    }
  }

  return EFI_SUCCESS;
}

/**
  Create an event to install a smbios type 14 (Group associations) data for the Intel firmware version info.

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

  @retval Status                 EFI_SUCCESS
  @return Assert, otherwise.

**/
EFI_STATUS
EFIAPI
SmbiosUpdateInit (
  IN      EFI_HANDLE              ImageHandle,
  IN      EFI_SYSTEM_TABLE        *SystemTable
 )
{
//  VOID                       *Registration;
  EFI_EVENT                  EndOfDxeEvent;
  EFI_STATUS                 Status;
  EFI_SMBIOS_TABLE_HEADER    *Record;
  UINT8                      Index;
  H2O_CP_HANDLE              CpHandle;

  Smbios           = NULL;
  Record           = NULL;
  Index            = 0;

  //
  // Find the SMBIOS Protocol
  //
  Status = gBS->LocateProtocol (&gEfiSmbiosProtocolGuid, NULL, (VOID **) &Smbios);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR | DEBUG_INFO, "Locate SmbiosProtocol failed in CreateFviGroupCallBack event.\n"));
    return EFI_NOT_FOUND;
  }

  //
  // Create Smbios Type 20 information
  //
//  Status = SmbiosType20(Smbios);
  Status = gBS->CreateEventEx (
                  EVT_NOTIFY_SIGNAL,
                  TPL_NOTIFY,
                  SmbiosType20,
                  NULL,
                  &gEfiEndOfDxeEventGroupGuid,
                  &EndOfDxeEvent
                  );
  ASSERT_EFI_ERROR (Status);
  
//  //
//  // Create an ExitPmAuth protocol callback event to generate SMBIOS table - Group Associations Type14
//  //
//  EfiCreateProtocolNotifyEvent (
//          &gExitPmAuthProtocolGuid,
//          TPL_CALLBACK - 1,
//          CreateType14FviGroupCallback,
//          NULL,
//          &Registration
//          );
  //
  // Register End of DXE event
  //
  if (FeaturePcdGet (PcdH2OBdsCpReadyToBootAfterSupported)) {
    Status  = H2OCpRegisterHandler (
                  &gH2OBdsCpReadyToBootAfterGuid,
                  AddSmbiosT14Callback,
                  H2O_CP_MEDIUM,
                  &CpHandle
                  );
    if (EFI_ERROR (Status)) {
      DEBUG_CP ((DEBUG_ERROR, "Checkpoint Register Fail: %g (%r)\n", &gH2OBdsCpReadyToBootAfterGuid, Status));
      ASSERT_EFI_ERROR (Status);
      return Status;
    }
    DEBUG_CP ((DEBUG_ERROR, "Checkpoint Registered: %g (%r)\n", &gH2OBdsCpReadyToBootAfterGuid, Status));
  }

  return EFI_SUCCESS;
}