/** @file
  MultiSkuDistinctionDxe.c -- Lenovo HR650X Multi-SKU Distinction DXE driver.

  This DXE driver determines the system type (SKU) from a custom HOB
  created during PEI.  Depending on the SKU value it either:
    - Installs a SKU protocol on a new handle, and records the SKU in the
      configuration table (when the SKU type byte is not ASCII 'U', 0x55);
    - Registers an ExitBootServices callback (when the SKU type byte IS 'U'),
      which later installs the SKU protocol just before boot handoff.

  It also contains a small debug-print / assert abstraction that wraps
  a "MsDebug" protocol discovered via LocateProtocol.

  File paths in ASSERT strings reference the EDK-II build tree:
    e:\hs\MdePkg\Library\UefiBootServicesTableLib\...
    e:\hs\MdePkg\Library\UefiRuntimeServicesTableLib\...
    e:\hs\MdePkg\Library\DxeHobLib\...
    e:\hs\MdePkg\Library\BaseLib\Unaligned.c
**/

#include "MultiSkuDistinctionDxe.h"

//
// ---------------------------------------------------------------------------
// Module globals
// ---------------------------------------------------------------------------
EFI_HANDLE              gImageHandle      = NULL;
EFI_SYSTEM_TABLE       *gSystemTable      = NULL;
EFI_BOOT_SERVICES      *gBootServices     = NULL;
EFI_RUNTIME_SERVICES   *gRuntimeServices   = NULL;

//
// Cached MsDebugProtocol pointer (0 until first successful Lookup).
//
MS_DEBUG_PROTOCOL       *gMsDebugProtocol = NULL;

//
// Cached HOB list pointer (0 until first successful GetHobList).
//
VOID                    *gHobList         = NULL;

//
// ---------------------------------------------------------------------------
// GUID definitions
// ---------------------------------------------------------------------------
EFI_GUID gMsDebugProtocolGuid         = MS_DEBUG_PROTOCOL_GUID;
EFI_GUID gLenovoMultiSkuHobGuid       = LENOVO_MULTI_SKU_HOB_GUID;
EFI_GUID gLenovoMultiSkuConfigTableGuid = LENOVO_MULTI_SKU_CONFIG_TABLE_GUID;
EFI_GUID gLenovoSkuProtocolGuid       = LENOVO_SKU_PROTOCOL_GUID;

// ---------------------------------------------------------------------------
// Helper: unaligned read
// ---------------------------------------------------------------------------

/**
  Reads an unaligned QWORD from the given buffer.

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

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

// ---------------------------------------------------------------------------
// Helper: GUID comparison
// ---------------------------------------------------------------------------

/**
  Compares two EFI_GUIDs as two 64-bit integers.

  @param[in]  Guid1  First GUID.
  @param[in]  Guid2  Second GUID.

  @retval TRUE   The GUIDs are equal.
  @retval FALSE  The GUIDs differ.
**/
BOOLEAN
CompareGuid (
  IN EFI_GUID  *Guid1,
  IN EFI_GUID  *Guid2
  )
{
  UINT64  Part1;
  UINT64  Part2;

  Part1 = ReadUnaligned64 (Guid1);
  Part2 = ReadUnaligned64 (Guid2);
  if (Part1 != Part2) {
    return FALSE;
  }
  Part1 = ReadUnaligned64 ((UINT8 *)Guid1 + 8);
  Part2 = ReadUnaligned64 ((UINT8 *)Guid2 + 8);
  return (Part1 == Part2);
}

// ---------------------------------------------------------------------------
// Locate and cache MsDebug protocol
// ---------------------------------------------------------------------------

/**
  Locates and caches the MsDebug protocol.

  Raises TPL to TPL_HIGH_LEVEL and, if the current TPL is <= TPL_NOTIFY,
  calls LocateProtocol to find the protocol by GUID.

  @return  Pointer to the MsDebug protocol interface, or NULL if not found
           or if the TPL was too high to safely locate.
**/
MS_DEBUG_PROTOCOL *
GetMsDebugProtocol (
  VOID
  )
{
  MS_DEBUG_PROTOCOL  *Protocol;
  EFI_STATUS          Status;
  UINTN               OldTpl;

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

  OldTpl = gBootServices->RaiseTPL (TPL_HIGH_LEVEL);
  gBootServices->RestoreTPL (OldTpl);

  if (OldTpl <= TPL_NOTIFY) {
    Status = gBootServices->LocateProtocol (
                              &gMsDebugProtocolGuid,
                              NULL,
                              (VOID **)&Protocol
                              );
    if (!EFI_ERROR (Status)) {
      gMsDebugProtocol = Protocol;
    } else {
      gMsDebugProtocol = NULL;
    }
    return gMsDebugProtocol;
  }

  return NULL;
}

// ---------------------------------------------------------------------------
// Debug print wrapper for MsDebug protocol
// ---------------------------------------------------------------------------

/**
  Debug print via MsDebug protocol.

  Reads CMOS index 0x4B to determine the system/chassis SKU type.
  The SKU byte controls which error-level masks are enabled for printing.

  @param[in]  ErrorLevel  Debug message error level mask.
  @param[in]  Format      Format string.
  @param[in]  ...         Variable arguments.

  @return  Number of characters printed, or FALSE (0) if not printed.
**/
UINTN
EFIAPI
MsDebugPrint (
  IN UINTN   ErrorLevel,
  IN CHAR8   *Format,
  ...
  )
{
  VA_LIST             Args;
  MS_DEBUG_PROTOCOL  *Protocol;
  UINT8               CmosIdx;
  UINT8               SkuByte;
  UINTN               Mask;

  Protocol = GetMsDebugProtocol ();
  if (Protocol == NULL) {
    return 0;
  }

  //
  // Read CMOS RTC index 0x4B -- system identification byte.
  //
  CmosIdx = IoRead8 (0x70);
  IoWrite8 (0x70, (CmosIdx & 0x80) | 0x4B);
  SkuByte = IoRead8 (0x71);

  //
  // Determine mask based on SKU byte.
  //
  if (SkuByte > 3) {
    if (SkuByte == 0) {
      //
      // Fallback: read a flag from fixed MMIO address 0xFDAF0490.
      //
      SkuByte = (*(volatile UINT8 *)0xFDAF0490 & 2) | 1;
    }
  }

  if ((UINT8)(SkuByte - 1) <= 0xFD) {
    if (SkuByte == 1) {
      Mask = 0x80000004;
    } else {
      Mask = 0x80000006;
    }
  } else {
    Mask = 4;
  }

  if ((Mask & ErrorLevel) != 0) {
    VA_START (Args, Format);
    //
    // Protocol vtable entry [0] is the DebugPrint function.
    //
    return Protocol->DebugPrint (ErrorLevel, Format, Args);
  }

  return 0;
}

// ---------------------------------------------------------------------------
// Assert wrapper for MsDebug protocol
// ---------------------------------------------------------------------------

/**
  Assertion handler via MsDebug protocol.

  @param[in]  FileName     Source file name.
  @param[in]  LineNumber   Line number.
  @param[in]  Description  Assertion description.
**/
VOID
EFIAPI
MsDebugAssert (
  IN CHAR8  *FileName,
  IN UINTN  LineNumber,
  IN CHAR8  *Description
  )
{
  MS_DEBUG_PROTOCOL  *Protocol;

  Protocol = GetMsDebugProtocol ();
  if (Protocol != NULL) {
    //
    // Protocol vtable entry [1] is the DebugAssert function.
    //
    Protocol->DebugAssert (FileName, LineNumber, Description);
  }
}

// ---------------------------------------------------------------------------
// HOB list retrieval from configuration table
// ---------------------------------------------------------------------------

/**
  Retrieves and caches the HOB list pointer from the System Table.

  Searches the configuration table for the HOB list GUID.

  @return  Pointer to the HOB list, or NULL if not found.
**/
VOID *
GetHobList (
  VOID
  )
{
  UINTN           Index;
  UINT64          TableCount;
  EFI_GUID       *GuidPtr;
  VOID           *VendorTable;

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

  gHobList = NULL;

  TableCount = gSystemTable->NumberOfTableEntries;
  if (TableCount == 0) {
    //
    // No configuration table entries.
    //
    MsDebugPrint (0x80000000, "\nASSERT_EFI_ERROR (Status = %r)\n", 0x800000000000000Eull);
    MsDebugAssert (
      "e:\\hs\\MdeModulePkg\\Library\\DxeHobLib\\HobLib.c",
      54,
      "!EFI_ERROR (Status)"
      );
    return NULL;
  }

  for (Index = 0; Index < TableCount; Index++) {
    GuidPtr      = gSystemTable->ConfigurationTable[Index].VendorGuid;
    VendorTable  = gSystemTable->ConfigurationTable[Index].VendorTable;

    if (CompareGuid (&gEfiHobListGuid, GuidPtr)) {
      gHobList = VendorTable;
      break;
    }
  }

  if (gHobList == NULL) {
    MsDebugPrint (0x80000000, "\nASSERT_EFI_ERROR (Status = %r)\n", 0x800000000000000Eull);
    MsDebugAssert (
      "e:\\hs\\MdeModulePkg\\Library\\DxeHobLib\\HobLib.c",
      55,
      "mHobList != ((void *) 0)"
      );
  }

  return gHobList;
}

// ---------------------------------------------------------------------------
// HOB GUID-extension search
// ---------------------------------------------------------------------------

/**
  Finds a GUID-extension HOB (type 0x04) whose name matches the given GUID.

  @param[in]  Guid      Pointer to the GUID to search for.
  @param[in]  HobStart  Pointer to the start of the HOB list.

  @return  Pointer to the matching HOB, or NULL if not found.
**/
EFI_HOB_GENERIC_HEADER *
FindHobByGuid (
  IN CONST EFI_GUID  *Guid,
  IN VOID            *HobStart
  )
{
  EFI_HOB_GENERIC_HEADER  *Hob;

  Hob = (EFI_HOB_GENERIC_HEADER *)HobStart;
  if (Hob == NULL) {
    return NULL;
  }

  while (1) {
    if (Hob->HobType == EFI_HOB_TYPE_END_OF_HOB_LIST) {
      return NULL;
    }
    if (Hob->HobType == EFI_HOB_TYPE_GUID_EXT) {
      return Hob;
    }
    Hob = (EFI_HOB_GENERIC_HEADER *)((UINT8 *)Hob + Hob->HobLength);
  }
}

// ---------------------------------------------------------------------------
// Extract Multi-SKU info from HOB
// ---------------------------------------------------------------------------

/**
  Walks the HOB list looking for the Lenovo Multi-SKU HOB
  ({8A0D6B86-64F3-4DC2-A483-498771DFB7FD}) and extracts the SKU type
  bytes from the HOB data.

  @param[out]  SkuInfo  Receives the SKU identification bytes.

  @retval EFI_SUCCESS           SkuInfo was populated.
  @retval EFI_NOT_FOUND         Required HOB was not found.
**/
EFI_STATUS
GetMultiSkuInfoFromHob (
  OUT MULTI_SKU_INFO  *SkuInfo
  )
{
  VOID                   *HobList;
  EFI_HOB_GENERIC_HEADER *Hob;

  HobList = GetHobList ();
  if (HobList == NULL) {
    return EFI_NOT_FOUND;
  }

  //
  // Walk the HOB list looking for a GUID-ext HOB whose data matches
  // gLenovoMultiSkuHobGuid.
  //
  Hob = FindHobByGuid (&gLenovoMultiSkuHobGuid, HobList);
  while (Hob != NULL) {
    if (CompareGuid (
          (EFI_GUID *)((UINT8 *)Hob + sizeof(EFI_HOB_GENERIC_HEADER)),
          &gLenovoMultiSkuHobGuid
          ))
    {
      break;
    }
    Hob = FindHobByGuid (
            &gLenovoMultiSkuHobGuid,
            (VOID *)((UINT8 *)Hob + Hob->HobLength)
            );
  }

  if (Hob == NULL) {
    return EFI_NOT_FOUND;
  }

  //
  // The 3-byte SKU type is at offset 48 from the HOB base.
  // (8 bytes generic header + 16 bytes GUID + 24 bytes additional data)
  //
  SkuInfo->SkuType[0] = *((UINT8 *)Hob + 48);
  SkuInfo->SkuType[1] = *((UINT8 *)Hob + 49);
  SkuInfo->SkuType[2] = *((UINT8 *)Hob + 50);

  return EFI_SUCCESS;
}

// ---------------------------------------------------------------------------
// SKU protocol installation
// ---------------------------------------------------------------------------

/**
  Installs the SKU protocol.

  Publishes the SKU information via HandleProtocol using
  gLenovoSkuProtocolGuid.

  @retval EFI_SUCCESS  Protocol was installed (or HandleProtocol returned OK).
  @retval others       LocateProtocol failed or HandleProtocol error.
**/
EFI_STATUS
InstallSkuProtocol (
  VOID
  )
{
  EFI_STATUS       Status;
  MULTI_SKU_INFO   SkuInfo;
  EFI_HANDLE       NewHandle;

  NewHandle = NULL;

  MsDebugPrint (64, "InstallSkuProtocol is called\n");

  Status = GetMultiSkuInfoFromHob (&SkuInfo);
  if (EFI_ERROR (Status)) {
    //
    // SKU info not available -- still try to publish protocol
    // with whatever we have (or the HOB info is not critical).
    //
  }

  if (SkuInfo.SkuType[0] != 0x55) {
    //
    // Not a unified SKU byte: publish the protocol.
    //
    Status = gBootServices->HandleProtocol (
                              &NewHandle,
                              &gLenovoSkuProtocolGuid,
                              (VOID **)&SkuInfo
                              );
  }

  return Status;
}

// ---------------------------------------------------------------------------
// Module entry point
// ---------------------------------------------------------------------------

/**
  Main entry point for MultiSkuDistinctionDxe.

  Determines the system SKU/chassis type from the HOB list and decides
  how to publish the SKU information:

    1. Locates the HOB list and extracts Multi-SKU data.
    2. If SKU type byte == 0x55 (ASCII 'U'), registers an ExitBootServices
       callback via CreateEvent.  The callback (InstallSkuProtocol) installs
       the SKU protocol at boot-handoff time.
    3. Otherwise (SKU is not 'U'), InstallSkuProtocol is called directly.

  @param[in]  ImageHandle  Handle of the loaded driver image.
  @param[in]  SystemTable  Pointer to the UEFI system table.

  @retval EFI_SUCCESS           The driver executed without error.
  @retval EFI_OUT_OF_RESOURCES  Could not create the callback event.
  @retval others                Error reading HOB data.
**/
EFI_STATUS
EFIAPI
MultiSkuDistinctionDxeEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS       Status;
  MULTI_SKU_INFO   SkuInfo;
  EFI_EVENT        ExitBootEvent;

  Status = EFI_SUCCESS;

  //
  // Save global protocol references.
  //
  gImageHandle    = ImageHandle;
  if (!ImageHandle) {
    MsDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiBootServicesTableLib\\UefiBootServicesTableLib.c",
      51,
      "gImageHandle != ((void *) 0)"
      );
  }

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

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

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

  //
  // Pre-warm the HOB list cache.
  //
  GetHobList ();

  //
  // Read SKU info from HOB.
  //
  Status = GetMultiSkuInfoFromHob (&SkuInfo);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Check the SKU type byte.
  // 0x55 = 'U' (unified / unknown SKU) -- delay installation.
  //
  if (SkuInfo.SkuType[0] == 0x55) {
    //
    // Create an ExitBootServices notification event.
    //
    Status = gBootServices->CreateEvent (
                              EVT_NOTIFY_SIGNAL,
                              TPL_NOTIFY,
                              (EFI_EVENT_NOTIFY)InstallSkuProtocol,
                              NULL,
                              &ExitBootEvent
                              );
    if (!EFI_ERROR (Status)) {
      //
      // Register the configuration table entry for the SKU GUID.
      //
      gBootServices->InstallConfigurationTable (
                       &gLenovoMultiSkuConfigTableGuid,
                       NULL
                       );
    }

    return Status;
  }

  //
  // Non-'U' SKU: install protocol immediately.
  //
  return InstallSkuProtocol ();
}