/** @file
  PciDxeInit - PCI DXE Initialization Driver

  This driver initializes the PCI subsystem during the DXE phase. It is part of
  the AMI ModulePkg for the Lenovo HR650X platform (Purley/Intel grantley-EP).

  Responsibilities:
    - Installs the AmiPciBusSetupOverride protocol with device-specific hooks
    - Manages the PCI bus override table (per-device config callbacks)
    - Handles VGA device enumeration and selection (onboard vs offboard)
    - Manages S3 boot script save/restore via SmmLockBox
    - Initializes PCIe segments and MMIO base
    - Provides RAID driver binding for onboard SATA/AHCI controllers
    - Restricts MMIO resources (e.g., USB under 4GB)
    - Port degradation for 10GbE BAR

  Build info: DEBUG_VS2015 X64, HR6N0XMLK
  Source path: AmiModulePkg\PCI\PciDxeInit

  Copyright (c) American Megatrends Inc. (AMI). All rights reserved.
**/

#include "PciDxeInit.h"

//
// ==========================================================================
// Global Data
// ==========================================================================
//

//
// EFI library globals
//
EFI_HANDLE               gImageHandle_Private = NULL;
EFI_SYSTEM_TABLE         *gSystemTable = NULL;
EFI_RUNTIME_SERVICES     *gRuntimeServices = NULL;

//
// PCIe MMIO base address (from PCI Express Base Address PCD)
//
UINT64                   gPciExpressBase = 0;

//
// Protocol notification events
//
EFI_EVENT                gEventAmiPciBusSetupOverride = NULL;
EFI_EVENT                gEventDxeSmmReadyToLock = NULL;
EFI_EVENT                gEventS3BootScriptSave = NULL;

//
// S3 boot script state
//
S3_BOOT_SCRIPT_CONTEXT   *gS3BootScriptContext = NULL;
S3_BOOT_SCRIPT_CONTEXT   *gS3BootScriptContextBackup = NULL;
BOOLEAN                  gBootScriptBufferAllocated = FALSE;
BOOLEAN                  gBootScriptLockedBySmm = FALSE;

//
// VGA state
//
BOOLEAN                  gAtleastOneVideoFound = FALSE;
UINTN                    gNumOffboardVideoPresent = 0;

//
// SMM LockBox communication
//
BOOLEAN                  gSmmLockBoxInitialized = FALSE;

//
// Setup variable support globals
//
BOOLEAN                  gSetupVariableExists = FALSE;

//
// ==========================================================================
// Debug Support
// ==========================================================================
//

/**
  Reads the current debug level from CMOS.

  @return Debug level mask.
**/
UINTN
EFIAPI
GetCurrentDebugLevel (
    VOID
    )
{
    UINT8   CmosData;

    //
    // Read CMOS index 0x4B (debug level register)
    //
    IoWrite8 (0x70, (IoRead8 (0x70) & 0x80) | 0x4B);
    CmosData = IoRead8 (0x71);

    if (CmosData > 3) {
        if (CmosData == 0) {
            CmosData = (MmioRead8 (0xFDAF0490) & 2) | 1;
        }
    }

    if ((CmosData - 1) > 0xFD) {
        return 0;
    }

    if (CmosData == 1) {
        return BIT31 | 4;
    }

    return (UINTN)-1;
}

//
// ==========================================================================
// PciDxeInit Entry Point
// ==========================================================================
//

/**
  Module entry point for PciDxeInit.

  @param  ImageHandle  The module's image handle.
  @param  SystemTable  The UEFI system table.

  @return EFI_SUCCESS or EFI error code.
**/
EFI_STATUS
EFIAPI
PciDxeInitEntryPoint (
    IN EFI_HANDLE        ImageHandle,
    IN EFI_SYSTEM_TABLE  *SystemTable
    )
{
    EFI_BOOT_SERVICES     *BootServices;
    EFI_RUNTIME_SERVICES  *RuntimeServices;
    EFI_STATUS            Status;
    EFI_HANDLE            Handle;

    //
    // Initialize the driver
    //
    PciDxeInitDriverInit (ImageHandle);

    //
    // Save EFI system table pointers (may already be set by lib constructor)
    //
    if (gSystemTable == NULL) {
        gSystemTable        = SystemTable;
        BootServices        = SystemTable->BootServices;
        RuntimeServices     = SystemTable->RuntimeServices;
        gBS                 = BootServices;
        gRT                 = RuntimeServices;
    }

    //
    // Install the AmiPciBusSetupOverride protocol
    //
    Handle = NULL;
    Status = gBS->InstallMultipleProtocolInterfaces (
                    &Handle,
                    &gAmiPciBusSetupOverrideProtocolGuid,
                    &gAmiPciBusSetupOverride,
                    NULL
                    );
    if (EFI_ERROR (Status)) {
        DEBUG ((EFI_D_ERROR, "\nASSERT_EFI_ERROR (Status = %r)\n", Status));
        ASSERT_EFI_ERROR (FALSE);
        return Status;
    }

    return Status;
}

//
// ==========================================================================
// Driver Initialization
// ==========================================================================
//

/**
  Driver initialization routine.

  @param  ImageHandle  The module's image handle.

  @return EFI_STATUS.
**/
EFI_STATUS
PciDxeInitDriverInit (
    IN  EFI_HANDLE  ImageHandle
    )
{
    EFI_STATUS                  Status;
    UINT64                      PciExpressBaseAddress;
    UINTN                       SegBusTableSize;
    UINT8                       SetupByte;
    UINT16                      CmosFlags;
    BOOLEAN                     InterruptsEnabled;
    UINT64                      DelayStart;

    //
    // Save global handles
    //
    gImageHandle_Private = ImageHandle;

    if (gImageHandle_Private == NULL) {
        DEBUG ((EFI_D_ERROR, "gImageHandle != NULL\n"));
    }
    if (gST == NULL) {
        DEBUG ((EFI_D_ERROR, "gST != NULL\n"));
    }
    if (gBS == NULL) {
        DEBUG ((EFI_D_ERROR, "gBS != NULL\n"));
    }
    if (gRT == NULL) {
        DEBUG ((EFI_D_ERROR, "gRT != NULL\n"));
    }

    //
    // Get DxeServicesTable
    //
    Status = gBS->LocateProtocol (
                    &gEfiDxeServicesTableProtocolGuid,
                    NULL,
                    (VOID **)&gDS
                    );
    ASSERT_EFI_ERROR (Status);
    if (gDS == NULL) {
        DEBUG ((EFI_D_ERROR, "gDS != NULL\n"));
    }

    //
    // Get PCI Express base address PCD or locate MmPciBase protocol
    //
    gPciExpressBase = PcdGet64 (PcdPciExpressBaseAddress);

    if (gPciExpressBase == 0) {
        Status = gBS->LocateProtocol (
                        &gEfiMmPciBaseProtocolGuid,
                        NULL,
                        (VOID **)&gPciExpressBase
                        );
        if (Status == EFI_SUCCESS && gPciExpressBase == 0) {
            DEBUG ((EFI_D_ERROR, "mPciUsra != NULL\n"));
        }
    }

    //
    // Initialize HOB list
    //
    Status = EfiGetSystemConfigurationTable (
               &gEfiHobListGuid,
               (VOID **)&gHobList
               );
    ASSERT_EFI_ERROR (Status);
    if (gHobList == NULL) {
        DEBUG ((EFI_D_ERROR, "mHobList != NULL\n"));
    }

    //
    // Enable PCI config space access
    //
    CmosFlags = (UINT16)AsmReadEflags ();
    AsmEnableInterrupt ();

    //
    // Small delay using RDTSC (~10us)
    //
    DelayStart = AsmReadTsc () & 0xFFFFFF;
    AsmDisableInterrupt ();

    //
    // Write PCI configuration enable bit
    //
    IoWrite32 (0xCF8, 0x80000000);

    //
    // Wait for enable to stabilize
    //
    while ((((UINT32)DelayStart + 357 - (UINT32)AsmReadTsc ()) & 0x800000) == 0) {
        AsmPause ();
    }
    AsmEnableInterrupt ();

    //
    // Restore interrupt state
    //
    if ((CmosFlags & BIT9) != 0) {
        AsmEnableInterrupt ();
    } else {
        AsmDisableInterrupt ();
    }

    //
    // Initialize S3 boot script library
    //
    Status = S3BootScriptLibInitialize ();
    ASSERT_EFI_ERROR (Status);

    //
    // Get PCIe segment/bus table size from PCD
    //
    SegBusTableSize = PcdGetSize (PcdPcieSegBusTable);
    ASSERT (sizeof (PCIE_SEG_BUS_TABLE) >= SegBusTableSize);

    PciExpressBaseAddress = PcdGet64 (PcdPcieSegBusTable);

    //
    // Allocate and initialize segment/bus table
    //
    if (PciExpressBaseAddress != 0) {
        gPcieSegBusTable = AllocatePool (SegBusTableSize);
        if (gPcieSegBusTable != NULL) {
            CopyMem (
              gPcieSegBusTable,
              (VOID *)(UINTN)PciExpressBaseAddress,
              SegBusTableSize
              );
        }
    }

    //
    // Read PCI Bus Setup Override setup variable
    //
    Status = GetSetupVariableByIndex (260, 0, &SetupByte);
    if (EFI_ERROR (Status)) {
        SetupByte = 0xFE;
    }

    //
    // Create readiness notification event for PCI enum
    //
    Status = gBS->CreateEvent (
                    EVT_NOTIFY_SIGNAL,
                    TPL_CALLBACK,
                    (EFI_EVENT_NOTIFY)PciDevCommonHook,
                    NULL,
                    &gEventAmiPciBusSetupOverride
                    );
    if (EFI_ERROR (Status)) {
        DEBUG ((EFI_D_ERROR, "Event != NULL\n"));
    }

    return Status;
}

//
// ==========================================================================
// PCI Device Common Hook Dispatcher
// ==========================================================================
//

/**
  Dispatches hooks registered in the override table for a given device.

  @param  This        AmiPciBusSetupOverride protocol instance.
  @param  Context     Hook context (output).
  @param  OverrideTable Current override entry.

  @return EFI_STATUS.
**/
EFI_STATUS
EFIAPI
PciInitHookDispatcher (
    IN  VOID                   *PciDevicePrivateData,
    OUT VOID                   **HookContext,
    IN  PCI_BUS_OVERRIDE_ENTRY *OverrideTable
    )
{
    PCI_BUS_OVERRIDE_ENTRY    *Entry;
    PCI_DEVICE_CALLBACK        Callback;
    EFI_STATUS                 Status;
    EFI_STATUS                 ReturnStatus;
    UINTN                      Index;

    ReturnStatus = EFI_NOT_FOUND;
    Entry = gPciOverrideTable;

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

    if (OverrideTable == NULL) {
        return EFI_INVALID_PARAMETER;
    }

    if (*(UINT64 *)OverrideTable != PCI_BUS_OVERRIDE_TABLE_SIGNATURE) {
        return EFI_INVALID_PARAMETER;
    }

    for (Index = 0; ; Index += sizeof (PCI_BUS_OVERRIDE_ENTRY)) {
        Callback = (PCI_DEVICE_CALLBACK)(UINTN)Entry[Index].Callback;
        if (Callback == NULL) {
            break;
        }

        if (((*(UINT16 *)(PciDevicePrivateData + FIELD_OFFSET)) & Entry[Index].Match.VenMask) == Entry[Index].Match.VenId &&
            ((*(UINT16 *)(PciDevicePrivateData + FIELD_OFFSET + 2)) & Entry[Index].Match.DevMask) == Entry[Index].Match.DevId) {

            DEBUG ((
                EFI_D_INFO,
                "DevCmnHook: B%X|D%X|F%X (VID=0x%X; DID=0x%X; Mask(0x%X:0x%X) Step=%d",
                *(UINT8 *)(PciDevicePrivateData + 219),   // Bus
                *(UINT8 *)(PciDevicePrivateData + 218),   // Device
                *(UINT8 *)(PciDevicePrivateData + 217),   // Function
                Entry[Index].Match.VenId,
                Entry[Index].Match.DevId,
                Entry[Index].Match.VenMask,
                Entry[Index].Match.DevMask,
                Entry[Index].InitStep
                ));

            Status = Callback (PciDevicePrivateData, HookContext, OverrideTable);
            DEBUG ((EFI_D_INFO, ")=%r\n", Status));

            if (EFI_ERROR (Status) && Status != EFI_NOT_FOUND) {
                return Status;
            }

            if (!EFI_ERROR (Status)) {
                ReturnStatus &= ~EFI_ERROR_BIT;
            }
        }
    }

    return ReturnStatus;
}

//
// ==========================================================================
// S3 Boot Script Management
// ==========================================================================
//

/**
  Initialize the S3 boot script library.

  Allocates the boot script buffer and registers events for:
  - DXE SMM Ready To Lock (close boot script)
  - S3 ready (save boot script to lockbox)

  @return EFI_STATUS.
**/
EFI_STATUS
S3BootScriptLibInitialize (
    VOID
    )
{
    EFI_STATUS             Status;

    DEBUG ((EFI_D_INFO, "%a() in %a module\n", __FUNCTION__, "PciDxeInit"));

    if (gEventDxeSmmReadyToLock == NULL) {
        //
        // Allocate boot script buffer
        //
        gS3BootScriptContext = (S3_BOOT_SCRIPT_CONTEXT *)AllocatePages (EFI_SIZE_TO_PAGES (0x20));
        if (gS3BootScriptContext == NULL) {
            //
            // Fallback: allocate with AllocatePages
            //
            Status = gBS->AllocatePages (
                            AllocateAnyPages,
                            EfiBootServicesData,
                            1,
                            &gS3BootScriptContext
                            );
            ASSERT_EFI_ERROR (Status);

            gS3BootScriptContext = (S3_BOOT_SCRIPT_CONTEXT *)(UINTN)gS3BootScriptContext;
            gBootScriptBufferAllocated = TRUE;
        }

        if (gS3BootScriptContext != NULL) {
            ZeroMem (gS3BootScriptContext, sizeof (S3_BOOT_SCRIPT_CONTEXT));
        }

        //
        // Create ReadyToLock event
        //
        Status = gBS->CreateEvent (
                        EVT_NOTIFY_SIGNAL,
                        TPL_CALLBACK,
                        (EFI_EVENT_NOTIFY)S3BootScriptEntryClose,
                        NULL,
                        &gEventDxeSmmReadyToLock
                        );
        if (Status == EFI_SUCCESS) {
            DEBUG ((EFI_D_INFO, "mEventDxeSmmReadyToLock != NULL\n"));
        }

        //
        // Try to locate SMM Communication protocol
        //
        Status = gBS->LocateProtocol (
                        &gEfiSmmCommunicationProtocolGuid,
                        NULL,
                        (VOID **)&gSmmCommunication
                        );
        if (Status == EFI_SUCCESS) {
            Status = gSmmCommunication->Register (
                                        gSmmCommunication,
                                        (EFI_EVENT_NOTIFY)SmmLockBoxEntryClose,
                                        &gEventS3BootScriptSave
                                        );
            ASSERT_EFI_ERROR (Status);

            Status = gSmmCommunication->Register (
                                        gSmmCommunication,
                                        (EFI_EVENT_NOTIFY)SmmLockBoxEntryClose,
                                        &gEventS3BootScriptSave
                                        );
            ASSERT_EFI_ERROR (Status);

            Status = gSmmCommunication->Register (
                                        gSmmCommunication,
                                        (EFI_EVENT_NOTIFY)SmmLockBoxBackupEntry,
                                        &gEventS3BootScriptSave
                                        );
            ASSERT_EFI_ERROR (Status);
        }
    }

    if (gSmmCommunication != NULL && gS3BootScriptContextBackup == NULL) {
        //
        // Allocate backup buffer from SMM communication buffer
        //
        Status = gSmmCommunication->GetBuffer (
                                    gSmmCommunication,
                                    32,
                                    (VOID **)&gS3BootScriptContextBackup
                                    );
        if (EFI_ERROR (Status)) {
            Status = gSmmCommunication->AllocateBuffer (
                                        gSmmCommunication,
                                        32,
                                        (VOID **)&gS3BootScriptContextBackup
                                        );
            ASSERT_EFI_ERROR (Status);

            gSmmLockBoxInitialized = TRUE;
            ZeroMem (gS3BootScriptContextBackup, 0x20);
        }
    }

    return EFI_SUCCESS;
}

/**
  Deinitialize the S3 boot script library on ReadyToBoot.

  @return EFI_STATUS.
**/
EFI_STATUS
S3BootScriptLibDeinitialize (
    VOID
    )
{
    EFI_STATUS  Status;

    Status = gBS->CloseEvent ((EFI_EVENT)&gEventAmiPciBusSetupOverride);
    ASSERT_EFI_ERROR (Status);

    if (gBootScriptBufferAllocated) {
        Status = gBS->FreePages (
                        (EFI_PHYSICAL_ADDRESS)(UINTN)gS3BootScriptContext,
                        1
                        );
        ASSERT_EFI_ERROR (Status);

        Status = PcdSet64 (PcdPcieSegBusTable, 0);
        ASSERT_EFI_ERROR (Status);
    }

    if (gSmmCommunication != NULL && gSmmLockBoxInitialized) {
        Status = gSmmCommunication->FreeBuffer (
                                    gSmmCommunication,
                                    (VOID *)gS3BootScriptContextBackup
                                    );
        ASSERT_EFI_ERROR (Status);

        Status = PcdSet64 (PcdPcieSegBusTableAlt, 0);
        ASSERT_EFI_ERROR (Status);
    }

    return EFI_SUCCESS;
}