/** @file
  Implementation of PchAhciPei module for Crisis Recovery

;******************************************************************************
;* Copyright (c) 2014 - 2016, Insyde Software Corp. All Rights Reserved.
;*
;* You may not reproduce, distribute, publish, display, perform, modify, adapt,
;* transmit, broadcast, present, recite, release, license or otherwise exploit
;* any part of this publication in any form, by any means, without the prior
;* written permission of Insyde Software Corporation.
;*
;******************************************************************************
*/
#include <PiPei.h>
#include <Guid/H2OCp.h>
#include <Library/DebugLib.h>
#include <Library/PeiServicesLib.h>
#include <Register/SataRegs.h>
#include <Library/IoLib.h>
#include <Library/PciSegmentLib.h>
#include <Library/H2OCpLib.h>
#include <IndustryStandard/Pci.h>
#include <Library/SataSocLib.h>
//[-start-191225-IB16740000-add]// for SataDevNumber & SataFuncNumber function
#include <Library/PchPciBdfLib.h>
//[-end-191225-IB16740000-add]//

//[-start-200506-IB17800061-add]//
//ADL RC 1151 change satalib satasoclib, but satasoclib does not include SataRegBase functioin .
// keep orignal design and move SataRegBase to here.
/**
  Get SATA controller address that can be passed to the PCI Segment Library functions.

  @param[in]  SataCtrlIndex       SATA controller index

  @retval SATA controller address in PCI Segment Library representation
**/
UINT64
SataRegBase (
  IN UINT32      SataCtrlIndex
  )
{
  ASSERT (SataCtrlIndex < MaxSataControllerNum ());

  return SataPciCfgBase (SataCtrlIndex);
}
//[-end-200506-IB17800061-add]//


/**
  Handle PCH AHCI controller init.

  @param[in] Event           A pointer to the Event that triggered the callback.
  @param[in] Handle          Checkpoint handle.

**/
VOID
EFIAPI
PciEnumUpdateDevResourcesHandler (
  IN EFI_EVENT                              Event,
  IN H2O_CP_HANDLE                          Handle
  )
{
  EFI_STATUS                                           Status;
  H2O_PEI_CP_PCI_ENUM_UPDATE_DEV_RESOURCES             *CpPciEnumUpdateDevResourceData;
  UINT32                                               Index;
  UINT32                                               ControllerNum;
  UINTN                                                PciSataRegBase;
  UINT32                                               AhciBar;
  UINT8                                                Mask;
  UINT32                                               Register;

  Status = H2OCpLookup (Handle, (VOID **)&CpPciEnumUpdateDevResourceData, NULL);
  if (EFI_ERROR (Status)) {
    DEBUG_CP ((DEBUG_ERROR, "Checkpoint Data Not Found: %x (%r)\n", Handle, Status));
    DEBUG_CP ((DEBUG_ERROR, "   %a\n", __FUNCTION__));
    return;
  }

  if (CpPciEnumUpdateDevResourceData->Bus != 0) {
    //
    // Assuming AHCI HC is behind bus 0
    //
    return;
  }

  ControllerNum = MaxSataControllerNum ();
  for (Index = 0; Index < ControllerNum; Index++) {
    if ((CpPciEnumUpdateDevResourceData->Device == SataDevNumber (Index)) &&
        (CpPciEnumUpdateDevResourceData->Function == SataFuncNumber (Index))) {
        break;
    }
  }

  if (Index == ControllerNum) {
    return;
  }

  //
  // Get SATA controller address
  //
  PciSataRegBase = (UINTN)SataRegBase (Index);
  //
  // Set PCS: enable all, PCH-LP [2:0], PCH-H [7:0]
  //
  PciSegmentOr16 (PciSataRegBase + R_SATA_CFG_PCS, (UINT16)B_SATA_CFG_PCS_PXE_MASK);
  //
  // Set PCS: Present all, PCH-LP [18:16], PCH-H [23:16]
  //
  PciSegmentOr16 (PciSataRegBase + R_SATA_CFG_PCS + 2, (UINT16)B_SATA_CFG_PCS_PXE_MASK);
  DEBUG ((DEBUG_ERROR, "PchAhci: R_SATA_CFG_PCS = 0x%x \n", PciSegmentRead32 (PciSataRegBase + R_SATA_CFG_PCS)));
  //
  // Assign base address register to AHCI ABAR
  //
  AhciBar = CpPciEnumUpdateDevResourceData->PciBar[5];
  PciSegmentWrite32 (PciSataRegBase + R_SATA_CFG_AHCI_BAR, AhciBar);
  DEBUG ((DEBUG_ERROR, "PchAhci: R_SATA_CFG_AHCI_BAR = 0x%x \n", PciSegmentRead32 (PciSataRegBase + R_SATA_CFG_AHCI_BAR)));
  //
  // Enable AHCI
  //
  Mask = (EFI_PCI_COMMAND_BUS_MASTER | EFI_PCI_COMMAND_MEMORY_SPACE | EFI_PCI_COMMAND_IO_SPACE);
  PciSegmentOr16 (PciSataRegBase + PCI_COMMAND_OFFSET, (UINT16)Mask);
  DEBUG ((DEBUG_ERROR, "PchAhci: PCI_COMMAND_OFFSET = 0x%x \n", MmioRead16((PciSataRegBase + PCI_COMMAND_OFFSET))));
  //
  // Assign ABAR PI register
  //
  Mask = PciSegmentRead8 (PciSataRegBase + R_SATA_CFG_PCS + 2);
  Register = MmioRead32 (AhciBar + R_SATA_MEM_AHCI_PI);
  Register &= (UINT32)(~B_SATA_MEM_AHCI_PI_PORT_MASK);
  Register |= (UINT32)Mask;
  MmioWrite32 (AhciBar + R_SATA_MEM_AHCI_PI, Register);
  DEBUG ((DEBUG_ERROR, "PchAhci: R_SATA_MEM_AHCI_PI = 0x%x (0x%x)\n", MmioRead32 (AhciBar + R_SATA_MEM_AHCI_PI), Register));
  //
  // Assign ABAR CAP register
  // Update the Host Capabilites Register
  // NOTE: Many of the bits in this register are R/WO (Read/Write Once)
  //
  Register = PcdGet32 (PcdAbarCapDefault);
  MmioWrite32 (AhciBar + R_SATA_MEM_AHCI_CAP, Register);
  DEBUG ((DEBUG_ERROR, "PchAhci: R_SATA_MEM_AHCI_CAP = 0x%x (0x%x)\n", MmioRead32(AhciBar + R_SATA_MEM_AHCI_CAP), Register));

  //
  // Updating checkpoint data if need
  //
  CpPciEnumUpdateDevResourceData->Status = H2O_CP_TASK_SKIP;
  DEBUG ((DEBUG_INFO, "Checkpoint Task Skip: %g\n", &gH2OPeiCpPciEnumUpdateDevResourcesGuid));
  DEBUG ((DEBUG_INFO, "   %a\n", __FUNCTION__));
}

/**
  PEIM entry.

  @param  FileHandle             Handle of the file being invoked.
  @param  PeiServices            Describes the list of possible PEI Services.

  @retval EFI_SUCCESS            The driver is successfully initialized.
  @retval Others                 Can't initialize the driver.

**/
EFI_STATUS
InitializePchAhci (
  IN       EFI_PEI_FILE_HANDLE        FileHandle,
  IN CONST EFI_PEI_SERVICES     **PeiServices
  )
{
  EFI_STATUS                    Status;
  H2O_CP_HANDLE                 CpHandle;

  if (FeaturePcdGet (PcdH2OPeiCpPciEnumUpdateDevResourcesSupported)) {
    Status = H2OCpRegisterHandler (
               &gH2OPeiCpPciEnumUpdateDevResourcesGuid,
               PciEnumUpdateDevResourcesHandler,
               H2O_CP_LOW,
               &CpHandle
               );
    if (EFI_ERROR (Status)) {
      DEBUG_CP ((DEBUG_ERROR, "Checkpoint Register Fail: %g (%r)\n", &gH2OPeiCpPciEnumUpdateDevResourcesGuid, Status));
      return Status;
    }
    DEBUG_CP ((DEBUG_INFO, "Checkpoint Registered: %g (%r)\n", &gH2OPeiCpPciEnumUpdateDevResourcesGuid, Status));
  }

  return EFI_SUCCESS;
}