/** @file
  AmiCpuPeiPreMem.c

  Recovered AmiCpuPeiPreMem module body.

  Module-owned flow stays here. Copied helper bodies live in
  AmiCpuPeiPreMemHelpers.c. Shared platform accessors such as GetPcdPpi are
  treated as external dependencies, not local recovery targets.
**/

#include "AmiCpuPeiPreMem.h"

#include <PiPei.h>

//
// Copied helper implementations are provided by AmiCpuPeiPreMemHelpers.c.
// GetPcdPpi remains an imported dependency boundary.
//
VOID *EFIAPI GetPcdPpi (VOID *PeiServices);
UINT64 EFIAPI ReadUnaligned64 (CONST VOID *Buffer);
BOOLEAN EFIAPI CompareGuid (UINTN Guid1, UINTN Guid2);
INTN EFIAPI CpuId (INTN FunctionId, UINT32 *CpuidEaxOut);
UINT64 EFIAPI GetMask23 (VOID);
UINT64 EFIAPI RShift18 (UINTN LowBits, UINT32 HighBits);
VOID *EFIAPI GetPeiServicesTablePointer (VOID);

#define PEI_FFS_TYPE_FIRMWARE   0x01
#define PEI_FFS_ALIGNMENT       7U
#define PEI_FFS_ALIGN_MASK      0xFFFFFFF8U
#define PEI_FFS_SIZE_MASK       0x00FFFFFFU
#define PEI_FFS_HEADER_SIZE     24U
#define PEI_MAX_DEP_COUNT       0x14U
#define PEI_DEFAULT_FFS_SIZE    2048U

//
// Local function prototypes.
//
static INT32 EFIAPI
InitializePatchRegionFromPcd (
  IN VOID *PeiServices
  );

static INT32 EFIAPI
PeiFindFvByGuid (
  IN INT32       FvBaseAddr,
  OUT UINT32    *FvGuidValOut,
  OUT INT32     *SizeOutPtr
  );

//
// UEFI/PEI module entry point thunk.
//
EFI_STATUS
EFIAPI
ModuleEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  (VOID)ImageHandle;
  (VOID)SystemTable;

  return (EFI_STATUS)InitializePatchRegionFromPcd (NULL);
}

//
// Module body: discover patch region values from PCD PPI and, if missing,
// locate the fallback FV by GUID.
//
static INT32 EFIAPI
InitializePatchRegionFromPcd (
  IN VOID *PeiServices
  )
{
  VOID    *PcdPpi;
  UINT64   PatchRegionAddress;
  UINT64   PatchRegionSize;
  INT32    Status;

  PcdPpi            = GetPcdPpi (GetPeiServicesTablePointer ());
  PatchRegionAddress = ((UINT64 (EFIAPI *)(INT32))((VOID **)PcdPpi)[4]) (113);
  PcdPpi            = GetPcdPpi (GetPeiServicesTablePointer ());
  PatchRegionSize    = ((UINT64 (EFIAPI *)(INT32))((VOID **)PcdPpi)[4]) (114);

  if (PatchRegionSize == 0) {
    Status = PeiFindFvByGuid (
               (INT32)-15663104,
               (UINT32 *)&PatchRegionAddress,
               (INT32 *)&PatchRegionSize
               );
    if (Status < 0) {
      Status = PeiFindFvByGuid (
                 (INT32)-4194304,
                 (UINT32 *)&PatchRegionAddress,
                 (INT32 *)&PatchRegionSize
                 );
    }

    if (Status >= 0) {
      PcdPpi = GetPcdPpi (PeiServices);
      ((VOID (EFIAPI *)(INT32, UINT32, UINT32))((VOID **)PcdPpi)[18]) (
        113,
        (UINT32)PatchRegionAddress,
        (UINT32)(PatchRegionAddress >> 32)
        );

      PcdPpi = GetPcdPpi (PeiServices);
      ((VOID (EFIAPI *)(INT32, UINT32, UINT32))((VOID **)PcdPpi)[18]) (
        114,
        (UINT32)PatchRegionSize,
        (UINT32)(PatchRegionSize >> 32)
        );
    }
  }

  return 0;
}

//
// Locate a firmware volume by GUID; when scanning the default base address,
// apply CPUID affinity checks before accepting the section payload.
//
static INT32 EFIAPI
PeiFindFvByGuid (
  IN INT32   FvBaseAddr,
  OUT UINT32 *FvGuidValOut,
  OUT INT32  *SizeOutPtr
  )
{
  UINT32          FvEnd;
  UINT32          FvStartCandidate;
  UINT32          FvStart;
  UINT32          FfsSize;
  UINT64          MsrValue;
  UINT64          Mask;
  UINT32          CpuClusterId;
  UINT32          ScanLimit;
  UINT32          SectionAddr;
  UINT32          SectionCount;
  UINT32          SectionSize;
  UINT32          DependencyCount;
  UINT32          DependencyAddr;
  UINT8           DependencyIndex;
  UINT32          NextSectionAddr;
  UINT8           ClusterAffinityMask;
  INT32           CpuidEaxValue;

  static const UINT32 TargetGuid[4] = {
    386434418U,
    1156527999U,
    (UINT32)-1615835505,
    1889552127U
  };

  FvEnd = (UINT32)(*(UINT32 *)(UINTN)(FvBaseAddr + 32) + (UINTN)FvBaseAddr - 1U);
  if (*(UINT32 *)(UINTN)(FvBaseAddr + 40) != 0x4856465Fu) {
    return (INT32)0x80000012;
  }

  FvStartCandidate = (UINT32)((UINTN)FvBaseAddr + *(UINT16 *)(UINTN)(FvBaseAddr + 48));
  if (FvEnd <= FvStartCandidate) {
    return (INT32)0x80000012;
  }

  if (*(UINT16 *)(UINTN)(FvBaseAddr + 52) != 0) {
    FvStartCandidate = (UINT32)(
      (UINTN)FvBaseAddr +
      *(UINT16 *)(UINTN)(FvBaseAddr + 52) +
      *(UINT32 *)(UINTN)((UINTN)FvBaseAddr + *(UINT16 *)(UINTN)(FvBaseAddr + 52) + 16)
      );
  }

  FvStart = (FvStartCandidate + PEI_FFS_ALIGNMENT) & PEI_FFS_ALIGN_MASK;
  if (FvStart >= FvEnd) {
    return (INT32)0x80000012;
  }

  while (TRUE) {
    FfsSize = *(UINT32 *)(UINTN)(FvStart + 20) & PEI_FFS_SIZE_MASK;
    if (CompareGuid ((UINTN)FvStart, (UINTN)TargetGuid)) {
      break;
    }

    if ((UINT32)(FfsSize + FvStart) > FvStart) {
      FvStart = (FfsSize + FvStart + PEI_FFS_ALIGNMENT) & PEI_FFS_ALIGN_MASK;
      if (FvStart < FvEnd) {
        continue;
      }
    }

    return (INT32)0x80000012;
  }

  if (FvBaseAddr != (INT32)-15663104) {
    *FvGuidValOut = FvStart + PEI_FFS_HEADER_SIZE;
    *SizeOutPtr = (INT32)FfsSize;
    FvGuidValOut[1] = 0;
    SizeOutPtr[1] = 0;
    return 0;
  }

  MsrValue = __readmsr (0x17u);
  Mask = GetMask23 ();
  CpuClusterId = (UINT32)RShift18 (
                          (UINTN)((UINT32)MsrValue & (UINT32)~(UINT32)Mask),
                          (UINT32)((UINT32)(MsrValue >> 32) & (UINT32)~(UINT32)(Mask >> 32))
                          );
  ClusterAffinityMask = (UINT8)(1U << CpuClusterId);
  ScanLimit = FvStart + FfsSize;

  CpuId ((INTN)CpuClusterId, (UINT32 *)&CpuidEaxValue);
  SectionAddr = FvStart + PEI_FFS_HEADER_SIZE;
  if (*(UINT32 *)(UINTN)SectionAddr != PEI_FFS_TYPE_FIRMWARE) {
    return (INT32)0x80000012;
  }

  while (TRUE) {
    SectionCount = *(UINT32 *)(UINTN)(SectionAddr + 28);
    if (SectionCount != 0) {
      SectionSize = *(UINT32 *)(UINTN)(SectionAddr + 32);
    } else {
      SectionSize = PEI_DEFAULT_FFS_SIZE;
    }

    if ((*(UINT32 *)(UINTN)(SectionAddr + 12) == (UINT32)CpuidEaxValue) &&
        ((ClusterAffinityMask & *(UINT8 *)(UINTN)(SectionAddr + 24)) != 0)) {
      *FvGuidValOut = SectionAddr;
      FvGuidValOut[1] = 0;
      *SizeOutPtr = (INT32)SectionSize;
      SizeOutPtr[1] = 0;
      return 0;
    }

    if (*(UINT32 *)(UINTN)(SectionAddr + 32) > (UINT32)(SectionCount + 48)) {
      break;
    }

    NextSectionAddr = (SectionSize + 15U) & 0xFFFFFFF0U;
    if (SectionAddr + NextSectionAddr < ScanLimit) {
      SectionAddr += NextSectionAddr;
      if (*(UINT32 *)(UINTN)SectionAddr == PEI_FFS_TYPE_FIRMWARE) {
        continue;
      }
    }

    return (INT32)0x80000012;
  }

  DependencyCount = *(UINT32 *)(UINTN)(SectionCount + SectionAddr + 48);
  DependencyAddr = SectionCount + SectionAddr + 68;
  if (DependencyCount >= PEI_MAX_DEP_COUNT) {
    return (INT32)0x80000012;
  }

  DependencyIndex = 0;
  if (DependencyCount == 0) {
    ScanLimit = FvStart + FfsSize;
  } else {
    while ((*(UINT32 *)(UINTN)DependencyAddr != (UINT32)CpuidEaxValue) ||
           ((ClusterAffinityMask & *(UINT8 *)(UINTN)(DependencyAddr + 4)) == 0)) {
      DependencyAddr += 12;
      if ((UINT8)(++DependencyIndex) >= DependencyCount) {
        ScanLimit = FvStart + FfsSize;
        goto ScanNextSection;
      }
    }
  }

  *FvGuidValOut = SectionAddr;
  FvGuidValOut[1] = 0;
  *SizeOutPtr = (INT32)SectionSize;
  SizeOutPtr[1] = 0;
  return 0;

ScanNextSection:
  NextSectionAddr = (SectionSize + 15U) & 0xFFFFFFF0U;
  if (SectionAddr + NextSectionAddr < ScanLimit) {
    SectionAddr += NextSectionAddr;
    if (*(UINT32 *)(UINTN)SectionAddr == PEI_FFS_TYPE_FIRMWARE) {
      goto ScanNextSection;
    }
  }

  return (INT32)0x80000012;
}