/**@file
  Power Loss Notification

@copyright
  INTEL CONFIDENTIAL
  Copyright 2020 - 2021 Intel Corporation.

  The source code contained or described herein and all documents related to the
  source code ("Material") are owned by Intel Corporation or its suppliers or
  licensors. Title to the Material remains with Intel Corporation or its suppliers
  and licensors. The Material may contain trade secrets and proprietary and
  confidential information of Intel Corporation and its suppliers and licensors,
  and is protected by worldwide copyright and trade secret laws and treaty
  provisions. No part of the Material may be used, copied, reproduced, modified,
  published, uploaded, posted, transmitted, distributed, or disclosed in any way
  without Intel's prior express written permission.

  No license under any patent, copyright, trade secret or other intellectual
  property right is granted to or conferred upon you by disclosure or delivery
  of the Materials, either expressly, by implication, inducement, estoppel or
  otherwise. Any license under such intellectual property rights must be
  express and approved by Intel in writing.

  Unless otherwise agreed by Intel in writing, you may not remove or alter
  this notice or any other notice embedded in Materials by Intel or
  Intel's suppliers or licensors in any way.

  This file contains a 'Sample Driver' and is licensed as such under the terms
  of your license agreement with Intel or your vendor. This file may be modified
  by the user, subject to the additional terms of the license agreement.

@par Specification Reference:
**/

#include "PowerLossNotify.h"

GLOBAL_REMOVE_IF_UNREFERENCED  VOID *mNvmePassthruRegistration  = NULL;
GLOBAL_REMOVE_IF_UNREFERENCED  BOOLEAN mIsPlnConfigured = FALSE;

/**
  Set feature command.

  @param[in]   NvmeDevice       The pointer to the NVME_PASS_THRU_DEVICE data structure.
  @param[in]   NameSpaceId      Namespace ID
  @param[in]   EnablePLN        Indicate if PLN is enabled or not.

  @return EFI_SUCCESS           Successfully send command to device.
  @return EFI_DEVICE_ERROR      Fail to send command to device.

**/
EFI_STATUS
NvmeSetFeature (
  IN  EFI_NVM_EXPRESS_PASS_THRU_PROTOCOL    *NvmeDevice,
  IN  UINT32                                NameSpaceId,
  IN  UINT8                                 EnablePln
  )
{
  EFI_NVM_EXPRESS_PASS_THRU_COMMAND_PACKET CommandPacket;
  EFI_NVM_EXPRESS_COMMAND                  Command;
  EFI_NVM_EXPRESS_COMPLETION               Completion;
  EFI_STATUS                               Status;
  SAVE_FEATURE_DW3                         *ResultDw3;
  SAVE_FEATURE_DW0                         *ResultDw0;
  NVME_ADMIN_SET_FEATURES                  SetFeatureCdw10;

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

  ZeroMem (&CommandPacket, sizeof (EFI_NVM_EXPRESS_PASS_THRU_COMMAND_PACKET));
  ZeroMem (&Command, sizeof (EFI_NVM_EXPRESS_COMMAND));
  ZeroMem (&Completion, sizeof (EFI_NVM_EXPRESS_COMPLETION));
  ZeroMem (&SetFeatureCdw10, sizeof (UINT32));

  CommandPacket.NvmeCmd        = &Command;
  CommandPacket.NvmeCompletion = &Completion;

  SetFeatureCdw10.Fid = PLN_FEATURE_ID;
  SetFeatureCdw10.Sv = BIT0;

  Command.Cdw0.Opcode          = NVME_ADMIN_SET_FEATURES_CMD;
  CommandPacket.NvmeCmd->Cdw10 = *((UINT32*) &SetFeatureCdw10);
  CommandPacket.NvmeCmd->Cdw11 = (UINT32)(EnablePln & BIT0);
  CommandPacket.NvmeCmd->Flags = CDW10_VALID | CDW11_VALID;
  CommandPacket.NvmeCmd->Nsid  = NVME_CONTROLLER_ID;
  CommandPacket.CommandTimeout = EFI_TIMER_PERIOD_SECONDS (5);
  CommandPacket.QueueType      = NVME_ADMIN_QUEUE;

  Status = NvmeDevice->PassThru (
                         NvmeDevice,
                         NameSpaceId,
                         &CommandPacket,
                         NULL
                         );

  if (!EFI_ERROR (Status)) {
    ResultDw3 = (SAVE_FEATURE_DW3*) &Completion.DW3;
    ResultDw0 = (SAVE_FEATURE_DW0*) &Completion.DW0;
    if ((ResultDw3->StatusCode == 0) && (ResultDw3->StatusCodeType == 0) && ((UINT8)ResultDw0->PlnEnable == EnablePln)) {
      return EFI_SUCCESS;
    }
    return EFI_UNSUPPORTED;
  }

  return Status;
}

/**
  Get feature command.

  @param[in]   NvmeDevice       The pointer to the NVME_PASS_THRU_DEVICE data structure.
  @param[in]   Sel              Select field for get feature command.
  @param[in]   NameSpaceId      Namespace ID
  @param[out]  Result           capability status

  @return EFI_SUCCESS           Successfully send command to device.
  @return EFI_DEVICE_ERROR      Fail to send command to device.

**/
EFI_STATUS
NvmeGetFeature (
  IN  EFI_NVM_EXPRESS_PASS_THRU_PROTOCOL    *NvmeDevice,
  IN  UINT8                                 Sel,
  IN  UINT32                                NameSpaceId,
  OUT UINT32                                *Result
  )
{
  EFI_NVM_EXPRESS_PASS_THRU_COMMAND_PACKET CommandPacket;
  EFI_NVM_EXPRESS_COMMAND                  Command;
  EFI_NVM_EXPRESS_COMPLETION               Completion;
  EFI_STATUS                               Status;
  NVME_ADMIN_GET_FEATURES                  GetFeatureCdw10;

  if (Result == NULL || NvmeDevice == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  ZeroMem (&CommandPacket, sizeof (EFI_NVM_EXPRESS_PASS_THRU_COMMAND_PACKET));
  ZeroMem (&Command, sizeof (EFI_NVM_EXPRESS_COMMAND));
  ZeroMem (&Completion, sizeof (EFI_NVM_EXPRESS_COMPLETION));
  ZeroMem (&GetFeatureCdw10, sizeof (UINT32));

  CommandPacket.NvmeCmd        = &Command;
  CommandPacket.NvmeCompletion = &Completion;

  GetFeatureCdw10.Fid = PLN_FEATURE_ID;
  GetFeatureCdw10.Sel = Sel;

  Command.Cdw0.Opcode          = NVME_ADMIN_GET_FEATURES_CMD;
  CommandPacket.NvmeCmd->Cdw10 = *((UINT32*) &GetFeatureCdw10);
  CommandPacket.NvmeCmd->Flags = CDW10_VALID;
  CommandPacket.NvmeCmd->Nsid  = NVME_CONTROLLER_ID;
  CommandPacket.CommandTimeout = EFI_TIMER_PERIOD_SECONDS (5);
  CommandPacket.QueueType      = NVME_ADMIN_QUEUE;

  Status = NvmeDevice->PassThru (
                         NvmeDevice,
                         NameSpaceId,
                         &CommandPacket,
                         NULL
                         );

  if (!EFI_ERROR (Status)) {
    CopyMem (Result, &Completion.DW0, sizeof (UINT32));
  }

  return Status;
}

/**
  Enable or disable PLN with EC command on CS or non-CS mode.

  @param[in]  EnablePln        Enable / Disable PLN feature.
  @param[in]  Cs               Connected Standby, 1: enabled, 0: disabled.

**/
VOID
EcControlPln (
  IN  UINT8                     EnablePln,
  IN  UINT8                     Cs
  )
{
  SendEcCommand (EC_C_PLN_ENABLE);

  if (EnablePln == 0) {
    SendEcData(EC_PLN_DISABLED);  // feature is not enabled
  } else if (Cs == 0) {
    SendEcData(EC_PLN_PB_OVR_NO_CS);  // 4 sec
  } else {
    SendEcData(EC_PLN_PB_OVR_CS);  // 10 sec
  }

}

/**
  Enable or disable PLN with GPIO on CS or non-CS mode.

  @param[in]  EnablePln        Enable / Disable PLN feature.
  @param[in]  Cs               Connected Standby, 1: enabled, 0: disabled.

**/
VOID
GpioControlPln (
  IN  UINT8                     EnablePln,
  IN  UINT8                     Cs
  )
{
  UINT32                            PlnDelaySelGpio;
  UINT32                            PlnDelayEnGpio;

  PlnDelayEnGpio = PcdGet32 (PcdPlnDelayEnableGpio);
  PlnDelaySelGpio = PcdGet32 (PcdPlnDelaySelectionGpio);

  //
  // There is no discrete logic to control PLN.
  //
  if (PlnDelayEnGpio == 0 || PlnDelaySelGpio == 0) {
    return;
  }

  if (EnablePln == 0) {
    GpioSetOutputValue (PlnDelayEnGpio, GPIO_PLN_DISABLED);
  } else {
    GpioSetOutputValue (PlnDelayEnGpio, GPIO_PLN_ENABLED);
    if (Cs == 0) {
      GpioSetOutputValue (PlnDelaySelGpio, GPIO_PLN_PB_OVR_NO_CS);
    } else {
      GpioSetOutputValue (PlnDelaySelGpio, GPIO_PLN_PB_OVR_CS);
    }
  }
}

/**
  Enable or disable PLN on device if device support PLN

  @param[in]  NvmePassthru        The pointer to the NVME_PASS_THRU_DEVICE data structure.
  @param[in]  NameSpaceId         Name space ID.
  @param[in]  EnablePln           Enable / Disable PLN feature.

**/
EFI_STATUS
ConfigureDevicePln (
  IN  EFI_NVM_EXPRESS_PASS_THRU_PROTOCOL    *NvmePassthru,
  IN  UINT32                                NameSpaceId,
  IN  UINT8                                 EnablePln
  )
{
  EFI_STATUS                         Status;
  UINT8                              PlnVal;
  GET_FEATURE_ATTR_CUR               CurrentAttr;

  //
  // Test if PLN is supported
  //
  Status = NvmeGetFeature (NvmePassthru, SEL_CURRENT, NameSpaceId, (UINT32*) &CurrentAttr);

  if (!EFI_ERROR(Status)) {
    PlnVal = (UINT8) CurrentAttr.PlnEnable;
    DEBUG ((DEBUG_INFO, "PLN - Device support PLN, namespaceid= %d, status: %d -> %d.\n", NameSpaceId, PlnVal, EnablePln));

    if (EnablePln != PlnVal) {
      Status = NvmeSetFeature (NvmePassthru, NameSpaceId, EnablePln);
      if (EFI_ERROR (Status)) {
        DEBUG ((DEBUG_ERROR, "  Configured PLN to value: %d was failed. Status: %r \n", Status));
      }
    }
  }

  return Status;
}

/**
  Nvme Pass Thru Protocol notification event handler.

  @param[in] Event    Event whose notification function is being invoked.
  @param[in] Context  Pointer to the notification function's context.

**/
VOID
EFIAPI
NvmePassthruNotificationEvent (
  IN  EFI_EVENT       Event,
  IN  VOID            *Context
  )
{
  EFI_STATUS                         Status;
  UINTN                              BufferSize;
  EFI_HANDLE                         Handle;
  EFI_NVM_EXPRESS_PASS_THRU_PROTOCOL *NvmePassthru;
  UINTN                              VarDataSize;
  SETUP_DATA                         SystemConfiguration;
  UINT32                             NamespaceId;

  BufferSize = sizeof (EFI_HANDLE);
  Status = gBS->LocateHandle (
                  ByRegisterNotify,
                  NULL,
                  mNvmePassthruRegistration,
                  &BufferSize,
                  &Handle
                  );
  if (EFI_ERROR (Status)) {
    return ;
  }

  //
  // Get NvmePassThruProtocol interface
  //
  Status = gBS->HandleProtocol (
                  Handle,
                  &gEfiNvmExpressPassThruProtocolGuid,
                  (VOID **) &NvmePassthru
                  );
  if (EFI_ERROR (Status)) {
    return;
  }

  //
  // Check PLN SETUP option
  //
  VarDataSize = sizeof (SETUP_DATA);
  Status = gRT->GetVariable (
                  L"Setup",
                  &gSetupVariableGuid,
                  NULL,
                  &VarDataSize,
                  &SystemConfiguration
                  );
  if (EFI_ERROR (Status)) {
    return;
  }

  //
  // Configured device with desired PLN setting if device support PLN, follow the below flow
  // 1. send command to controller, if it is failed then
  // 2. send command to the rest name space
  //
  Status = ConfigureDevicePln (NvmePassthru, NVME_CONTROLLER_ID, SystemConfiguration.PlnEnable);
  if (EFI_ERROR (Status)) {
    NamespaceId = NVME_ALL_NAMESPACES;
    while (TRUE) {
      Status = NvmePassthru->GetNextNamespace (
                               NvmePassthru,
                               (UINT32 *)&NamespaceId
                               );
      if (EFI_ERROR (Status)) {
        break;
      }

      Status = ConfigureDevicePln (NvmePassthru, NamespaceId, SystemConfiguration.PlnEnable);
    }
  }

  if (!mIsPlnConfigured) {

    DEBUG ((DEBUG_INFO, "PLN - Configure EC and GPIO.\n"));

    //
    // 1. Send command to EC to notify the power button override time if go with EC path
    // 2. Set GPIO
    //
    EcControlPln (SystemConfiguration.PlnEnable, SystemConfiguration.LowPowerS0Idle);
    GpioControlPln (SystemConfiguration.PlnEnable, SystemConfiguration.LowPowerS0Idle);
    mIsPlnConfigured = TRUE;
  }
  return;
}

/**
  Driver entry point to register Nvme PassThru Protocol callback.

  @param[in] ImageHandle      The firmware allocated handle for the EFI image.
  @param[in] SystemTable      A pointer to the EFI System Table.

  @retval EFI_SUCCESS         Notify event has been created
**/
EFI_STATUS
EFIAPI
PowerLostNotifyEntryPoint (
  IN EFI_HANDLE                       ImageHandle,
  IN EFI_SYSTEM_TABLE                 *SystemTable
  )
{
  EFI_STATUS                         Status;
  UINTN                              VarDataSize;
  SETUP_DATA                         SystemConfiguration;

  DEBUG ((DEBUG_INFO, "PLN - PowerLostNotifyEntryPoint\n"));

  //
  // Check PLN SETUP option
  //
  VarDataSize = sizeof (SETUP_DATA);
  Status = gRT->GetVariable (
                  L"Setup",
                  &gSetupVariableGuid,
                  NULL,
                  &VarDataSize,
                  &SystemConfiguration
                  );
  if (!EFI_ERROR (Status) && SystemConfiguration.PlnEnable != CTL_PLN_DEFAULT) {
    //
    // Register NvmePassthruNotificationEvent() notify function.
    //
    EfiCreateProtocolNotifyEvent (
      &gEfiNvmExpressPassThruProtocolGuid,
      TPL_CALLBACK,
      NvmePassthruNotificationEvent,
      NULL,
      &mNvmePassthruRegistration
      );
  }

  return EFI_SUCCESS;
}