/** @file

;******************************************************************************
;* Copyright (c) 2018 - 2020, 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 <SataDevSleep.h>
#include <Library/PchInfoLib.h>
#include <SetupVariable.h>
#include <Library/SataSocLib.h>

EFI_STATUS
EFIAPI
SataDevSleepEntryPoint (
  IN EFI_HANDLE         ImageHandle,
  IN EFI_SYSTEM_TABLE   *SystemTable
  )
{
  EFI_STATUS           Status;
  EFI_EVENT            Event;
  VOID                 *Registration;

  Status = EFI_SUCCESS;
  DEBUG((DEBUG_INFO | DEBUG_ERROR, "SATADevSleepEntryPoint Entry\n"));

//  if (IsPchLp()) {
  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  SATADevSleepCallback,
                  NULL,
                  &Event
                  );
  if (!EFI_ERROR (Status)) {
    Status = gBS->RegisterProtocolNotify (
                    &gEfiAtaPassThruProtocolGuid,
                    Event,
                    &Registration
                    );
  }
//  } else {
//    DEBUG ((DEBUG_ERROR, "Platform does not support DevSleep \n"));
//  }
  DEBUG((DEBUG_INFO | DEBUG_ERROR, "SATADevSleepEntryPoint End\n"));
  return EFI_SUCCESS;
}

/**
  Enable DEVSLP command of the disk if supported.
  The following steps are according to PCH spec <DevSleep Enabling>

**/
VOID
EFIAPI
SATADevSleepCallback (
  IN EFI_EVENT                          Event,
  IN VOID                               *Context
  )
{
  EFI_STATUS                         Status;
  UINTN                              HandleCount;
  UINT32                             HandleIndex;
  EFI_HANDLE                         *HandleBuffer;
  EFI_ATA_PASS_THRU_PROTOCOL         *AtaPassThruPtr;
  EFI_DEVICE_PATH_PROTOCOL           *AtaPassThruDevPath;
  EFI_DEVICE_PATH_PROTOCOL           *TmpDevicePath;
  BOOLEAN                            IsMatched;
  UINT32                             AhciBar;
  UINT32                             PciD31F2RegBase;
  UINT32                             Data32Or;
  UINT16                             PortIndex;
  DEV_SATA_DVSLEEP_DATA              PortData[PCH_MAX_SATA_PORTS] = {0};
  PCH_SETUP                          *PchSetup;
  UINTN                              PchSetupSize;
  BOOLEAN                            UpdatePchSetup;
  UINT8                              MaxSataPorts;
  BOOLEAN                            SetupPortDevSleepEnable;
  UINT32                             Capability2;
  UINT32                             Data32;

  DEBUG((DEBUG_INFO | DEBUG_ERROR, "SATADevSleepCallback Start\n"));

  Status = EFI_SUCCESS;
  HandleCount  = 0;
  HandleBuffer = NULL;
  AtaPassThruDevPath = NULL;
  IsMatched = FALSE;
  PortIndex = 0;
  PchSetup = NULL;
  UpdatePchSetup = FALSE;
  MaxSataPorts = MaxSataPortNum (SATA_1_CONTROLLER_INDEX);
  SetupPortDevSleepEnable = FALSE;

  PciD31F2RegBase = (UINT32) MmPciAddress (0, 0, PCI_DEVICE_NUMBER_PCH_SATA, PCI_FUNCTION_NUMBER_PCH_SATA, 0);
  AhciBar = Mmio32 (PciD31F2RegBase , R_SATA_CFG_AHCI_BAR);
  Capability2     = Mmio32 (AhciBar, R_SATA_MEM_AHCI_CAP2);

  PchSetupSize = sizeof (PCH_SETUP);
  PchSetup = AllocateZeroPool (PchSetupSize);
  if (PchSetup == NULL) {
    return;
  }
  Status = gRT->GetVariable (
                  PCH_SETUP_VARIABLE_NAME,
                  &gPchSetupVariableGuid,
                  NULL,
                  &PchSetupSize,
                  PchSetup
                  );
  if (EFI_ERROR (Status)) {
    goto Error2;
  }

  //
  //  Check there has any port enable DevSlp function on platform or not.
  //
  for (PortIndex = 0; PortIndex < MaxSataPorts; PortIndex++) {
    if (PchSetup->PxDevSlp[PortIndex] == TRUE) {
      PortData[PortIndex].State.SetupDevsleepEnable = PchSetup->PxDevSlp[PortIndex];
      SetupPortDevSleepEnable = TRUE;
    }
  }
  if (!SetupPortDevSleepEnable) {
   goto Error2;
  }

  Status = gBS->LocateHandleBuffer (
          ByProtocol,
          &gEfiAtaPassThruProtocolGuid,
          NULL,
          &HandleCount,
          &HandleBuffer
          );
  if (EFI_ERROR(Status) || (HandleCount == 0)) {
    DEBUG((DEBUG_ERROR, "LocateHandleBuffer Error!, Status = %r\n", Status));
    goto Error1;
  }

  for (HandleIndex = 0; HandleIndex < HandleCount; HandleIndex++) {

    Status = gBS->OpenProtocol (
                    HandleBuffer[HandleIndex],
                    &gEfiDevicePathProtocolGuid,
                    (VOID **) &AtaPassThruDevPath,
                    NULL,
                    NULL,
                    EFI_OPEN_PROTOCOL_GET_PROTOCOL
                    );
    if (EFI_ERROR (Status)) {
      continue;
    }
    TmpDevicePath = AtaPassThruDevPath;
    while (!IsDevicePathEnd (TmpDevicePath)) {
      if((TmpDevicePath->Type == HARDWARE_DEVICE_PATH) && (TmpDevicePath->SubType == HW_PCI_DP)) {
        if((((PCI_DEVICE_PATH*)TmpDevicePath)->Device == PCI_DEVICE_NUMBER_PCH_SATA) &&
            (((PCI_DEVICE_PATH*)TmpDevicePath)->Function) == PCI_FUNCTION_NUMBER_PCH_SATA) {
            //
            // Controller matched
            //
            IsMatched = TRUE;
            break;
        }
      }
      TmpDevicePath = NextDevicePathNode (TmpDevicePath);
    }
    if(IsMatched) {
      break;
    }

  }
  Status = gBS->HandleProtocol (
                  HandleBuffer[HandleIndex],
                  &gEfiAtaPassThruProtocolGuid,
                  (VOID**)&AtaPassThruPtr
                  );
  if (EFI_ERROR (Status)) {
    DEBUG((DEBUG_ERROR, " Handle Protocolr for gEfiAtaPassThruProtocolGuid Error!, Status = %r\n", Status));
    goto Error1;
  }
  //
  // Do not enable DevSlp if DevSlp and Aggressive DevSlp Management is not supported
  //
  if ((Capability2 & (B_SATA_MEM_AHCI_CAP2_SDS | B_SATA_MEM_AHCI_CAP2_SADM)) == 0) {
    DEBUG ((DEBUG_INFO, "ERROR: DevSlp and Aggressive DevSlp Management is not supported.\n"));
    goto Error1;
  }
  //
  // For any port, Check activitys before opening the feature.
  //
  for (PortIndex = 0; PortIndex < MaxSataPorts; PortIndex++) {
    if (!PortData[PortIndex].State.SetupDevsleepEnable) {
      continue;
    }
    //
    // Do not enable DevSlp if DevSlp is not present
    // Do not enable DevSlp if Hot Plug or Mechanical Presence Switch is supported
    //
    if ((((Mmio32(AhciBar, R_SATA_MEM_AHCI_P0DEVSLP + (0x80 * PortIndex))) & B_SATA_MEM_AHCI_PxDEVSLP_DSP) == 0) ||
        ((Mmio32(AhciBar, R_SATA_MEM_AHCI_P0CMD + (0x80 * PortIndex)) &
         (B_SATA_MEM_AHCI_PxCMD_HPCP | B_SATA_MEM_AHCI_PxCMD_MPSP)) != 0)) {
      DEBUG ((DEBUG_INFO, "ERROR: DevSlp is not present, Hot Plug or Mechanical Presence Switch is supported\n"));
      continue;
    }
    //
    // Detect whether the device supports DevSlp
    //
    GetSATAIdentifyDeviceData (AtaPassThruPtr, &PortIndex, &PortData[PortIndex]);
    if (!PortData[PortIndex].State.DevsleepCap) {
      //
      // DevSlp is not supported by the device
      //
      DEBUG ((DEBUG_INFO, "DevSlp feature is not supported for device at port [%d]!\n", PortIndex));
      continue;
    }
    if (!PortData[PortIndex].State.DevsleepPwr) {
      //
      // DevSlp to Reduced PwrState is not supported
      //
      DEBUG ((DEBUG_INFO, "DevSlp to Reduced PwrState is not supported for device at port [%d]\n", PortIndex));
      continue;
    }
    //
    // DevSlp feature wasn't enabled for device
    //
    if (!PortData[PortIndex].State.DevsleepEn) {
      //
      // Enable Device DevSlp feature via SET FEATURES command.
      //
      Status = EnableSATADevsleep (AtaPassThruPtr, &PortIndex);
      if (EFI_ERROR (Status)) {
        DEBUG((DEBUG_ERROR, "EnableSATADevsleep Error!, Status = %r\n", Status));
        continue;
      }
    }

    Status = GetSATADevsleepTimes (AtaPassThruPtr, &PortIndex, &PortData[PortIndex]);
    if (EFI_ERROR (Status)) {
      DEBUG((DEBUG_ERROR, "GetSATADevsleepTimes Error!, Status = %r\n", Status));
    }
    //
    // Makesure PxCMD.ST and PxDEVSLP.ADSE are cleared to '0' before updating PxDEVSLP.DITO and PxDEVSLP.MDAT value.
    //
    Data32 = Mmio32 (AhciBar, R_SATA_MEM_AHCI_P0CMD + (0x80 * PortIndex));
    Mmio32And (AhciBar, R_SATA_MEM_AHCI_P0CMD + (0x80 * PortIndex), (UINT32)~(B_SATA_MEM_AHCI_PxCMD_ST));
    Mmio32And (AhciBar, R_SATA_MEM_AHCI_P0DEVSLP + (0x80 * PortIndex), (UINT32)~(B_SATA_MEM_AHCI_PxDEVSLP_ADSE));
    Mmio32And (AhciBar, R_SATA_MEM_AHCI_P0DEVSLP + (0x80 * PortIndex), (UINT32)~((0x1f) << 10));
    Mmio32And (AhciBar, R_SATA_MEM_AHCI_P0DEVSLP + (0x80 * PortIndex), (UINT32)~((0xff) << 2));
    //
    // Update SATA PxDEVSLP.DETO and PxDEVSLP.MDAT
    //
    Data32Or = (((PortData[PortIndex].DETO & 0x1FF) << 2) | ((PortData[PortIndex].MDAT & 0x1F) << 10));
    Mmio32AndThenOr (AhciBar, R_SATA_MEM_AHCI_P0DEVSLP + (0x80 * PortIndex), 0xFFFF8003, Data32Or);
    //
    // Restore SATA PxCMD.ST and Update PxDEVSLP.ADSE
    //
    Mmio32Or (AhciBar, R_SATA_MEM_AHCI_P0CMD + (0x80 * PortIndex), (Data32 & B_SATA_MEM_AHCI_PxCMD_ST));
    Mmio32Or (AhciBar, R_SATA_MEM_AHCI_P0DEVSLP + (0x80 * PortIndex), B_SATA_MEM_AHCI_PxDEVSLP_ADSE);

    if (PortData[PortIndex].State.Valid) {
      PchSetup->SataType[PortIndex]  = PortData[PortIndex].State.SolidStateDrive;
      UpdatePchSetup = TRUE;
    }
  }

  DEBUG((DEBUG_INFO | DEBUG_ERROR, "SATADevSleepCallback End\n"));
  if (UpdatePchSetup == TRUE) {
    Status = gRT->SetVariable (
                    PCH_SETUP_VARIABLE_NAME,
                    &gPchSetupVariableGuid,
                    EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                    PchSetupSize,
                    PchSetup
                    );
  }

Error1:
    FreePool (HandleBuffer);
Error2:
    FreePool (PchSetup);
}

EFI_STATUS
EnableSATADevsleep (
  IN  EFI_ATA_PASS_THRU_PROTOCOL         *AtaPassThruPtr,
  IN  UINT16                             *PortNumber
  )
{
  EFI_STATUS                         Status;
  EFI_ATA_STATUS_BLOCK               Asb;
  EFI_ATA_COMMAND_BLOCK              Acb;
  EFI_ATA_PASS_THRU_COMMAND_PACKET   AtaPassThruCmdPacket;

  Status = EFI_SUCCESS;

  ZeroMem (&AtaPassThruCmdPacket, sizeof (EFI_ATA_PASS_THRU_COMMAND_PACKET));
  ZeroMem (&Asb, sizeof (EFI_ATA_STATUS_BLOCK));
  ZeroMem (&Acb, sizeof (EFI_ATA_COMMAND_BLOCK));

  AtaPassThruCmdPacket.Asb = &Asb;
  AtaPassThruCmdPacket.Acb = &Acb;
  AtaPassThruCmdPacket.Acb->AtaCommand       = ATA_COMMAND_SET_FEATURES;
  AtaPassThruCmdPacket.Acb->AtaFeatures      = ATA_COMMAND_SET_FEATURES_ENABLE;
  AtaPassThruCmdPacket.Acb->AtaSectorCount   = ATA_COMMAND_FEATURES_DEVSLEEP;
  AtaPassThruCmdPacket.Timeout               = 100000000;
  AtaPassThruCmdPacket.Protocol              = EFI_ATA_PASS_THRU_PROTOCOL_ATA_NON_DATA;
  AtaPassThruCmdPacket.Length                = EFI_ATA_PASS_THRU_LENGTH_NO_DATA_TRANSFER;
  Status = AtaPassThruPtr->PassThru (
                             AtaPassThruPtr,
                             *PortNumber,
                             0,
                             &AtaPassThruCmdPacket,
                             NULL
                             );
  return Status;
}

EFI_STATUS
GetSATADevsleepTimes (
  IN  EFI_ATA_PASS_THRU_PROTOCOL         *AtaPassThruPtr,
  IN  UINT16                             *PortNumber,
  OUT DEV_SATA_DVSLEEP_DATA              *PortData
  )
{
  EFI_STATUS                         Status;
  EFI_ATA_STATUS_BLOCK               Asb;
  EFI_ATA_COMMAND_BLOCK              Acb;
  EFI_ATA_PASS_THRU_COMMAND_PACKET   AtaPassThruCmdPacket;
  EFI_READ_LOG_EXT_DATA              *SataLog;
  UINT32                             DevSlpData;

  Status = gBS->AllocatePool (
                  EfiBootServicesData,
                  sizeof (EFI_READ_LOG_EXT_DATA),
                  (VOID **)&SataLog
                  );
  ZeroMem (SataLog, sizeof (EFI_READ_LOG_EXT_DATA));
  ZeroMem (&AtaPassThruCmdPacket, sizeof (EFI_ATA_PASS_THRU_COMMAND_PACKET));
  ZeroMem (&Asb, sizeof (EFI_ATA_STATUS_BLOCK));
  ZeroMem (&Acb, sizeof (EFI_ATA_COMMAND_BLOCK));

  AtaPassThruCmdPacket.Asb                   = &Asb;
  AtaPassThruCmdPacket.Acb                   = &Acb;
  AtaPassThruCmdPacket.Acb->AtaCommand       = ATA_COMMAND_READ_LOG_EXT;
  AtaPassThruCmdPacket.Acb->AtaSectorNumber  = 0x30;
  AtaPassThruCmdPacket.Acb->AtaCylinderLow   = 0x8;
  AtaPassThruCmdPacket.Acb->AtaSectorCount   = (UINT8) 0x1;
  AtaPassThruCmdPacket.InTransferLength      = sizeof (EFI_READ_LOG_EXT_DATA);
  AtaPassThruCmdPacket.InDataBuffer          = SataLog;
  AtaPassThruCmdPacket.Protocol              = EFI_ATA_PASS_THRU_PROTOCOL_PIO_DATA_IN;
  AtaPassThruCmdPacket.Length                = EFI_ATA_PASS_THRU_LENGTH_BYTES;
  AtaPassThruCmdPacket.Timeout               = 100000000;
  Status = AtaPassThruPtr->PassThru (
                             AtaPassThruPtr,
                             *PortNumber,
                             0,
                             &AtaPassThruCmdPacket,
                             NULL
                             );
  //
  // Set PxDEVSLP.DETO to 20ms and PxDEVSLP.MDAT to 10ms if the Identify Device Data log (30h, 8)
  // reports 0 or if the command fails
  //
  DevSlpData = SataLog->Buffer[12];
  if (EFI_ERROR (Status) || ((DevSlpData & 0xFF00) == 0)) {
    PortData->DETO = 20;
  } else {
    PortData->DETO = ((DevSlpData & 0xFF00) >> 8);
  }

  if (EFI_ERROR (Status) || ((DevSlpData & 0x1F) == 0)) {
    PortData->MDAT = 10;
  } else {
    PortData->MDAT = DevSlpData & 0x1F;
  }

   FreePool (SataLog);

  return Status;
}


EFI_STATUS
GetSATAIdentifyDeviceData (
  IN  EFI_ATA_PASS_THRU_PROTOCOL         *AtaPassThruPtr,
  IN  UINT16                             *PortNumber,
  DEV_SATA_DVSLEEP_DATA                         *PortData
  )
{
  EFI_STATUS                             Status;
  EFI_ATA_STATUS_BLOCK                   Asb;
  EFI_ATA_COMMAND_BLOCK                  Acb;
  EFI_ATA_PASS_THRU_COMMAND_PACKET       AtaPassThruCmdPacket;
  EFI_ATA8_IDENTIFY_DATA                 *IdentifyTable;

  PortData->State.Valid           = 0;
  PortData->State.SolidStateDrive = 0;
  PortData->State.DevsleepCap     = 0;
  PortData->State.DevsleepEn      = 0;

  Status = gBS->AllocatePool (
                  EfiBootServicesData,
                  0x200,
                  (VOID **)&IdentifyTable
                  );

  ZeroMem (&AtaPassThruCmdPacket, sizeof (EFI_ATA_PASS_THRU_COMMAND_PACKET));
  ZeroMem (&Asb, sizeof (EFI_ATA_STATUS_BLOCK));
  ZeroMem (&Acb, sizeof (EFI_ATA_COMMAND_BLOCK));

  AtaPassThruCmdPacket.Asb                = &Asb;
  AtaPassThruCmdPacket.Acb                = &Acb;
  AtaPassThruCmdPacket.Acb->AtaCommand    = ATA_COMMAND_IDENTIFY_DEVICE;
  AtaPassThruCmdPacket.InDataBuffer       = (VOID*) IdentifyTable;
  AtaPassThruCmdPacket.InTransferLength   = ATA_LOG_SIZE;
  AtaPassThruCmdPacket.Protocol           = EFI_ATA_PASS_THRU_PROTOCOL_PIO_DATA_IN;
  AtaPassThruCmdPacket.Length             = EFI_ATA_PASS_THRU_LENGTH_BYTES;
  AtaPassThruCmdPacket.Timeout            = 100000000;
  Status = AtaPassThruPtr->PassThru (
                             AtaPassThruPtr,
                             *PortNumber,
                             0,
                             &AtaPassThruCmdPacket,
                             NULL
                             );
  if (!EFI_ERROR (Status)) {
    PortData->State.Valid = 1;
    //
    // Detect if the Device is Non-rotating media (e.g., SSD)
    //
    if (IdentifyTable->reserved_210_255[7] == 0x0001) {
      PortData->State.SolidStateDrive = 1;
      //
      // Word 77, bit 7: DevSlp to Reduced PwrState
      //
      PortData->State.DevsleepPwr = ((IdentifyTable->reserved_69_79[8] & 0x80) >> 7);
      //
      // Word 78, bit 8: Support DevSlp
      //
      PortData->State.DevsleepCap = ((IdentifyTable->reserved_69_79[9] & 0x100) >> 8);
      //
      // Word 79, bit 8: Enable DevSlp
      //
      PortData->State.DevsleepEn = ((IdentifyTable->reserved_69_79[10] & 0x100) >> 8);
    }
  }

  FreePool (IdentifyTable);

  return EFI_SUCCESS;
}