/** @file
  IioCfgUpdateDxeNeonCityEPECB - IIO Configuration Update Driver for NeonCity EPECB

  This DXE driver registers platform-specific IIO (Integrated I/O) configuration
  data via the UBA (Universal BIOS Agent) protocol for the NeonCity EP EC B
  platform. It publishes four IIO configuration tables that describe PCIe port
  topology, bifurcation settings, and link configuration per IIO stack.

  The driver flow:
    1. ModuleEntryPoint initializes UEFI service table globals (gImageHandle,
       gST, gBS, gRT) with assertions.
    2. Retrieves the HOB list from the system table's firmware resource array.
    3. Locates the UBA IIO configuration protocol via gBS->LocateProtocol.
    4. Calls the protocol's SetIioConfig function four times, each time
       registering a 48-byte IIO configuration block identified by a GUID.

  The IIO configuration data is embedded as a PIIO-format table in the .data
  section, containing per-stack PCIe port configuration entries.

  Copyright (c) Intel Corporation. All rights reserved.
  SPDX-License-Identifier: BSD-2-Clause-Patent

 **/

#include <Uefi.h>
#include <Library/BaseLib.h>
#include <Library/DebugLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/HobLib.h>
#include <Protocol/Ubiquito.h>

#include "IioCfgUpdateDxeNeonCityEPECB.h"

//
// UEFI Protocol GUID definitions
//
// These GUIDs identify the IIO configuration tables and protocols
// used by this driver for the NeonCity EP EC B platform.
//

EFI_GUID gIioConfigTable0Guid      = IIO_CONFIG_TABLE_0_GUID;
EFI_GUID gIioConfigTable1Guid      = IIO_CONFIG_TABLE_1_GUID;
EFI_GUID gIioConfigTable2Guid      = IIO_CONFIG_TABLE_2_GUID;
EFI_GUID gIioConfigTable3Guid      = IIO_CONFIG_TABLE_3_GUID;
EFI_GUID gIioConfigProtocolGuid    = IIO_CONFIG_PROTOCOL_GUID;
EFI_GUID gUbaPlatformBoardGuid     = UBA_PLATFORM_BOARD_PROTOCOL_GUID;

//
// HOB discovery GUID pair.
//
// The system table contains a firmware resource array. Each entry is a 24-byte
// descriptor. To find the HOB list, the driver scans these descriptors looking
// for a pair where the first QWORD of descriptor[N] matches gHobListGuid1 and
// the first QWORD of descriptor[N] matches gHobListGuid2. The third QWORD of
// the matching descriptor contains the HOB list pointer.
//
// gHobListGuid1 = gEfiDxeServicesTableGuid
//   { 0x7739F24C, 0x93D7, 0x11D4, { 0x9A, 0x3A, 0x00, 0x90, 0x27, 0x3F, 0xC1, 0x4D } }
//
EFI_GUID gHobListGuid1             = { 0x7739F24C, 0x93D7, 0x11D4, { 0x9A, 0x3A, 0x00, 0x90, 0x27, 0x3F, 0xC1, 0x4D } };
EFI_GUID gHobListGuid2             = { 0x90003A9A, 0x3F27, 0x4DC1, { 0x82, 0xD0, 0x3B, 0x12, 0x01, 0x32, 0x5C, 0x46 } };

//
// Module globals
//
VOID    *mHobList            = NULL;  ///< Cached HOB list pointer
VOID    *mUbaDebugProtocol   = NULL;  ///< Cached UBA debug protocol pointer

//
// IIO Configuration Table (platform data for NeonCity EP EC B)
//
// The configuration data is stored in PIIO format (signature "PIIO", version 1,
// length 0xD48). It contains IIO port configuration entries including PCIe port
// numbers, bifurcation settings, link speeds, and topology information for each
// IIO stack on the NeonCity EP EC B platform.
//
#pragma pack(1)
typedef struct {
  UINT32    Signature;          ///< "PIIO" = 0x4F494950 (Platform IIO)
  UINT32    Version;            ///< Table version (1)
  UINT64    Length;             ///< Total config data length (0xD48 = 3400 bytes)
  UINT8     Data[];             ///< IIO configuration entries
} IIO_CFG_TABLE;

IIO_CFG_TABLE mIioCfgTable = {
  .Signature = 0x4F494950,  // "PIIO"
  .Version   = 1,
  .Length    = 0xD48,
  .Data      = {
    // IIO port configuration entries follow
    // (see .data section at 0xC1C for full contents)
  },
};
#pragma pack()

//
// Internal function prototypes
//
UINT64
ReadUnaligned64 (
  IN CONST VOID  *Buffer
  );

BOOLEAN
IsMatchingHobGuidPair (
  IN VOID  *HobGuid1,
  IN VOID  *HobGuid2
  );

/**
  Reads a 64-bit value from an unaligned pointer.

  Equivalent to ReadUnaligned64 from BaseLib on x64 since the CPU supports
  unaligned access natively. Asserts that Buffer is not NULL.

  @param[in] Buffer  Pointer to read from (must not be NULL).

  @return The 64-bit value at the given address.
**/
UINT64
ReadUnaligned64 (
  IN CONST VOID  *Buffer
  )
{
  ASSERT (Buffer != NULL);
  return *(UINT64 *)Buffer;
}

/**
  Checks whether two GUID values read from HOB descriptors match the
  expected GUID pair, by comparing their first 8 bytes (QWORD) only.

  This is an optimization that reduces the comparison to two UINT64
  checks instead of a full 16-byte GUID comparison.

  @param[in] HobGuid1  Pointer to the first GUID to compare.
  @param[in] HobGuid2  Pointer to the second GUID to compare.

  @return TRUE if both GUID comparisons match, FALSE otherwise.
**/
BOOLEAN
IsMatchingHobGuidPair (
  IN VOID  *HobGuid1,
  IN VOID  *HobGuid2
  )
{
  return (ReadUnaligned64 (HobGuid1) == ReadUnaligned64 (&gHobListGuid1) &&
          ReadUnaligned64 (HobGuid2) == ReadUnaligned64 (&gHobListGuid2));
}

/**
  Retrieves the HOB (Hand-Off Block) list from the system table.

  Searches the firmware resource array in the system table for a matching
  GUID pair (gHobListGuid1/gHobListGuid2). Each firmware resource descriptor
  is 24 bytes (3 * UINT64). When a descriptor's first QWORD matches gHobListGuid1
  and the following descriptor's first QWORD matches gHobListGuid2, the third
  QWORD of the first descriptor is the HOB list pointer.

  The result is cached in mHobList for subsequent calls.

  @return A pointer to the HOB list root, or NULL if not found.
**/
VOID *
GetHobList (
  VOID
  )
{
  UINTN   Index;
  UINTN   FwResourceCount;
  VOID    *FwResourceArray;

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

  mHobList = NULL;

  //
  // The system table at offset 0x68 (104) contains the number of firmware
  // resource descriptors, and at offset 0x70 (112) contains the array pointer.
  //
  FwResourceCount   = *(UINTN *)((UINTN)gST + 104);
  FwResourceArray   = *(VOID **)((UINTN)gST + 112);

  if (FwResourceArray != NULL) {
    for (Index = 0; Index < FwResourceCount; Index++) {
      if (IsMatchingHobGuidPair (
            (VOID *)((UINTN)FwResourceArray + (Index * 24)),
            (VOID *)((UINTN)FwResourceArray + (Index * 24) + 8)
            ))
      {
        //
        // The third QWORD of the matching descriptor is the HOB list pointer.
        //
        mHobList = *(VOID **)((UINTN)FwResourceArray + (Index * 24) + 16);
        break;
      }
    }
  }

  if (mHobList == NULL) {
    //
    // HOB list not found: report error via debug and assert.
    //
    DEBUG ((EFI_D_ERROR, "\nASSERT_EFI_ERROR (Status = %r)\n", EFI_NOT_FOUND));
    ASSERT_EFI_ERROR (EFI_NOT_FOUND);
  }

  ASSERT (mHobList != NULL);
  return mHobList;
}

/**
  Locates the UBA debug protocol interface.

  Tests whether sufficient page allocation capacity is available by
  allocating and freeing 0x10 pages (64 KB) of boot services data.
  If the allocation succeeds (returned pointer <= 0x10 meaning the
  AllocatePages call consumed a low address), the UBA debug protocol
  is located via gBS->LocateProtocol.

  The result is cached in mUbaDebugProtocol.

  @return A pointer to the UBA debug protocol interface, or NULL if not found.
**/
VOID *
GetDebugProtocol (
  VOID
  )
{
  UINT64        Pages;
  EFI_STATUS    Status;

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

  //
  // Probe page allocator: try to allocate 0x10 pages (64 KB).
  // The pointer value is checked against 0x10 to determine if the
  // allocator function is available (low addresses are typical for
  // early DXE phase allocations, and a return value <= 0x10 suggests
  // the allocation backend is functional).
  //
  Pages = (UINT64)(UINTN)gBS->AllocatePages (AllocateAnyPages, EfiBootServicesData, 0x10);
  gBS->FreePages ((EFI_PHYSICAL_ADDRESS)Pages, 0x10);

  if (Pages <= 0x10) {
    Status = gBS->LocateProtocol (&gUbaPlatformBoardGuid, NULL, (VOID **)&mUbaDebugProtocol);
    if (EFI_ERROR (Status)) {
      mUbaDebugProtocol = NULL;
    }
  } else {
    mUbaDebugProtocol = NULL;
  }

  return mUbaDebugProtocol;
}

/**
  Debug assertion handler.

  When the UBA debug protocol is available, forwards the assertion
  information to the protocol's assertion handler function at offset 8
  in the protocol vtable.

  @param[in] FileName     The source file name where the assertion occurred.
  @param[in] LineNumber   The line number of the assertion.
  @param[in] Description  The assertion description string.
**/
VOID
EFIAPI
DebugAssert (
  IN CONST CHAR8  *FileName,
  IN UINTN        LineNumber,
  IN CONST CHAR8  *Description
  )
{
  VOID    *DebugProtocol;
  UINT64  (*AssertFunc)(CONST CHAR8 *, UINTN, CONST CHAR8 *);

  DebugProtocol = GetDebugProtocol ();
  if (DebugProtocol != NULL) {
    //
    // The assertion handler is at offset 8 in the protocol vtable.
    //
    AssertFunc = *(UINT64 (*)(CONST CHAR8 *, UINTN, CONST CHAR8 *))((UINTN)DebugProtocol + 8);
    AssertFunc (FileName, LineNumber, Description);
  }
}

/**
  Debug print handler.

  Checks the platform's CMOS debug level and filters the message based on the
  error level. When the UBA debug protocol is available and the message passes
  the filter, it is forwarded to the protocol's print function at offset 0 in
  the vtable.

  CMOS offset 0x4B determines the debug output level:
    0 = Disabled (falls back to hardware strap)
    1 = Errors only (EFI_D_ERROR)
    2 = Warnings
    3 = Info
    >3 = Reserved (treated as level from strap)

  @param[in]  ErrorLevel  The error level of the debug message.
  @param[in]  Format      The format string.
  @param[in]  ...         Variable arguments for the format string.

  @return The number of characters printed, or 0 if no handler was available
          or the message was filtered.
**/
UINTN
EFIAPI
DebugPrint (
  IN UINTN        ErrorLevel,
  IN CONST CHAR8  *Format,
  ...
  )
{
  VOID      *DebugProtocol;
  UINT64    (*PrintFunc)(UINTN, CONST CHAR8 *, ...);
  UINT8     CmosIndex;
  UINT8     CmosData;
  UINT8     DebugLevel;
  UINT64    FilterMask;
  BOOLEAN   ShouldPrint;
  VA_LIST   Args;

  DebugProtocol = GetDebugProtocol ();
  if (DebugProtocol != NULL) {
    //
    // Read the current CMOS index register (I/O port 0x70),
    // then select CMOS offset 0x4B (BIOS debug level).
    //
    CmosIndex = __inbyte (0x70);
    __outbyte (0x70, CmosIndex & 0x80 | 0x4B);
    CmosData = __inbyte (0x71);
    DebugLevel = CmosData;

    //
    // Determine effective debug level.
    // If DebugLevel > 3 (e.g., uninitialized CMOS), fall back to
    // hardware strap: read from I/O port 0xFDAF0490 bit 1, then OR with 1.
    //
    if (DebugLevel > 3) {
      if (DebugLevel == 0) {
        DebugLevel = (*(volatile UINT8 *)0xFDAF0490 & 2) | 1;
      }
    }

    //
    // Compute the filter mask based on debug level.
    // Level 1 (errors only):  0x80000004
    // Level 2+ (errors+warn): 0x80000006
    //
    ShouldPrint = FALSE;
    if ((DebugLevel - 1) <= 0xFD) {
      FilterMask = (DebugLevel == 1) ? 0x80000004ULL : 0x80000006ULL;
      if ((FilterMask & ErrorLevel) != 0) {
        ShouldPrint = TRUE;
      }
    }

    if (ShouldPrint) {
      VA_START (Args, Format);
      //
      // The print function is at offset 0 in the protocol vtable.
      //
      PrintFunc = *(UINT64 (*)(UINTN, CONST CHAR8 *, ...))((UINTN)DebugProtocol);
      return PrintFunc (ErrorLevel, Format, Args);
    }
  }

  return 0;
}

/**
  Registers IIO configuration data via the UBA protocol.

  This function performs the following steps:
    1. Issues a DEBUG message identifying the platform type.
    2. Locates the UBA IIO configuration protocol via gBS->LocateProtocol.
    3. Calls the protocol's SetIioConfig function four times to register
       each of the four IIO configuration tables.

  The four configuration tables are identified by:
    - gIioConfigTable0Guid  (IIO Config Table 0)
    - gIioConfigTable1Guid  (IIO Config Table 1)
    - gIioConfigTable2Guid  (IIO Config Table 2)
    - gIioConfigTable3Guid  (IIO Config Table 3)

  Each table is 48 bytes of platform-specific IIO configuration data sourced
  from the embedded PIIO table (mIioCfgTable).

  @return EFI_SUCCESS      All four configuration tables were registered.
  @return EFI_NOT_FOUND    The UBA protocol could not be located.
  @return Error codes      From the UBA protocol SetIioConfig function.
**/
EFI_STATUS
SetIioConfiguration (
  VOID
  )
{
  EFI_STATUS  Status;
  VOID        *UbaProtocol;

  //
  // Report the platform type for debug purposes.
  //
  DEBUG ((EFI_D_INFO, "UBA:IioCfgUpdate-TypeNeonCityEPECB\n"));

  //
  // Locate the UBA IIO configuration protocol.
  // The protocol GUID is IIO_CONFIG_PROTOCOL_GUID.
  //
  Status = gBS->LocateProtocol (&gIioConfigProtocolGuid, NULL, &UbaProtocol);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Call the protocol's SetIioConfig function at offset 16 in the vtable
  // (second function pointer on x64: sizeof(VOID*) = 8, so offset 16 = 2nd entry).
  // Signature: EFI_STATUS (*)(VOID *This, EFI_GUID *CfgGuid, VOID *CfgData, UINTN DataSize)
  //
  // Register configuration table 0.
  //
  Status = ((EFI_STATUS (*)(VOID *, EFI_GUID *, VOID *, UINTN))((UINTN)UbaProtocol + 16))(
             UbaProtocol,
             &gIioConfigTable0Guid,
             &mIioCfgTable,
             48
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register configuration table 1.
  //
  Status = ((EFI_STATUS (*)(VOID *, EFI_GUID *, VOID *, UINTN))((UINTN)UbaProtocol + 16))(
             UbaProtocol,
             &gIioConfigTable1Guid,
             &mIioCfgTable,
             48
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register configuration table 2.
  //
  Status = ((EFI_STATUS (*)(VOID *, EFI_GUID *, VOID *, UINTN))((UINTN)UbaProtocol + 16))(
             UbaProtocol,
             &gIioConfigTable2Guid,
             &mIioCfgTable,
             48
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register configuration table 3.
  //
  Status = ((EFI_STATUS (*)(VOID *, EFI_GUID *, VOID *, UINTN))((UINTN)UbaProtocol + 16))(
             UbaProtocol,
             &gIioConfigTable3Guid,
             &mIioCfgTable,
             48
             );
  return Status;
}

/**
  Stub function that always returns zero.
  Used as a placeholder for unused function pointer slots or no-op callbacks.

  @return 0
**/
UINTN
NoOpStub (
  VOID
  )
{
  return 0;
}

/**
  The module entry point.

  Initializes the UEFI global variables (gImageHandle, gST, gBS, gRT) with
  assertions for NULL safety, retrieves the HOB list from the system table's
  firmware resource array, and invokes the IIO configuration registration
  routine.

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

  @return EFI_SUCCESS     IIO configuration was registered successfully.
  @return EFI_NOT_FOUND   The UBA IIO config protocol was not available.
  @return Other error     From SetIioConfiguration() or LocateProtocol().
**/
EFI_STATUS
EFIAPI
ModuleEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS  Status;

  //
  // Save the global image handle.
  //
  ASSERT (ImageHandle != NULL);
  gImageHandle = ImageHandle;

  //
  // Save the UEFI system table pointer.
  //
  ASSERT (SystemTable != NULL);
  gST = SystemTable;

  //
  // Save the boot services table pointer.
  //
  ASSERT (SystemTable->BootServices != NULL);
  gBS = SystemTable->BootServices;

  //
  // Save the runtime services table pointer.
  //
  ASSERT (SystemTable->RuntimeServices != NULL);
  gRT = SystemTable->RuntimeServices;

  //
  // Retrieve the HOB list (needed by platform configuration code).
  //
  GetHobList ();

  //
  // Register the IIO configuration data via the UBA protocol.
  //
  Status = SetIioConfiguration ();
  return Status;
}