/** @file
  SlotDataUpdateDxeNeonCityFPGA - Platform Slot Table (PSLT) update driver
  for NeonCity FPGA platforms.

  This driver is part of the Lenovo UBA (Universal BIOS Architecture) framework.
  It provides platform-specific PCIe slot configuration data for the NeonCity
  FPGA platform by registering PSLT entries through the UBA protocol.

  Copyright (C) 2026, Lenovo. All rights reserved.
  SPDX-License-Identifier: BSD-2-Clause-Patent
**/

#include "SlotDataUpdateDxeNeonCityFPGA.h"

//
// =============================================================================
// GUID Definitions
// =============================================================================
//

EFI_GUID gUbaProtocolGuid          = UBA_PROTOCOL_GUID;
EFI_GUID gUbaSlotDataPsl1Guid      = UBA_SLOT_DATA_PSLT1_GUID;
EFI_GUID gUbaSlotDataPsl2Guid      = UBA_SLOT_DATA_PSLT2_GUID;
EFI_GUID gUbaSlotDataPsl3Guid      = UBA_SLOT_DATA_PSLT3_GUID;
EFI_GUID gEfiDxeServicesTableGuid  = DXE_SERVICES_TABLE_GUID;

//
// =============================================================================
// Global Variables
// =============================================================================
//

EFI_HANDLE            gImageHandle      = NULL;  ///< 0xC80: Image handle from entry point
EFI_SYSTEM_TABLE     *gSystemTable      = NULL;  ///< 0xC70: System table from entry point
EFI_BOOT_SERVICES    *gBootServices     = NULL;  ///< 0xC78: Boot services table pointer
EFI_RUNTIME_SERVICES *gRuntimeServices  = NULL;  ///< 0xC88: Runtime services table pointer
VOID                 *gUbaProtocol      = NULL;  ///< 0xC90: Cached UBA protocol interface
VOID                 *gHobList          = NULL;  ///< 0xC98: Cached HOB list pointer
UINT8                 gSlotPlatformId   = 0;     ///< 0xCA0: Cached platform slot type

//
// =============================================================================
// Platform Slot Data (embedded in .data section)
//
// Three PSLT entries are defined statically and registered with the UBA
// protocol during initialization. The PSLT format:
//
//   Offset  Size  Field
//   ------  ----  -----
//   0       4     Signature   ("PSLT")
//   4       4     Version     (1)
//   8       4     Length      (total size of this entry)
//   12      4     SlotCount   (function pointer offset to GetSlotCount)
//   16      4     Reserved1
//   20      4     SlotDataOffset (function pointer offset to GetSlotData)
//   24      16    ConfigData  (OEM-specific bitfield configuration)
//
// =============================================================================
//

// PSLT Entry 1 (at 0xC50, size 32): Primary PSLT with full length 0xC40
STATIC CONST UINT8 mPslEntry1[32] = {
  'P', 'S', 'L', 'T',             // Signature
  0x01, 0x00, 0x00, 0x00,         // Version = 1
  0x40, 0x0C, 0x00, 0x00,         // Length = 0xC40 (3136 bytes)
  0xE8, 0x04, 0x00, 0x00,         // SlotCount = 0x4E8 (GetSlotCount function)
  0x00, 0x00, 0x00, 0x00,         // Reserved1 = 0
  0xEC, 0x04, 0x00, 0x00,         // SlotDataOffset = 0x4EC (GetSlotData function)
  0x01, 0x01, 0x02, 0x01,         // ConfigData[0..3]
  0x02, 0x01, 0xFF, 0x05,         // ConfigData[4..7]
  0x04, 0x00, 0x00, 0x00          // ConfigData[8..11] / Pad
};

// PSLT Entry 2 (at 0xBE0, size 40): Secondary PSLT with length 0xC40
STATIC CONST UINT8 mPslEntry2[40] = {
  'P', 'S', 'L', 'T',             // Signature
  0x01, 0x00, 0x00, 0x00,         // Version = 1
  0x40, 0x0C, 0x00, 0x00,         // Length = 0xC40
  0xE8, 0x04, 0x00, 0x00,         // SlotCount = 0x4E8
  0x00, 0x00, 0x00, 0x00,         // Reserved1 = 0
  0xEC, 0x04, 0x00, 0x00,         // SlotDataOffset = 0x4EC
  0x00, 0x00, 0x00, 0x00,         // ConfigData[0..3] = 0
  0x00, 0x00, 0x00, 0x00,         // ConfigData[4..7] = 0
  0x00, 0x00, 0x00, 0x00,         // ConfigData[8..11] = 0
  0x00, 0x00, 0x00, 0x00,         // ConfigData[12..15] = fill to size 40
};

// PSLT Entry 3 (at 0xC18, size 40): Tertiary PSLT with alternate length 0xC08
STATIC CONST UINT8 mPslEntry3[40] = {
  'P', 'S', 'L', 'T',             // Signature
  0x01, 0x00, 0x00, 0x00,         // Version = 1
  0x08, 0x0C, 0x00, 0x00,         // Length = 0xC08 (3080 bytes)
  0xE8, 0x04, 0x00, 0x00,         // SlotCount = 0x4E8
  0x00, 0x00, 0x00, 0x00,         // Reserved1 = 0
  0xEC, 0x04, 0x00, 0x00,         // SlotDataOffset = 0x4EC
  0x01, 0x01, 0x02, 0x01,         // ConfigData[0..3]
  0x00, 0x01, 0xFF, 0x05,         // ConfigData[4..7]
  0x04, 0x00, 0x00, 0x00,         // ConfigData[8..11]
  0x00, 0x00, 0x00, 0x00,         // ConfigData[12..15]
};

//
// =============================================================================
// Function Implementations
// =============================================================================
//

/**
  Returns the number of PCIe slots on this platform.

  For NeonCity FPGA, the platform always has 3 slots.

  @return UINT8  3 (fixed slot count for this platform).
**/
UINT8
GetSlotCount (
  VOID
  )
{
  return 3;
}

/**
  Returns slot configuration data based on slot number and platform type.

  The platform slot type is determined from RTC CMOS register 0x4B (read
  via I/O ports 0x70/0x71). If N3 (the platform type identifier) is 0 or 3,
  the function returns 15, which indicates the slot uses "other" or standard
  PCIe slot classification. Otherwise, the provided DefaultData is returned
  unchanged.

  @param[in] SlotNumber   The slot number (0-indexed) to query.
  @param[in] DefaultData  Default slot data value to return if no special
                          mapping is required for this platform type.

  @return UINT8  Slot configuration data byte.
**/
UINT8
GetSlotData (
  IN UINT8 SlotNumber,
  IN UINT8 DefaultData
  )
{
  //
  // Check if the platform type (N3) is 0 or 3, which maps to
  // a "standard/other" slot classification (value 15 = 0x0F).
  //
  if ((SlotNumber == 0) || (SlotNumber == 3)) {
    return 15;
  }

  return DefaultData;
}

/**
  Locates and caches the UBA protocol interface.

  This function allocates a small test buffer (EfiBootServicesData pool,
  17 bytes) and frees it to verify that enough pool memory is available
  for protocol operations. If the allocation succeeds (requested size > 0x10),
  it proceeds to locate the UBA protocol via gBS->LocateProtocol().

  The located protocol interface is cached in the global gUbaProtocol.

  @return Pointer to the UBA protocol interface, or NULL if:
          - Memory pool allocation failed (not enough available pool)
          - LocateProtocol failed to find the UBA protocol
**/
VOID *
UbaGetProtocolInterface (
  VOID
  )
{
  //
  // Return cached protocol if already located.
  //
  if (gUbaProtocol != NULL) {
    return gUbaProtocol;
  }

  //
  // Allocate a small test buffer to verify pool availability.
  // Size 31 with EfiBootServicesData type via gBS->AllocatePool.
  // (BootServices + 24 is AllocatePool, + 32 is FreePool)
  // The allocation of 31 bytes is a probe to verify the UEFI memory
  // system is operational. If the allocation returns > 0x10 bytes,
  // we proceed. If the allocation fails or returns <= 0x10 bytes,
  // we bail out (return NULL).
  //
  {
    UINTN   TestBufferSize;
    VOID   *TestBuffer;

    TestBufferSize = (UINTN)gBootServices->AllocatePool (EfiBootServicesData, 31, &TestBuffer);
    gBootServices->FreePool (TestBuffer);

    if (TestBufferSize <= 0x10) {
      return NULL;
    }
  }

  //
  // Locate the UBA protocol using its GUID.
  //
  {
    EFI_STATUS  Status;

    Status = gBootServices->LocateProtocol (
                              &gUbaProtocolGuid,   // GUID: 76232936-C80E-A131-3A3A-F2FC1C3932
                              NULL,                 // No registration key
                              &gUbaProtocol         // Output: located protocol interface
                              );

    if (EFI_ERROR (Status)) {
      gUbaProtocol = NULL;
    }
  }

  return gUbaProtocol;
}

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

  Reads the platform type from RTC CMOS register 0x4B (I/O ports 0x70/0x71).
  The platform type (N3 byte from CMOS) determines the debug message routing:

    - If N3 > 3 and N3 != 0: normal platform routing
    - If N3 == 0: reads hardware strap from F000:0490 bit 1 | 1
    - If N3 == 1: debug output to UBA_DEBUG_ROUTE_A (0x80000004)
    - Otherwise: debug output to UBA_DEBUG_ROUTE_B (0x80000006)

  The actual debug output is performed through the UBA protocol's function
  index 1 (offset 8 from the protocol base), which takes (DebugLevel, Format, Args).

  @param[in] DebugLevel  Debug level mask. Only messages with matching bits
                          in the routing mask are output (typically 0x80000000).
  @param[in] Format      Format string for the debug message.
  @param[in] ...         Variable argument list for format substitution.
**/
VOID
EFIAPI
UbaDebugPrint (
  IN UINTN        DebugLevel,
  IN CONST CHAR8 *Format,
  ...
  )
{
  UINTN                      DebugRoute;
  UINT8                      PlatformType;
  UINT8                      N3;
  UBA_PROTOCOL              *Uba;
  UBA_DEBUG_PRINT            UbaDebugFunc;
  VA_LIST                    Va;

  //
  // Get the UBA protocol interface (lazy init through cached result).
  //
  Uba = (UBA_PROTOCOL *)UbaGetProtocolInterface ();
  DebugRoute = 0;

  if (Uba != NULL) {
    //
    // Read platform type from RTC CMOS register 0x4B.
    // I/O port 0x70 is the CMOS address port, 0x71 is the data port.
    // We preserve bit 7 (NMI mask) of the address register.
    //
    PlatformType = IoRead8 (0x70);
    IoWrite8 (0x70, (PlatformType & 0x80) | 0x4B);  // Select CMOS reg 0x4B
    N3 = IoRead8 (0x71);                              // Read platform type byte

    //
    // Determine debug routing based on platform type.
    //
    if (N3 > 3) {
      // N3 is in valid range > 3, use as-is
    } else if (N3 == 0) {
      // Platform type 0: read hardware strap from F000:0490 bit 1, OR with 1
      N3 = (*(volatile UINT8 *)0xFDAF0490 & 0x02) | 0x01;
    }
    // else N3 is 1, 2, or 3: use directly

    //
    // Map platform type to debug route value.
    //
    if ((N3 - 1) <= 0xFD) {
      // N3 is 1..3 (valid range check passed)
      if (N3 == 1) {
        DebugRoute = 0x80000004;    // UBA_DEBUG_ROUTE_A
      } else {
        DebugRoute = 0x80000006;    // UBA_DEBUG_ROUTE_B
      }
    }

    //
    // If the debug level matches the routing mask, call UBA debug output.
    //
    if ((DebugRoute & DebugLevel) != 0) {
      //
      // UBA protocol function index 1 (offset 0x08): DebugPrint
      // Signature: VOID (*)(UINTN DebugLevel, CONST CHAR8 *Format, VA_LIST Va)
      //
      UbaDebugFunc = (UBA_DEBUG_PRINT)(*(UINT64 *)((UINT8 *)Uba + 8));

      VA_START (Va, Format);
      UbaDebugFunc (DebugLevel, Format, (UINT64 *)Va);
      VA_END (Va);
    }
  }
}

/**
  UBA-aware ASSERT replacement.

  Prints the assertion failure information through UBA protocol's debug
  output function and then triggers a system halt.

  @param[in] FileName     Source file name where the assertion occurred.
  @param[in] LineNumber   Line number where the assertion occurred.
  @param[in] AssertString The assertion expression that failed.
**/
VOID
EFIAPI
UbaDebugAssert (
  IN CONST CHAR8 *FileName,
  IN UINTN        LineNumber,
  IN CONST CHAR8 *AssertString
  )
{
  UBA_PROTOCOL    *Uba;
  UBA_DEBUG_ASSERT UbaAssertFunc;

  //
  // Get the UBA protocol and call its assertion handler.
  // UBA protocol function index 1 (offset 0x08) is the assert handler,
  // which takes (FileName, LineNumber, AssertString).
  //
  Uba = (UBA_PROTOCOL *)UbaGetProtocolInterface ();
  if (Uba != NULL) {
    UbaAssertFunc = (UBA_DEBUG_ASSERT)(*(UINT64 *)((UINT8 *)Uba + 8));
    UbaAssertFunc ((UINT64)FileName, LineNumber, (UINT64)AssertString);
  }
}

/**
  Reads a UINT64 value from a potentially unaligned memory address.

  Wraps the BaseLib ReadUnaligned64 function with a NULL pointer check.
  If Buffer is NULL, an assertion is triggered.

  @param[in] Buffer  Pointer to the memory to read. Must not be NULL.

  @return The UINT64 value at the given address.
**/
UINT64
ReadUnaligned64 (
  IN CONST VOID *Buffer
  )
{
  if (Buffer == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\BaseLib\\Unaligned.c",
      192,
      "Buffer != ((void *) 0)"
      );
  }

  return *(const UINT64 *)Buffer;
}

/**
  Checks whether two GUIDs are equal by comparing their 64-bit halves.

  This performs a double 64-bit compare rather than a byte-by-byte or
  128-bit compare, reading the GUIDs as two unaligned 64-bit values.
  This is sufficient for GUID equality checks where both pointers are
  valid and properly aligned for 64-bit access.

  @param[in] Guid1  Pointer to the first GUID to compare.
  @param[in] Guid2  Pointer to the second GUID to compare.

  @return TRUE if the GUIDs are equal, FALSE otherwise.
**/
BOOLEAN
IsHobGuidMatch (
  IN CONST EFI_GUID *Guid1,
  IN CONST EFI_GUID *Guid2
  )
{
  UINT64  Guid1Low;
  UINT64  Guid1High;
  UINT64  Guid2Low;
  UINT64  Guid2High;

  //
  // Read each GUID as two unaligned 64-bit values and compare.
  //
  Guid1Low  = ReadUnaligned64 (Guid1);
  Guid1High = ReadUnaligned64 ((CONST UINT8 *)Guid1 + 8);
  Guid2Low  = ReadUnaligned64 (Guid2);
  Guid2High = ReadUnaligned64 ((CONST UINT8 *)Guid2 + 8);

  return (BOOLEAN)(Guid1Low == Guid2Low && Guid1High == Guid2High);
}

/**
  Retrieves the HOB (Hand-Off Block) list pointer by locating the
  DXE Services HOB.

  Walks the system table's HOB list entries (SystemTable + 104 is the
  HOB list size, SystemTable + 112 is the HOB list pointer). Each HOB
  entry is 24 bytes: 8 bytes for GUID, 8 bytes for length/flags, and
  8 bytes for data pointer. The function searches for the HOB whose
  GUID matches gEfiDxeServicesTableGuid.

  The located HOB list pointer is cached globally in gHobList.

  @return Pointer to the start of the HOB list, or NULL if:
          - The system table has no HOB entries (count = 0)
          - No HOB matches the DXE Services GUID
**/
VOID *
GetHobList (
  VOID
  )
{
  UINTN    HobCount;
  UINT8   *HobEntries;
  UINTN    Index;

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

  //
  // SystemTable + 104 (0x68) = NumberOfTableEntries (UINTN)
  // SystemTable + 112 (0x70) = ConfigurationTable (EFI_CONFIGURATION_TABLE *)
  //
  HobCount    = *(UINTN *)((UINT8 *)gSystemTable + 104);
  HobEntries  = *(UINT8 **)((UINT8 *)gSystemTable + 112);

  gHobList = NULL;

  if (HobCount > 0) {
    //
    // Walk the configuration table to find the DXE Services HOB entry.
    // Each entry is sizeof(EFI_CONFIGURATION_TABLE) = 24 bytes:
    //   0: EFI_GUID VendorGuid  (16 bytes)
    //  16: VOID    *VendorTable (8 bytes)
    //
    for (Index = 0; Index < HobCount; Index++) {
      if (IsHobGuidMatch (
            (EFI_GUID *)(HobEntries + Index * 24),
            &gEfiDxeServicesTableGuid     // {0x7739F24C, 0x93D7, 0x11D4, ...}
            )) {
        //
        // Found the DXE Services HOB - cache its data pointer.
        //
        gHobList = *(VOID **)(HobEntries + Index * 24 + 16);
        break;
      }
    }

    if (gHobList == NULL) {
      //
      // DXE Services HOB not found - report error.
      //
      UbaDebugPrint (0x80000000LL, "\nASSERT_EFI_ERROR (Status = %r)\n", 0x800000000000000EULL);
      UbaDebugAssert (
        "e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c",
        54,
        "!EFI_ERROR (Status)"
        );
    }
  }

  if (gHobList == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c",
      55,
      "mHobList != ((void *) 0)"
      );
  }

  return gHobList;
}

/**
  Installs PSLT Slot Data entries via the UBA protocol.

  This function performs the following actions:
    1. Prints a debug banner via UbaDebugPrint.
    2. Locates the UBA protocol via UbaGetProtocolInterface().
    3. Calls the UBA protocol's SetData function (index 2, offset 0x10)
       three times with different PSLT configuration entries.

  The PSLT entries registered:
    - PSLT1 (GUID B93613E1-...): Primary PSLT, 32 bytes, length 0xC40
    - PSLT2 (GUID 8185B70E-...): Secondary PSLT, 40 bytes, length 0xC40
    - PSLT3 (GUID A87C540B-...): Tertiary PSLT, 40 bytes, length 0xC08

  @return EFI_STATUS  Result of the installation.
                       Returns the first error encountered, or EFI_SUCCESS
                       if all three PSLT entries are registered successfully.
**/
EFI_STATUS
InstallSlotData (
  VOID
  )
{
  EFI_STATUS      Status;
  UBA_PROTOCOL   *Uba;
  UBA_SET_DATA    UbaSetData;

  //
  // Print debug banner identifying this module.
  //
  UbaDebugPrint (0x80000000, "UBA:SlotDataUpdate-TypeNeonCityFPGA\n");

  //
  // Locate the UBA protocol interface.
  //
  Status = gBootServices->LocateProtocol (
                            &gUbaProtocolGuid,   // UBA protocol GUID
                            NULL,
                            (VOID **)&Uba
                            );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // UBA protocol function index 2 (offset 0x10): SetData
  // Signature: EFI_STATUS (*)(UBA_PROTOCOL *This, EFI_GUID *DataGuid,
  //                           VOID *Data, UINTN DataSize)
  //
  UbaSetData = (UBA_SET_DATA)(*(UINT64 *)((UINT8 *)Uba + 16));

  //
  // Register PSLT Entry 1: Primary slot table with length 0xC40
  //
  Status = UbaSetData (Uba, &gUbaSlotDataPsl1Guid, (VOID *)mPslEntry1, 32);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register PSLT Entry 2: Secondary slot table with length 0xC40
  //
  Status = UbaSetData (Uba, &gUbaSlotDataPsl2Guid, (VOID *)mPslEntry2, 40);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register PSLT Entry 3: Tertiary slot table with alternate length 0xC08
  //
  Status = UbaSetData (Uba, &gUbaSlotDataPsl3Guid, (VOID *)mPslEntry3, 40);

  return Status;
}

/**
  Entry point of the SlotDataUpdateDxeNeonCityFPGA driver.

  This is the standard UEFI DXE driver entry point. It:
    1. Validates and caches ImageHandle, SystemTable, BootServices,
       and RuntimeServices pointers.
    2. Initializes the HOB list via GetHobList().
    3. Installs platform slot data via InstallSlotData().

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

  @retval EFI_SUCCESS           Slot data successfully installed.
  @retval EFI_NOT_FOUND         UBA protocol or required HOBs not found.
  @retval EFI_INVALID_PARAMETER One or more parameters are invalid.
  @retval Others                Error from InstallSlotData or protocol operations.
**/
EFI_STATUS
EFIAPI
SlotDataUpdateDxeEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS  Status;

  //
  // Cache UEFI global variables.
  //
  gImageHandle = ImageHandle;
  if (gImageHandle == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiBootServicesTableLib\\UefiBootServicesTableLib.c",
      51,
      "gImageHandle != ((void *) 0)"
      );
  }

  gSystemTable = SystemTable;
  if (gSystemTable == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiBootServicesTableLib\\UefiBootServicesTableLib.c",
      57,
      "gST != ((void *) 0)"
      );
  }

  gBootServices = SystemTable->BootServices;
  if (gBootServices == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiBootServicesTableLib\\UefiBootServicesTableLib.c",
      63,
      "gBS != ((void *) 0)"
      );
  }

  gRuntimeServices = SystemTable->RuntimeServices;
  if (gRuntimeServices == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiRuntimeServicesTableLib\\UefiRuntimeServicesTableLib.c",
      47,
      "gRT != ((void *) 0)"
      );
  }

  //
  // Retrieve the HOB list (cached globally in gHobList).
  //
  GetHobList ();

  //
  // Install platform-specific slot data via UBA protocol.
  //
  Status = InstallSlotData ();

  return Status;
}