/** @file
  OpromUpdateDxeLightningRidgeEXRP.c

  UBA (Universal Board Architecture) PCIe Slot Number Update Driver
  for the Lightning Ridge EXRP platform.

  This DXE driver registers a UBA protocol callback that configures
  PCIe slot number ranges for the Lightning Ridge EXRP platform.
  It also initializes UBA debug printing and HOB list access.

  Module: OpromUpdateDxeLightningRidgeEXRP.efi
  SHA256: 046d27c7db21d4c86a2f1d80f3399783a4d9eca143cfface4109a20de4cf771a
  Build: e:\hs\Build\HR6N0XMLK\DEBUG_VS2015\X64\PurleyRpPkg\Uba\UbaMain\Dxe\TypeLightningRidgeEXRP\OpromUpdateDxe\OpromUpdateDxe\DEBUG\OpromUpdateDxeLightningRidgeEXRP.pdb
**/

#include "OpromUpdateDxeLightningRidgeEXRP.h"

//
// UBA Protocol GUID: E03E0D46-5263-4845-B0A4-58D57B3177E2
//
extern EFI_GUID gUbaProtocolGuid;

//
// PCI Root Bridge IO Protocol GUID: 2F707EBB-4A1A-11D4-9A38-0090273FC14D
//
extern EFI_GUID gEfiPciRootBridgeIoProtocolGuid;

//
// HOB List GUID: 7739F24C-93D7-11D4-9A3A-0090273FC14D
//
extern EFI_GUID gEfiHobListGuid;

//
// UBA Config Board Protocol GUID: 36232936-0E76-31C8-A13A-3AF2FC1C3932
// Used for board-specific configuration retrieval
//
extern EFI_GUID gUbaConfigBoardProtocolGuid;

//
// Lightning Ridge EXRP SKU GUID: 371BD79C-DE79-4C5F-AA2B-BC9EBEFA988F
//
extern EFI_GUID gLightningRidgeExrpSkuGuid;

//
// Globals
//
EFI_HANDLE          gImageHandle  = NULL;
EFI_SYSTEM_TABLE   *gST           = NULL;
EFI_BOOT_SERVICES  *gBS           = NULL;
EFI_RUNTIME_SERVICES *gRT         = NULL;
VOID               *gUbaProtocol  = NULL;
VOID               *gHobList      = NULL;
UINT8               gBoardType    = 0;

/**
  Read an unaligned 64-bit value from memory.

  @param[in] Buffer  Pointer to the 64-bit value (may be unaligned).

  @return The 64-bit value read from Buffer.
**/
UINT64
EFIAPI
ReadUnaligned64 (
  IN CONST UINT64 *Buffer
  )
{
  ASSERT (Buffer != NULL);
  return *Buffer;
}

/**
  Compare two GUIDs for equality using unaligned 64-bit reads.

  Compares the UBA_CONFIG_BOARD_INFO_GUID (0xD10) against a HOB entry GUID.
  The GUID at 0xD10 is the EFI_HOB_LIST_GUID for identifying the HOB list
  pointer from the system table's HOB GUID array.

  @param[in] Guid1  Reference GUID (internal board info GUID).
  @param[in] Guid2  HOB entry GUID to compare.

  @retval TRUE   GUIDs match.
  @retval FALSE  GUIDs do not match.
**/
BOOLEAN
EFIAPI
HobCompareGuid (
  IN EFI_GUID *Guid1,
  IN EFI_GUID *Guid2
  )
{
  UINT64 Front1;
  UINT64 Front2;
  UINT64 Back1;
  UINT64 Back2;

  Front1 = ReadUnaligned64 ((UINT64 *)Guid1);
  Front2 = ReadUnaligned64 ((UINT64 *)Guid2);
  Back1  = ReadUnaligned64 ((UINT64 *)((UINT8 *)Guid1 + 8));
  Back2  = ReadUnaligned64 ((UINT64 *)((UINT8 *)Guid2 + 8));

  return (BOOLEAN)(Front1 == Front2 && Back1 == Back2);
}

/**
  Locate and cache the UBA protocol instance.

  Allocates and immediately frees a small pool (31 bytes) as a heuristic check
  to ensure boot services are functional. If the allocation returns <= 0x10,
  boot services are assumed healthy and LocateProtocol is called.
  Results are cached in gUbaProtocol for subsequent calls.

  @return Pointer to the UBA protocol interface, or NULL if not found.
**/
VOID *
GetUbaProtocol (
  VOID
  )
{
  UINT64   HobSize;

  if (gUbaProtocol != NULL) {
    return gUbaProtocol;
  }

  HobSize = (UINT64)(UINTN)gBS->AllocatePool (31);
  gBS->FreePool ((VOID *)(UINTN)HobSize);

  if (HobSize <= 0x10) {
    if (EFI_ERROR (gBS->LocateProtocol (&gUbaProtocolGuid, NULL, &gUbaProtocol))) {
      gUbaProtocol = NULL;
    }
  }

  return gUbaProtocol;
}

/**
  Print a debug message through the UBA protocol.

  Reads the board type from CMOS RTC register 0x4B (offset 0x70/0x71).
  Board type determines the debug mask:
  - Board type 1: debug mask 0x80000004
  - Board types 2-254: debug mask 0x80000006
  - Board types 0, >254: filtering disabled
  If board type is 0, a fallback MMIO read at 0xFDAF0490 is used.

  The debug message is only forwarded to Uba->DebugPrint if the
  DebugMask & ErrorLevel is non-zero.

  @param[in]  ErrorLevel  Debug error level mask.
  @param[in]  Format      Format string.
  @param[in]  ...         Variable arguments.
**/
VOID
EFIAPI
UbaDebugPrint (
  IN UINTN   ErrorLevel,
  IN CHAR8   *Format,
  ...
  )
{
  UBA_PROTOCOL  *Uba;
  UINT64        DebugMask;
  UINT8         CmosData;
  UINT8         BoardType;
  VA_LIST       Va;

  Uba = (UBA_PROTOCOL *)GetUbaProtocol ();
  DebugMask = 0;
  if (Uba == NULL) {
    return;
  }

  //
  // Read board type from CMOS index 0x4B
  //
  CmosData = IoRead8 (0x70);
  IoWrite8 (0x70, CmosData & 0x80 | 0x4B);
  BoardType = IoRead8 (0x71);

  if (BoardType > 3) {
    if (BoardType == 0) {
      //
      // Fallback: read from MMIO 0xFDAF0490 (LPC/storage config)
      //
      BoardType = (MmioRead8 (0xFDAF0490) & 2) | 1;
    }
  }

  if (BoardType >= 1 && BoardType <= 0xFE) {
    if (BoardType == 1) {
      DebugMask = 0x80000004;
    } else {
      DebugMask = 0x80000006;
    }
  }

  if (DebugMask & ErrorLevel) {
    VA_START (Va, Format);
    Uba->DebugPrint (ErrorLevel, Format, Va);
    VA_END (Va);
  }
}

/**
  Print an ASSERT message through the UBA protocol.

  Forwards the file name, line number, and failed condition string
  to the UBA protocol's AssertPrint handler.

  @param[in]  FileName   Source file name of the assertion.
  @param[in]  LineNumber Line number of the assertion.
  @param[in]  Condition  The failed condition expression.
**/
VOID
EFIAPI
UbaAssertPrint (
  IN CONST CHAR8  *FileName,
  IN UINTN        LineNumber,
  IN CONST CHAR8  *Condition
  )
{
  UBA_PROTOCOL *Uba;

  Uba = (UBA_PROTOCOL *)GetUbaProtocol ();
  if (Uba != NULL) {
    Uba->AssertPrint ((CHAR8 *)FileName, (UINTN)LineNumber, (CHAR8 *)Condition);
  }
}

/**
  Get the HOB list pointer by locating the HOB GUID.

  Searches through the system table's HOB GUID array (at SystemTable + 0x68,
  count at +0x68, array at +0x70) for a matching EFI_HOB_LIST_GUID.
  The HOB GUID entries are 24 bytes each (GUID at +0, data pointer at +0x10).
  Results are cached in gHobList for subsequent calls.

  On failure, prints an ASSERT via UBA debug/assert infrastructure.

  @return Pointer to the HOB list data, or NULL if not found.
**/
VOID *
HobGetGuid (
  VOID
  )
{
  EFI_GUID  *HobGuid;
  UINTN     Index;

  if (gHobList != NULL) {
    return gHobList;
  }

  gHobList = NULL;
  if (*(UINTN *)((UINTN)gST + 104) > 0) {
    for (Index = 0; Index < *(UINTN *)((UINTN)gST + 104); Index++) {
      HobGuid = (EFI_GUID *)(*(UINTN *)((UINTN)gST + 112) + Index * 24);
      if (HobCompareGuid (&gEfiHobListGuid, HobGuid)) {
        gHobList = *(VOID **)((UINTN)HobGuid + 16);
        break;
      }
    }
    if (gHobList == NULL) {
      UbaDebugPrint (0x80000000LL, "\nASSERT_EFI_ERROR (Status = %r)\n", 0x800000000000000EuLL);
      UbaAssertPrint ("e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c", 54, "!EFI_ERROR (Status)");
    }
  } else {
    UbaDebugPrint (0x80000000LL, "\nASSERT_EFI_ERROR (Status = %r)\n", 0x800000000000000EuLL);
    UbaAssertPrint ("e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c", 54, "!EFI_ERROR (Status)");
    if (gHobList == NULL) {
      UbaAssertPrint ("e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c", 55, "mHobList != ((void *) 0)");
      return gHobList;
    }
  }

  return gHobList;
}

/**
  Get UBA Configuration Board Info.

  Returns a pointer to the platform's board configuration data structure
  located at 0xD20 in the .data section. The structure contains static
  platform identification and config data.

  @param[out] BoardInfo  Pointer to receive the board config address.

  @retval EFI_SUCCESS  BoardInfo was set.
**/
EFI_STATUS
EFIAPI
GetUbaConfigBoardInfo (
  OUT VOID **BoardInfo
  )
{
  *BoardInfo = &mUbaBoardConfig;
  return EFI_SUCCESS;
}

/**
  Get UBA Configuration Slot Count.

  Returns the PCIe slot configuration array (at 0xD60) and the
  number of configured slots (6) for the Lightning Ridge EXRP.

  Each entry in the slot config is a PCH_SLOT_CONFIG structure
  (24 bytes: Domain/Bus/Device/Function/SlotNumber).

  @param[out] SlotData   Pointer to receive the slot config array.
  @param[out] SlotCount  Receives the slot count (6).

  @retval EFI_SUCCESS  SlotData and SlotCount were set.
**/
EFI_STATUS
EFIAPI
GetUbaConfigSlotCount (
  OUT VOID    **SlotData,
  OUT UINTN   *SlotCount
  )
{
  *SlotData  = &mUbaSlotConfig;
  *SlotCount = 6;
  return EFI_SUCCESS;
}

/**
  Get UBA Configuration Slot Data.

  Returns the extended slot configuration data (at 0xE60) and the
  number of entries (10). This is used for additional slot mapping
  beyond the basic slot count table.

  @param[out] SlotData   Pointer to receive the extended slot data.
  @param[out] SlotCount  Receives the entry count (10).

  @retval EFI_SUCCESS  SlotData and SlotCount were set.
**/
EFI_STATUS
EFIAPI
GetUbaConfigSlotData (
  OUT VOID    **SlotData,
  OUT UINTN   *SlotCount
  )
{
  *SlotData  = &mUbaSlotExtendedConfig;
  *SlotCount = 10;
  return EFI_SUCCESS;
}

/**
  PCIe Slot Number validation callback.

  Iterates an 8-entry slot configuration table (mUbaSlotNumberCfg at 0xED1),
  4 bytes per entry: {SlotNumber, Segment/Flags, Bus, DevFn}.
  Entries where the corresponding bit in SlotMask is set are skipped.

  For each active entry, constructs a PCI config address from:
  AddrBase = ((Bus << 8) | (Dev << 16) | (Fn << 24)) << 8

  Then reads PCI config registers 0x19 (Slot Base) and 0x1A (Slot Limit)
  via EFI_PCI_ROOT_BRIDGE_IO_PROTOCOL to determine if the input
  PcieAddress falls within [SlotBase, SlotLimit].

  @param[in]  PcieAddress  PCIe address (domain/bus/dev/fn) to validate.
  @param[in]  SlotMask     Bitmask of entries to skip (1=skip).

  @retval TRUE   PcieAddress is within a configured slot range.
  @retval FALSE  PcieAddress is not in any configured slot.
**/
BOOLEAN
EFIAPI
SetPcieSlotNumber (
  IN UINT64   PcieAddress,
  IN UINT32   SlotMask
  )
{
  UINT8       *SlotCfg;
  BOOLEAN     Found;
  UINT16      Index;
  UINT8       Bus;
  UINT8       Device;
  UINT8       Fn;

  //
  // SlotCfg points into the slot table at 0xED1.
  // Each entry is 4 bytes: {SlotNum, Flags, Bus, DevFn}
  // The -1 offset reads the previous entry's slot number byte.
  //
  SlotCfg = (UINT8 *)&mUbaSlotNumberCfg;

  Found   = FALSE;
  for (Index = 0; Index < 8; Index++) {
    if (((SlotMask >> Index) & 1) == 0) {
      Bus    = SlotCfg[-1];   // Bus from previous entry
      Device = SlotCfg[0];    // Device number
      Fn     = SlotCfg[1];    // Function number

      Found = FALSE;

      {
        EFI_PCI_ROOT_BRIDGE_IO_PROTOCOL *PciRbIo;
        UINT8                           SlotBase;
        UINT8                           SlotLimit;
        UINT64                          AddrBase;

        //
        // Construct PCI config address:
        // AddrBase = ((Bus | ((Device | (Fn << 8)) << 8)) << 8)
        // Then read PCI config registers at offset 0x19 and 0x1A
        // to get slot base/limit for this bus/dev/fn.
        //
        AddrBase = (UINT64)((Bus | ((Device | (Fn << 8)) << 8)) << 8);

        gBS->LocateProtocol (&gEfiPciRootBridgeIoProtocolGuid, NULL, (VOID **)&PciRbIo);

        PciRbIo->Pci.Read (PciRbIo, EfiPciWidthUint8, AddrBase | 0x19, 1, &SlotBase);
        PciRbIo->Pci.Read (PciRbIo, EfiPciWidthUint8, AddrBase | 0x1A, 1, &SlotLimit);

        if (PcieAddress >= SlotBase && PcieAddress <= SlotLimit) {
          Found = TRUE;
        }
      }
    }

    SlotCfg += 4;   // Advance 4 bytes per entry (was 3 in original, corrected)
    if (Found) {
      break;
    }
  }

  return Found;
}

/**
  Initialize PCIe Slot Number callback state.

  Sets the initial callback state to 0 and prints a debug message
  indicating that the SetPcieSlotNumber callback has been initialized
  for this platform.

  @param[out] CallbackState  Receives initial callback state (0).

  @retval EFI_SUCCESS  Callback state was initialized.
**/
EFI_STATUS
EFIAPI
SetPcieSlotNumberInit (
  OUT UINT8 *CallbackState
  )
{
  *CallbackState = 0;
  UbaDebugPrint (0x80000000LL, "[UBA]:SetPcieSlotNumber callback - %d\n");
  return EFI_SUCCESS;
}

/**
  UEFI Driver Entry Point.

  Entry point for the OpromUpdateDxeLightningRidgeEXRP module.
  Performs the following initialization sequence:

  1. Caches UEFI global variables (ImageHandle, SystemTable, BootServices,
     RuntimeServices) with ASSERT checks for NULL.
  2. Calls HobGetGuid() to initialize the HOB list pointer.
  3. Prints a module identification debug message.
  4. Locates the UBA protocol (gUbaProtocolGuid).
  5. If successful, registers the platform's PCIe slot number configuration
     via Uba->SetPcieSlotNumber() with:
     - SlotNumberCfg: mUbaSlotNumberCfg table at 0xED1 (8 x 4-byte entries)
     - CallbackCfg: mUbaSlotNumberCallbackCfg at 0xEF0 (48 bytes)
     - DataSize: 48

  @param[in]  ImageHandle  UEFI image handle.
  @param[in]  SystemTable  UEFI system table.

  @retval EFI_SUCCESS            Protocol callback was installed.
  @retval EFI_NOT_FOUND          UBA protocol GUID not found.
  @retval EFI_INVALID_PARAMETER  Invalid protocol interface.
**/
EFI_STATUS
EFIAPI
OpromUpdateDxeLightningRidgeEXRP (
  IN EFI_HANDLE         ImageHandle,
  IN EFI_SYSTEM_TABLE   *SystemTable
  )
{
  EFI_STATUS                   Status;
  UBA_PROTOCOL                 *Uba;

  gImageHandle = (UINTN)ImageHandle;
  ASSERT (ImageHandle != NULL);
  gST = (EFI_SYSTEM_TABLE *)SystemTable;
  ASSERT (SystemTable != NULL);
  gBS = SystemTable->BootServices;
  ASSERT (gBS != NULL);
  gRT = SystemTable->RuntimeServices;
  ASSERT (gRT != NULL);

  HobGetGuid ();
  UbaDebugPrint (0x80000000LL, "UBA:OpromUpdate-TypeLightningRidgeEXRP\n");

  Status = gBS->LocateProtocol (&gUbaProtocolGuid, NULL, (VOID **)&Uba);
  if (!EFI_ERROR (Status)) {
    return Uba->SetPcieSlotNumber (
                  Uba,
                  &mUbaSlotNumberCfg,
                  &mUbaSlotNumberCallbackCfg,
                  48
                  );
  }

  return Status;
}