#include <PiPei.h>
#include <Ppi/MeUma.h>  /* GUID-defined PPI produced by this driver */
#include <Library/DebugLib.h>
#include <Library/PeiServicesLib.h>
#include <Library/HobLib.h>
#include <Library/IoLib.h>
#include <Library/PcdLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/BaseLib.h>

//
// ---------------------------------------------------------------------------
// Globals
// ---------------------------------------------------------------------------
//
// The PPI descriptor (GUID + EFI_PEI_PPI_DESCRIPTOR) for the ME UMA PPI.
// Located in the .data section at 0xffda3390.
//
extern EFI_GUID  gMeUmaPpiGuid;
extern VOID      *gMeUmaPpi;

UINT8  byte_FFDA339C;

//
// ---------------------------------------------------------------------------
// Function prototypes
// ---------------------------------------------------------------------------
EFI_STATUS
EFIAPI
ModuleEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  );

//
// Internal functions (previously sub_*)
//

/**
  Reads 64 bits from an unaligned buffer.
**/
UINT64
EFIAPI
ReadUnaligned64 (
  IN CONST VOID  *Buffer
  )
{
  ASSERT (Buffer != NULL);
  return *(const UINT64 *)Buffer;
}

/**
  Writes 64 bits to an unaligned buffer.
**/
UINT64
EFIAPI
WriteUnaligned64 (
  OUT VOID   *Buffer,
  IN  UINT64 Value
  )
{
  ASSERT (Buffer != NULL);
  *(UINT64 *)Buffer = Value;
  return Value;
}

/**
  Reads 16-bit I/O port.
**/
UINT16
EFIAPI
IoRead16 (
  IN UINT16  Port
  )
{
  ASSERT ((Port & 1) == 0);
  return *(volatile UINT16 *)(UINTN)Port;
}

/**
  Writes 16-bit I/O port.
**/
UINT16
EFIAPI
IoWrite16 (
  IN UINT16  Port,
  IN UINT16  Value
  )
{
  ASSERT ((Port & 1) == 0);
  *(volatile UINT16 *)(UINTN)Port = Value;
  return Value;
}

/**
  Reads 32-bit I/O port (dword).
**/
UINT32
EFIAPI
IoRead32 (
  IN UINT16  Port
  )
{
  ASSERT ((Port & 3) == 0);
  return __indword (Port);
}

/**
  Reads a 16-bit value from PCI configuration space.
**/
UINT16
ReadPciCfg16 (
  IN UINT8  Bus,
  IN UINT8  Dev,
  IN UINT8  Func
  )
{
  PCI_CONFIG_ACCESS  PciCfg;
  PciCfg.Reg  = 0x80000000 | ((UINT32)Bus << 16) | ((UINT32)Dev << 11) | ((UINT32)Func << 8);
  PciCfg.Reserved = 0;
  //
  // The actual read goes through the memory-mapped ECAM or I/O CFG mechanism.
  // We construct the address and decode it via the PCI Express library.
  //
  UINTN  Address = (PciCfg.Reg | 0) & 0xFFF;  // simplified: see PciExpressLib
  // ...
  return 0;
}

//
// ---------------------------------------------------------------------------
// Internal Helpers
// ---------------------------------------------------------------------------

/**
  Returns the EFI_PEI_SERVICES pointer via the IDT-based trick.
**/
EFI_PEI_SERVICES **
GetPeiServices (
  VOID
  )
{
  IA32_DESCRIPTOR  Idtr;
  AsmReadIdtr (&Idtr);
  return *(EFI_PEI_SERVICES ***)((UINTN)Idtr.Base - 4);
}

/**
  Returns the PCD pointer (PeiPcdLib).
**/
VOID *
GetPcdPtr (
  VOID *This
  )
{
  VOID   *PcdDb;
  EFI_STATUS  Status;

  Status = PeiServicesLocatePpi (&gPeiPcdPpiGuid, 0, NULL, &PcdDb);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "\nASSERT_EFI_ERROR (Status = %r)\n", Status));
    ASSERT_EFI_ERROR (Status);
    return NULL;
  }
  return PcdDb;
}

/**
  Calls the PcdGetX PPI method (index 4).
**/
UINT32
PcdGet32 (
  VOID *This,
  UINTN  Token
  )
{
  PCD_PPI  *PcdPpi;

  PcdPpi = (PCD_PPI *)GetPcdPtr (This);
  return PcdPpi->Get32 (Token);
}

/**
  Calls the PcdGetSize PPI method (index 5).
**/
UINTN
PcdGetSize (
  VOID *This,
  UINTN  Token
  )
{
  PCD_PPI  *PcdPpi;

  PcdPpi = (PCD_PPI *)GetPcdPtr (This);
  return PcdPpi->GetSize (Token);
}

/**
  Returns the HOB list pointer.
**/
VOID *
GetHobList (
  VOID
  )
{
  EFI_PEI_SERVICES  **PeiServices;
  EFI_STATUS         Status;
  VOID               *HobList;

  PeiServices = GetPeiServices ();
  Status = (*PeiServices)->GetHobList (PeiServices, &HobList);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "\nASSERT_EFI_ERROR (Status = %r)\n", Status));
    ASSERT_EFI_ERROR (Status);
  }
  ASSERT (HobList != NULL);
  return HobList;
}

/**
  Locates the first HOB with type EFI_HOB_TYPE_GUID_EXT (0x0004)
  whose GUID matches 'This'.
**/
EFI_HOB_GUID_TYPE *
FindGuidHob (
  IN EFI_GUID  *Guid
  )
{
  VOID             *HobList;
  EFI_HOB_GUID_TYPE *Hob;

  HobList = GetHobList ();
  for (Hob = (EFI_HOB_GUID_TYPE *)HobList; Hob->Header.HobType != EFI_HOB_TYPE_END_OF_HOB_LIST;
       Hob = (EFI_HOB_GUID_TYPE *)((UINT8 *)Hob + Hob->Header.HobLength)) {
    if (Hob->Header.HobType == EFI_HOB_TYPE_GUID_EXT &&
        CompareGuid (&Hob->Name, Guid)) {
      return Hob;
    }
  }
  return NULL;
}

/**
  Copies a GUID from Source to Destination via the unaligned helpers.
**/
VOID *
CopyGuid (
  OUT VOID       *Destination,
  IN  CONST VOID *Source
  )
{
  WriteUnaligned64 (Destination, ReadUnaligned64 (Source));
  WriteUnaligned64 ((UINT8 *)Destination + 8, ReadUnaligned64 ((UINT8 *)Source + 8));
  return Destination;
}

/**
  Compares two GUIDs.
  Returns TRUE if they are equal.
**/
BOOLEAN
CompareGuid (
  IN CONST EFI_GUID  *Guid1,
  IN CONST EFI_GUID  *Guid2
  )
{
  UINT64  A_lo, A_hi, B_lo, B_hi;

  A_lo = ReadUnaligned64 (Guid1);
  A_hi = ReadUnaligned64 ((UINT8 *)Guid1 + 8);
  B_lo = ReadUnaligned64 (Guid2);
  B_hi = ReadUnaligned64 ((UINT8 *)Guid2 + 8);

  return (A_lo == B_lo && A_hi == B_hi);
}

/**
  Builds a HOB with GUID type.
**/
VOID *
BuildGuidHob (
  IN EFI_GUID  *Guid,
  IN UINTN      DataLength
  )
{
  EFI_HOB_GUID_TYPE  *Hob;
  VOID               *Buffer;
  EFI_STATUS          Status;
  EFI_PEI_SERVICES  **PeiServices;

  ASSERT (Guid != NULL);
  ASSERT (DataLength <= 0xFFE0);

  PeiServices = (EFI_PEI_SERVICES **)GetPeiServices ();
  Status = (*PeiServices)->CreateHob (
                              PeiServices,
                              EFI_HOB_TYPE_GUID_EXT,
                              (UINT16)(DataLength + sizeof (EFI_HOB_GUID_TYPE)),
                              &Hob
                              );
  if (EFI_ERROR (Status)) {
    return NULL;
  }
  ASSERT (Hob != NULL);
  CopyGuid (&Hob->Name, Guid);
  Buffer = Hob + 1;   // points to Hob + 24
  ZeroMem (Buffer, DataLength);
  return (VOID *)((UINT8 *)Hob + sizeof (EFI_HOB_GUID_TYPE));
}

//
// ---------------------------------------------------------------------------
// UMA Parameter Validation
// ---------------------------------------------------------------------------

/**
  Validates the stored UMA location parameters against the Silicon registers.

  @param[in]  MeNcMemLowBaseAddr   Low UMA base address
  @param[in]  MeNcMemHighBaseAddr  High UMA base address (extension)
  @param[in]  MeNcMemLowLimit      Low UMA limit address
  @param[in]  MeNcMemHighLimit     High UMA limit address

  @retval  EFI_SUCCESS            UMA parameters are valid.
  @retval  EFI_INVALID_PARAMETER  UMA address or size mismatch.
**/
EFI_STATUS
ValidateUmaLocation (
  IN UINT32  MeNcMemLowBaseAddr,
  IN UINT32  MeNcMemHighBaseAddr,
  IN UINT32  MeNcMemLowLimit,
  IN UINT32  MeNcMemHighLimit
  )
{
  UINT32      MeUmaSize;
  UINT32      MeUmaSizeCalc;
  UINT32      UmaBase;
  UINT32      UmaBaseExt;
  EFI_STATUS  Status;

  Status = EFI_SUCCESS;

  //
  // Check if the ME is in a debug or error mode that skips UMA checking.
  //
  if (!(UINT8)ModuleEntryPoint (NULL, NULL)) {
    return EFI_SUCCESS;
  }

  DEBUG ((EFI_D_INFO, "ME UMA:  ====================================================\n"));
  DEBUG ((EFI_D_INFO, "ME UMA:  Stored UMA Location:\n"));
  DEBUG ((EFI_D_INFO, "ME UMA:  MeNcMemLowBaseAddr  = 0x%x\n", MeNcMemLowBaseAddr));
  DEBUG ((EFI_D_INFO, "ME UMA:  MeNcMemHighBaseAddr = 0x%x\n", MeNcMemHighBaseAddr));
  DEBUG ((EFI_D_INFO, "ME UMA:  MeNcMemLowLimit     = 0x%x (->0x%x)\n", MeNcMemLowLimit, MeNcMemLowLimit));
  DEBUG ((EFI_D_INFO, "ME UMA:  MeNcMemHighLimit    = 0x%x\n", MeNcMemHighLimit));

  //
  // Calculate UMA size:
  //   MeUmaSize  = (MeNcMemLowLimit & 0xFFF80000) + 0x80000 - MeNcMemLowBaseAddr
  //   MeUmaSize += MeNcMemHighLimit
  //
  MeUmaSize    = (MeNcMemLowLimit & 0xFFF80000) + 0x80000 - MeNcMemLowBaseAddr;
  MeUmaSizeCalc = MeUmaSize + MeNcMemHighLimit;

  DEBUG ((EFI_D_INFO, "ME UMA:  MeUmaSize           = 0x%x (%dM)\n", MeUmaSizeCalc, MeUmaSizeCalc >> 20));

  //
  // Read the UMA parameters from Silicon registers (via PCI config access).
  //
  UmaBase    = *(UINT32 *)((UINTN)PciCfgRead (0) + 240);  // offset 0xF0
  UmaBaseExt = *(UINT32 *)((UINTN)PciCfgRead (0) + 244);  // offset 0xF4

  DEBUG ((EFI_D_INFO, "ME UMA:  UMA parameters read:\n"));
  DEBUG ((EFI_D_INFO, "ME UMA:  MeUmaBase           = 0x%x\n", UmaBase));
  DEBUG ((EFI_D_INFO, "ME UMA:  MeUmaBaseExt        = 0x%x\n", UmaBaseExt));

  //
  // If base addresses are zero, skip validation.
  //
  if (MeNcMemLowBaseAddr == 0 || MeNcMemLowLimit == 0) {
    if (MeUmaSizeCalc != 0) {
      DEBUG ((EFI_D_WARN, "ME UMA:  WARNING: Lower part of UMA addresses is 0.\n"));
    }
    goto Done;
  }

  //
  // Validate the UMA base addresses match Silicon registers.
  //
  if (MeNcMemLowBaseAddr != UmaBase || MeNcMemHighBaseAddr != UmaBaseExt) {
    DEBUG ((EFI_D_ERROR, "ME UMA:  ERROR: UMA addresses error.\n"));
    Status = EFI_INVALID_PARAMETER;
  }

  //
  // Validate UMA size: must not exceed 64 MB, must be even (2 MB aligned).
  //
  if (MeUmaSizeCalc >> 20 > 64) {
    DEBUG ((EFI_D_ERROR, "ME UMA:  ERROR: UMA size error #01.\n"));
    Status = EFI_INVALID_PARAMETER;
  } else if ((MeUmaSizeCalc & 1) != 0) {
    DEBUG ((EFI_D_ERROR, "ME UMA:  ERROR: UMA size error #02.\n"));
    Status = EFI_INVALID_PARAMETER;
  }

Done:
  DEBUG ((EFI_D_INFO, "ME UMA:  ====================================================\n"));
  return Status;
}

//
// ---------------------------------------------------------------------------
// ME Type Detection
// ---------------------------------------------------------------------------

/**
  Determines the onboard ME type based on HOB and FS (Firmware Status) info.

  @return  The ME type:
           1  = SPS (Server Platform Services)
           2  = Consumer/Client ME
           15 = Debug Mode (DFX)
           255 = Unknown or Error
**/
UINT8
GetOnBoardMeType (
  VOID
  )
{
  UINT32  MeFirmwareStatus;
  UINT8   MeOperationMode;

  //
  // Check Debug Mode via PCH PMC (DWR flow).
  //
  if (IsPchDwrFlow ()) {
    DEBUG ((EFI_D_INFO, "HECI: GetOnBoardMeType() for DWR flow return Dfx type\n"));
    return ME_TYPE_DFX;  // 15
  }

  //
  // Read the ME Firmware Status (MEFS) register from PCI config.
  //
  MeFirmwareStatus = *(UINT32 *)((UINTN)PciCfgRead (0) + 64);  // offset 0x40

  if (MeFirmwareStatus == 0xFFFFFFFF) {
    //
    // Fall back to HOB if MEFS is invalid.
    //
    MeFirmwareStatus = GetMeFs1FromHob ();
    DEBUG ((EFI_D_INFO, "HECI: GetOnBoardMeType() reads Hfs info from HOB = %d\n"));
  }

  //
  // Evaluate the ME type from the status.
  //
  if ((MeFirmwareStatus & 0xF) == 0xF) {
    return ME_TYPE_DFX;  // Debug/DFX
  }

  if ((MeFirmwareStatus & 0xF) == 4) {
    return 255;  // Unknown type
  }

  MeOperationMode = (UINT8)((MeFirmwareStatus >> 16) & 0xF);
  DEBUG ((EFI_D_INFO, "HECI: MeOperationMode = %d\n", MeOperationMode));

  switch (MeOperationMode) {
    case 0:
    case 1:
      return ME_TYPE_CLIENT;  // 2

    case 2:
      return 255;

    case 3:
    case 4:
    case 5:
      return 255;

    case 7:
      return 255;

    case 15:
      return ME_TYPE_SPS;  // 1

    default:
      DEBUG ((EFI_D_ERROR, "HECI: ME type not recognized (MEFS1: 0x%08X)\n", MeFirmwareStatus));
      DEBUG ((EFI_D_ERROR, "                             (MEFS2: 0x%08X)\n", PciCfgRead (0)));
      return 0;
  }
}

/**
  Reads ME Firmware Status 1 (MEFS1) from the appropriate HOB.

  @return  The MEFS1 value, or -1 on error.
**/
UINT32
GetMeFs1FromHob (
  VOID
  )
{
  EFI_HOB_GUID_TYPE  *Hob;
  UINT32              Result;

  Result = (UINT32)-1;

  Hob = FindGuidHob (&gMeFwHobGuid);
  if (Hob == NULL || ((ME_FW_HOB *)GET_GUID_HOB_DATA (Hob))->Group[0].FunNumber != HECI1_DEVICE) {
    //
    // Could not read HOB, use fallback path.
    //
    DEBUG ((EFI_D_ERROR, "HECI: GetMeFs1FromHob() Can't read correctly MeFwHob info\n"));
  } else {
    Result = ((ME_FW_HOB *)GET_GUID_HOB_DATA (Hob))->Group[0].FunNumber;
  }

  DEBUG ((EFI_D_INFO, "HECI: GetMeFs1FromHob() returns MEFS1 = %d\n", Result));
  return Result;
}

//
// ---------------------------------------------------------------------------
// Performance Logging (PEI Performance)
// ---------------------------------------------------------------------------

/**
  Gets the PEI performance log and ID arrays from HOBs (or builds them).

  @param[out]  PeiPerformanceLog       The performance log array
  @param[out]  PeiPerformanceIdArray   The performance ID array
**/
VOID
GetPeiPerformance (
  OUT UINT32  **PeiPerformanceLog,
  OUT UINT32  **PeiPerformanceIdArray
  )
{
  EFI_HOB_GUID_TYPE  *GuidHob;

  ASSERT (PeiPerformanceLog != NULL);
  ASSERT (PeiPerformanceIdArray != NULL);

  GuidHob = FindGuidHob (&gEfiPeiPerformanceHobGuid);
  if (GuidHob != NULL) {
    //
    // Use existing HOB data.
    //
    *PeiPerformanceLog = (UINT32 *)((UINT8 *)GET_GUID_HOB_DATA (GuidHob));
    GuidHob = FindGuidHob (&gEfiPeiPerformanceIdArrayGuid);
    ASSERT (GuidHob != NULL);
    *PeiPerformanceIdArray = (UINT32 *)((UINT8 *)GET_GUID_HOB_DATA (GuidHob));
  } else {
    //
    // Build fresh HOBs for the performance log and ID array.
    //
    *PeiPerformanceLog = BuildGuidHob (&gEfiPeiPerformanceHobGuid, 40008);
    *PeiPerformanceLog = (UINT32 *)ZeroMem (*PeiPerformanceLog, 40008);
    *PeiPerformanceIdArray = BuildGuidHob (&gEfiPeiPerformanceIdArrayGuid, 4000);
    ZeroMem (*PeiPerformanceIdArray, 4000);
  }
}

/**
  Finds the performance log entry index for a given module ID.

  @param[in]  PerformanceLog  The performance log array
  @param[in]  ModuleId        The module identifier to find

  @return  The index of the found entry, or the array length if not found.
**/
UINT32
FindPerformanceLogIndex (
  IN UINT32  *PerformanceLog,
  IN UINT32   ModuleId
  )
{
  UINT32  i;
  UINT32  Count;

  Count = PerformanceLog[0];
  for (i = 0; i < Count; i++) {
    UINT32  Index = Count - i - 1;
    UINT32  *Entry = &PerformanceLog[10 * Index];

    if (Entry[0] == 0 && Entry[2] == ModuleId && Entry[3] == 0 &&
        !AsciiStrnCmp ("", "", 7) && !AsciiStrnCmp ("", "", 7)) {
      return Index;
    }
  }
  return i;
}

//
// ---------------------------------------------------------------------------
// HECI Message / PPI Installation
// ---------------------------------------------------------------------------

/**
  Checks if the current ME type prevents HECI communication.
**/
BOOLEAN
IsHeciSkipped (
  VOID
  )
{
  return GetOnBoardMeType () == 255;
}

/**
  Locates the HECI PPI and sends a message to query UMA configuration.

  @param[in]      Message       The HECI message buffer
  @param[out]     Response      The HECI response buffer

  @retval EFI_SUCCESS           Message sent and response received.
  @retval Others                Error
**/
EFI_STATUS
SendHeciMessage (
  IN  UINT8   *Message,
  OUT UINT8   *Response
  )
{
  EFI_STATUS                Status;
  EFI_PEI_PPI_DESCRIPTOR    *HeciPpi;
  HECI_PPI                  *Heci;

  Status = PeiServicesLocatePpi (&gHeciPpiGuid, 0, NULL, &HeciPpi);
  if (EFI_ERROR (Status)) {
    return EFI_UNSUPPORTED;
  }
  Heci = (HECI_PPI *)HeciPpi;

  return Heci->SendMessage (Heci, Message, Response);
}

/**
  Micro-delay using I/O port 0x508 polling.

  @param[in]  Microseconds  The number of microseconds to delay.
**/
VOID
MicroDelay (
  IN UINT32  Microseconds
  )
{
  UINT32  Count;
  UINT32  Base;
  UINT32  Current;

  Count = Microseconds >> 22;
  Base  = Microseconds & 0x3FFFFF;

  do {
    Current = Base + (IoRead32 (0x508) & 0xFFFFFF);
    Base = 0x400000;
    while (((Current - IoRead32 (0x508)) & 0x800000) == 0) {
      _mm_pause ();
    }
    Count--;
  } while (Count != 0xFFFFFFFF);
}

//
// ---------------------------------------------------------------------------
// Entry Point
// ---------------------------------------------------------------------------

/**
  ME UMA PPI Driver Entry Point.

  Installs the ME UMA PPI that exposes the UMA location information
  to other PEIMs. Validates the UMA parameters from Silicon registers.

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

  @retval EFI_SUCCESS           The PPI was installed successfully.
  @retval EFI_INVALID_PARAMETER UMA parameters are inconsistent.
  @retval Others                Error from PPI installation.
**/
EFI_STATUS
EFIAPI
ModuleEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS         Status;
  EFI_PEI_SERVICES **PeiServices;

  //
  // One-time initialization of the PCI Express register.
  //
  if ((IoRead32 (1024068) & 0x80) == 0) {
    InitializePciExpress ();
    IoWrite16 ((UINT16 *)(UINTN)(GetPcdPtr (NULL) + 1024068), 0x500);
    *(UINT8 *)(GetPcdPtr (NULL) + 1024068) |= 0x80;
  }

  DEBUG ((EFI_D_INFO, "ME UMA:  ME UMA PPI Driver EntryPoint\n"));

  //
  // Check if the system has valid ME firmware (skip on debug/error).
  //
  byte_FFDA339C = (IsHeciSkipped () < 0) ? 0 : byte_FFDA339C;

  //
  // Install the ME UMA PPI.
  //
  PeiServices = GetPeiServices ();
  Status = (*PeiServices)->InstallPpi (
                              PeiServices,
                              &gMeUmaPpiDescriptor
                              );
  DEBUG ((EFI_D_INFO, "ME UMA:  ME UMA PPI Installation status %r\n", Status));

  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "\nASSERT_EFI_ERROR (Status = %r)\n", Status));
    ASSERT_EFI_ERROR (Status);
  }

  return Status;
}