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

;******************************************************************************
;* Copyright (c) 2016, Insyde Software Corporation. 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 <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/DebugLib.h>
#include <Library/PcdLib.h>
#include <Library/PeiServicesLib.h>
#include <IndustryStandard/Pci.h>
#include <Ppi/NvmeController.h>
#include <Ppi/Stall.h>
#include <Ppi/EndOfPeiPhase.h>

#define CLASSCODE_REGISTER              0x08
#define BASE_ADDRESS_REGISTER           0x10
#define NVME_CLASSCODE                  0x01080200
#define ROOT_BRIDGE_BUS_REGISTER        0x18
#define ROOT_BRIDGE_ADDRESS_REGISTER    0x20

#define PEI_NVME_SIGNATURE              SIGNATURE_32 ('N', 'V', 'M', 'C')

typedef struct {
  UINTN                         Signature;
  PEI_NVME_CONTROLLER_PPI       NvmeControllerPpi;
  EFI_PEI_PPI_DESCRIPTOR        PpiList;
  EFI_PEI_NOTIFY_DESCRIPTOR     NotifyList;
  EFI_PEI_STALL_PPI             *StallPpi;
  UINTN                         TotalNvmeControllers;
  UINTN                         MemBase;
  UINTN                         RootBridge;
  UINTN                         PciAddress;
} PEI_NVME_DEVICE;

#define PEI_NVME_DEVICE_FROM_THIS(a)        CR(a, PEI_NVME_DEVICE, NvmeControllerPpi, PEI_NVME_SIGNATURE)
#define PEI_NVME_DEVICE_FROM_NOTIFY_DESC(a) CR(a, PEI_NVME_DEVICE, NotifyList, PEI_NVME_SIGNATURE)

EFI_STATUS
GetNvmeController (
  IN CONST EFI_PEI_SERVICES             **PeiServices,
  IN PEI_NVME_CONTROLLER_PPI            *This,
  IN UINT8                              NvmeControllerId,
  OUT EFI_PHYSICAL_ADDRESS              *BaseAddress
  );


EFI_STATUS
EFIAPI
EndOfPeiPpiNotifyCallback (
  IN CONST EFI_PEI_SERVICES             **PeiServices,
  IN CONST EFI_PEI_NOTIFY_DESCRIPTOR    *NotifyDescriptor,
  IN VOID                               *Ppi
  );

//
// Globals
//
EFI_PEI_PPI_DESCRIPTOR mPpiList = {
  (EFI_PEI_PPI_DESCRIPTOR_PPI | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gPeiNvmeControllerPpiGuid,
  NULL
};

EFI_PEI_NOTIFY_DESCRIPTOR mNotifyList = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEfiEndOfPeiSignalPpiGuid,
  EndOfPeiPpiNotifyCallback
};

//
// Helper function
//
EFI_STATUS
EnableNvmeController (
  IN CONST EFI_PEI_SERVICES   **PeiServices,
  IN PEI_NVME_DEVICE          *PeiPchNvmeDev,
  IN UINT8                    NvmeControllerId
  );

/**

 @param [in]   FfsHeader
 @param [in]   PeiServices


**/
EFI_STATUS
InitializePchNvme (
  IN EFI_PEI_FILE_HANDLE        FileHandle,
  IN CONST EFI_PEI_SERVICES     **PeiServices
  )
{
  EFI_STATUS              Status;
  EFI_PEI_STALL_PPI       *StallPpi;
  PEI_NVME_DEVICE         *PeiPchNvmeDev;
  
  if (PcdGet32 (PcdNvmeRootPortAddress) == 0) {
    return EFI_UNSUPPORTED;
  }
  //
  // Get StallPpi
  //
  Status = (**PeiServices).LocatePpi (
                             PeiServices,
                             &gEfiPeiStallPpiGuid,
                             0,
                             NULL,
                             (VOID **)&StallPpi
                             );
  if (EFI_ERROR (Status)) {
    return EFI_UNSUPPORTED;
  }
  PeiPchNvmeDev = (PEI_NVME_DEVICE *)AllocatePages (1);
  if (PeiPchNvmeDev == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  ZeroMem (PeiPchNvmeDev, sizeof(PeiPchNvmeDev));

  PeiPchNvmeDev->Signature            = PEI_NVME_SIGNATURE;
  PeiPchNvmeDev->NvmeControllerPpi.GetNvmeController = GetNvmeController;
  PeiPchNvmeDev->PpiList              = mPpiList;
  PeiPchNvmeDev->PpiList.Ppi          = &PeiPchNvmeDev->NvmeControllerPpi;
  PeiPchNvmeDev->StallPpi             = StallPpi;
  PeiPchNvmeDev->NotifyList           = mNotifyList;
  PeiPchNvmeDev->TotalNvmeControllers = 1;
  //
  // Assign resources and enable NVME controllers
  //
  PeiPchNvmeDev->MemBase = PcdGet32 ( PcdNvmeMemBaseAddress );
  Status = EnableNvmeController (PeiServices, PeiPchNvmeDev, 0);
  if (EFI_ERROR(Status)) {
    return Status;
  }
  //
  // Install NVME Controller PPI
  //
  Status = PeiServicesInstallPpi (
             &PeiPchNvmeDev->PpiList
             );
  if (EFI_ERROR(Status)) {
    return Status;
  }
  //
  // Install notification in order to reset the NVME
  //
  Status = PeiServicesNotifyPpi (
             &PeiPchNvmeDev->NotifyList
             );
  return Status;
}

/**
 Retrieve NVME controller information

 @param [in]   PeiServices              Pointer to the PEI Services Table.
 @param [in]   This                     Pointer to PEI_NVME_CONTROLLER_PPI
 @param [in]   NvmeControllerId         NVME Controller ID
 @param [out]  BaseAddress              Result NVME base address

 @retval EFI_INVALID_PARAMETER          Invalid AhciControllerId is given
 @retval EFI_SUCCESS                    NVME controller information is retrieved successfully

**/
EFI_STATUS
GetNvmeController (
  IN CONST EFI_PEI_SERVICES         **PeiServices,
  IN PEI_NVME_CONTROLLER_PPI        *This,
  IN UINT8                          NvmeControllerId,
  OUT EFI_PHYSICAL_ADDRESS          *BaseAddress
  )
{
  PEI_NVME_DEVICE         *PeiPchNvmeDev;

  PeiPchNvmeDev = PEI_NVME_DEVICE_FROM_THIS (This);

  if (NvmeControllerId >= PeiPchNvmeDev->TotalNvmeControllers) {
    return EFI_INVALID_PARAMETER;
  }
  *BaseAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)PeiPchNvmeDev->MemBase;
  return EFI_SUCCESS;
}

/**
 Enable NVME controller

 @param [in]   PeiServices              Pointer to the PEI Services Table.
 @param [in]   PeiPchNvmeDev            Pointer to the PEI_NVME_DEVICE instance
 @param [in]   NvmeControllerId         NVME Controller ID

 @retval EFI_INVALID_PARAMETER          Invalid AhciControllerId is given
 @retval EFI_SUCCESS                    NVME controller information is retrieved successfully

**/
EFI_STATUS
EnableNvmeController (
  IN CONST EFI_PEI_SERVICES   **PeiServices,
  IN PEI_NVME_DEVICE          *PeiPchNvmeDev,
  IN UINT8                    NvmeControllerId
  )
{
  EFI_PEI_PCI_CFG2_PPI  *PciCfgPpi;
  UINT32                RootPort;
  UINT32                Bus;
  UINT32                Register;
  UINT32                Bridge;
  UINT32                Device;
  UINT32                OrgSecBusNum;  

  if (NvmeControllerId >= PeiPchNvmeDev->TotalNvmeControllers) {
    return EFI_INVALID_PARAMETER;
  }
  PciCfgPpi = (**PeiServices).PciCfg;
  //
  // Setup appropriate value to PCIE bridge
  //
  RootPort  = PcdGet32 ( PcdNvmeRootPortAddress );
  Bus       = RootPort >> 16;
  RootPort  &= 0xffff;
  Bridge    = RootPort << 8;
  Device    = Bus << 24;

  //
  // backup the orignal secondary and subordinary bus number
  //
  Register = 0;
  PciCfgPpi->Read (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint32,
               Bridge | ROOT_BRIDGE_BUS_REGISTER,
               &Register
               );
  OrgSecBusNum = Register;
  

  //
  // Assign bus number to PCIE bridge
  //
  Register  = (Bus << 8) + (Bus << 16);
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint32,
               Bridge | ROOT_BRIDGE_BUS_REGISTER,
               &Register
               );
  //
  // Discover NVME
  //
  Register = 0;
  PciCfgPpi->Read (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               Device,
               &Register
               );
  if (Register == 0xffff) {
    //
    // NVME not found, clear bus number to PCIE bridge
    //
    goto error;
  }
  //
  // Check the class code 
  //
  Register = 0;
  PciCfgPpi->Read (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint32,
               Device | CLASSCODE_REGISTER,
               &Register
               );
  Register &= 0xffffff00;
  if (Register != NVME_CLASSCODE) {
    //
    // Not NVME, clear bus number to PCIE bridge
    //
    goto error;
  }
  PeiPchNvmeDev->RootBridge = Bridge;
  PeiPchNvmeDev->PciAddress = Device;
  //
  // Assign address range for root bridge
  //
  Register = ((PeiPchNvmeDev->MemBase + 0x00100000) & 0xfff00000) + ((PeiPchNvmeDev->MemBase & 0xfff00000) >> 16);
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint32,
               Bridge | ROOT_BRIDGE_ADDRESS_REGISTER,
               &Register
               );
  //
  // Assign address prefetchable range for root bridge
  //
  Register = ((PeiPchNvmeDev->MemBase + 0x00200000) & 0xfff00000) + (((PeiPchNvmeDev->MemBase + 0x00100000) & 0xfff00000) >> 16);
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint32,
               Bridge | ROOT_BRIDGE_ADDRESS_REGISTER + 4,
               &Register
               );
  //
  // Assign base address register to NVME
  //
  Register = PeiPchNvmeDev->MemBase;
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint32,
               Device | BASE_ADDRESS_REGISTER,
               &Register
               );
  //
  // Enable root bridge
  //
  Register = 0;
  PciCfgPpi->Read (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               Bridge | PCI_COMMAND_OFFSET,
               &Register
               );
  Register |= 0x06;
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               Bridge | PCI_COMMAND_OFFSET,
               &Register
               );
  //
  // Enable NVME
  //
  Register = 0;
  PciCfgPpi->Read (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               Device | PCI_COMMAND_OFFSET,
               &Register
               );
  Register |= 0x06;
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               Device | PCI_COMMAND_OFFSET,
               &Register
               );
  return EFI_SUCCESS;
error:
  Register = OrgSecBusNum;
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               Bridge | PCI_COMMAND_OFFSET,
               &Register
               );
  return EFI_NOT_FOUND;
}

/**
 Register notify ppi to reset the NVME.


 @param[in]   PeiServices               Pointer to the PEI Services Table.
 @param[in]   NotifyDescriptor          Pointer to the notify descriptor

 @retval EFI_SUCCESS

**/
EFI_STATUS
EFIAPI
EndOfPeiPpiNotifyCallback (
  IN CONST EFI_PEI_SERVICES               **PeiServices,
  IN CONST EFI_PEI_NOTIFY_DESCRIPTOR      *NotifyDescriptor,
  IN VOID                                 *Ppi
  )
{
  EFI_PEI_PCI_CFG2_PPI        *PciCfgPpi;
  UINT8                       *BaseAddress;
  UINT32                      Register;
  PEI_NVME_DEVICE             *PeiPchNvmeDev;
  
  PeiPchNvmeDev = PEI_NVME_DEVICE_FROM_NOTIFY_DESC (NotifyDescriptor);
  //
  // Disable the NVME
  //
  BaseAddress = (UINT8*)(UINTN)PeiPchNvmeDev->MemBase + 0x14;
  Register    = *(UINT32*)(UINTN)BaseAddress;
  Register    &= ~0x01;
  *(UINT32*)(UINTN)BaseAddress = Register;
  PeiPchNvmeDev->StallPpi->Stall (
                             PeiServices,
                             PeiPchNvmeDev->StallPpi,
                             500 * 1000
                             );
  PciCfgPpi = (**PeiServices).PciCfg;
  //
  // Disable NVME
  //
  Register = 0;
  PciCfgPpi->Read (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               PeiPchNvmeDev->PciAddress | PCI_COMMAND_OFFSET,
               &Register
               );
  Register &= ~0x06;
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               PeiPchNvmeDev->PciAddress | PCI_COMMAND_OFFSET,
               &Register
               );
  //
  // Disable root bridge
  //
  Register = 0;
  PciCfgPpi->Read (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               PeiPchNvmeDev->RootBridge | PCI_COMMAND_OFFSET,
               &Register
               );
  Register &= ~0x06;
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint16,
               PeiPchNvmeDev->RootBridge | PCI_COMMAND_OFFSET,
               &Register
               );
  //
  // Clear bus number for root bridge
  //
  Register = 0;
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint32,
               PeiPchNvmeDev->RootBridge | ROOT_BRIDGE_BUS_REGISTER,
               &Register
               );
  //
  // Clear address range for root bridge
  //
  PciCfgPpi->Write (
               PeiServices,
               PciCfgPpi,
               EfiPeiPciCfgWidthUint32,
               PeiPchNvmeDev->RootBridge | ROOT_BRIDGE_ADDRESS_REGISTER,
               &Register
               );
  return EFI_SUCCESS;
}