/** @file
  TigerLake iPCM test application

  This application tests Integrated Power Consumption Measurement Function 

Copyright (c) 2019, Intel Corporation. All rights reserved.<BR>
This software and associated documentation (if any) is furnished
under a license and may only be used or copied in accordance
with the terms of the license. Except as permitted by such
license, no part of this software or documentation may be
reproduced, stored in a retrieval system, or transmitted in any
form or by any means without the express written consent of
Intel Corporation.

**/          
#include <Uefi.h>
#include <Library/UefiLib.h>
#include <Library/DebugLib.h>
#include <Library/ShellCEntryLib.h>
#include <Protocol/Spi.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/BaseLib.h>
#include <Library/ShellLib.h>
#include <Protocol/SimpleFileSystem.h>
#include <Library/MmPciLib.h>

#include "..\..\Include\E3DongleLib.h"
#include "E3DataTest.h"


//usage
void usage(void)
{
  Print(L"Usage (all values are in hex): \n");
  Print(L"   E3DataTest.efi usbmode\n");
  Print(L"   E3DataTest.efi i2cmode\n");
  Print(L"   E3DataTest.efi writedongle  <i2cbus> <slaveaddr> <offset> <value>\n");
  Print(L"   E3DataTest.efi readdongle  <i2cbus> <slaveaddr> <offset>\n");
  Print(L"   E3DataTest.efi status\n");
  Print(L"   E3DataTest.efi genpdr <filename>\n");
  Print(L"   E3DataTest.efi gentestpdr <filename>\n");
  Print(L"   E3DataTest.efi dumpdongledata <i2cbus> <slaveaddr> <size> <filename>\n");
  Print(L"   E3DataTest.efi writefiletodongle <i2cbus> <slaveaddr> <filename>\n");
  Print(L"   E3DataTest.efi externalpower <value>\n");
  Print(L"   E3DataTest.efi readioexp <offset>\n");
  Print(L"   E3DataTest.efi writeioexp <offset> <value>\n");
  return;
}

/**
  This function write a buffer to a file

  @param[in] Buffer          Pointer to buffer
  @param[in] Size            Number of bytes
   @param[in]FileName     File name

  @retval EFI_SUCCESS        The function completed successfully.
**/
EFI_STATUS
WriteBufferToFile (
  IN VOID       *Buffer,
  IN UINTN      Size,
  IN CHAR16     *Name
  )
{
  EFI_STATUS              Status;
  SHELL_FILE_HANDLE       FileHandle;
  
  Status = ShellFileExists(Name);
  if(!EFI_ERROR(Status)){
    Print(L"File %s Is Existing\n",Name);
    return EFI_INVALID_PARAMETER;
  }

  Status = ShellOpenFileByName(Name, &FileHandle, EFI_FILE_MODE_READ|EFI_FILE_MODE_WRITE|EFI_FILE_MODE_CREATE, 0);
  if(EFI_ERROR(Status)){
    Print(L"Cannot create file %s to write!\n", Name);
    return Status;
  }

  Status = ShellWriteFile(FileHandle, &Size, Buffer);
  if(EFI_ERROR(Status)){
    Print(L"Cannot write the file %s!\n", Name);
    ShellCloseFile(&FileHandle);
    return Status;
  }

  ShellCloseFile(&FileHandle);
  return EFI_SUCCESS;
}

/**
  This function read a file to a buffer

  @param[out] Buffer          Pointer to buffer
  @param[in out] Size         Number of bytes
   @param[in]FileName        File name

  @retval EFI_SUCCESS        The function completed successfully.
**/
EFI_STATUS
ReadFileToBuffer (
  IN VOID       *Buffer,
  IN OUT UINTN  *Size,
  IN CHAR16     *Name
  )
{
  EFI_STATUS              Status;
  SHELL_FILE_HANDLE       FileHandle;
  
  Status = ShellFileExists(Name);
  if(EFI_ERROR(Status)){
    Print(L"File %s Is NOT Existing\n",Name);
    return EFI_INVALID_PARAMETER;
  }

  Status = ShellOpenFileByName(Name, &FileHandle, EFI_FILE_MODE_READ|EFI_FILE_MODE_WRITE, 0);
  if(EFI_ERROR(Status)){
    Print(L"Cannot open file %s to read!\n", Name);
    return Status;
  }

  Status = ShellReadFile(FileHandle, Size, Buffer);
  if(EFI_ERROR(Status)){
    Print(L"Cannot read the file %s!\n", Name);
    ShellCloseFile(&FileHandle);
    return Status;
  }

  ShellCloseFile(&FileHandle);
  return EFI_SUCCESS;
}

/**
  This function read a file to a buffer

  @param[out] Buffer          Pointer to buffer
  @param[in out] Size         Number of bytes
   @param[in]FileName        File name

  @retval EFI_SUCCESS        The function completed successfully.
**/
EFI_STATUS
GetFileSize (
  IN OUT UINTN  *Size,
  IN CHAR16     *Name
  )
{
  EFI_STATUS              Status;
  SHELL_FILE_HANDLE       FileHandle;
  UINT64                  FileSize;
  
  Status = ShellFileExists(Name);
  if(EFI_ERROR(Status)){
    Print(L"File %s Is NOT Existing\n",Name);
    return EFI_INVALID_PARAMETER;
  }

  Status = ShellOpenFileByName(Name, &FileHandle, EFI_FILE_MODE_READ|EFI_FILE_MODE_WRITE, 0);
  if(EFI_ERROR(Status)){
    Print(L"Cannot open file %s!\n", Name);
    return Status;
  }

  Status = ShellGetFileSize(FileHandle, &FileSize);
  if(EFI_ERROR(Status)){
    Print(L"Cannot get the size of the file %s!\n", Name);
    ShellCloseFile(&FileHandle);
    return Status;
  }

  *Size = (UINTN)FileSize;
  ShellCloseFile(&FileHandle);
  return EFI_SUCCESS;
}




/**
  This function calculates an UINT8 checksum.

  @param[in] Buffer          Pointer to buffer to checksum
  @param[in] Size            Number of bytes to checksum

  @retval EFI_SUCCESS        The function completed successfully.
**/
UINT8
Checksum (
  IN VOID       *Buffer,
  IN UINTN      Size
  )
{
  UINT8 Sum;
  UINT8 *Ptr;

  Sum = 0;
  ///
  /// Initialize pointer
  ///
  Ptr = Buffer;

  ///
  /// add all content of buffer
  ///
  while (Size--) {
    Sum = (UINT8) (Sum + (*Ptr++));
  }
  ///
  /// check checksum
  ///
  return ((0xFF-Sum+1) & 0xFF);
}


/**
  This function calculates an UINT8 checksum.

  @param[in] Buffer          Pointer to buffer to checksum
  @param[in] Size            Number of bytes to checksum

  @retval EFI_SUCCESS        The function completed successfully.
**/
EFI_STATUS
CheckChecksum (
  IN VOID       *Buffer,
  IN UINTN      Size
  )
{
  ///
  /// check checksum
  ///
  if (Checksum(Buffer, Size) != 0)
    return EFI_CRC_ERROR;

  return EFI_SUCCESS;
}

EFI_STATUS
DumpPDRData()
{
  EFI_STATUS                           Status;
  PCH_SPI_PROTOCOL                     *mSpiProtocol;
  E3_CALIBRATION_DATA_HEADER           E3CalibrationDataHeader;
  UINT8                                *Ptr = NULL; 
  UINT8                                *PdrBuffer = NULL;
  UINT16                               Length = 0;
  UINT16                               i;
  
  // Read the E3 calibration header from PDR (Platform Data Region)
  Status = gBS->LocateProtocol (
                  &gPchSpiProtocolGuid,
                  NULL,
                  (VOID **)&mSpiProtocol
                  );
  if (EFI_ERROR (Status)) {
    Print(L"Cannot locate PchSpiProtocol!!!\n");
    return Status;
  }
  Status = mSpiProtocol->FlashRead(
                           mSpiProtocol,
                           FlashRegionPlatformData,
                           (UINT32) E3_CALIBRATION_DATA_OFFSET_IN_PDR,
                           (UINT32) sizeof (E3_CALIBRATION_DATA_HEADER),
                           (UINT8*) &E3CalibrationDataHeader
                           );
  
  if (EFI_ERROR (Status)) {
    Print(L"Cannot read the header from PDR!!!\n");
    return Status;
  }

  Ptr = (UINT8*)&E3CalibrationDataHeader;
  Length = sizeof(E3CalibrationDataHeader);

  // Check if the signature in the header is valid
  if ((E3CalibrationDataHeader.Signature == E3_CALIBRATION_DATA_HEADER_SIGNATURE) &&
    (E3CalibrationDataHeader.Length > 0) && (E3CalibrationDataHeader.Length < 0xFF)) {   
    // Read the E3 calibration data from PDR (Platform Data Region)
    Status= gBS->AllocatePool(EfiBootServicesData, E3CalibrationDataHeader.Length, &PdrBuffer);
    if (EFI_ERROR (Status)) {
      Print(L"Cannot AllocatePool!!!\n");
      return Status;
    }

    CopyMem((VOID *)PdrBuffer, (VOID *)&E3CalibrationDataHeader, sizeof(E3CalibrationDataHeader));

    Status = mSpiProtocol->FlashRead(
                             mSpiProtocol,
                             FlashRegionPlatformData,
                             (UINT32) E3CalibrationDataHeader.DataOffset,
                             (UINT32) E3CalibrationDataHeader.DataLength,
                             (UINT8 *)(PdrBuffer+sizeof(E3CalibrationDataHeader))
                             );
    
    if (EFI_ERROR (Status)) {
      Print(L"Cannot read the data in PDR!!!\n");
      return Status;
    }

    Ptr = PdrBuffer;
    Length = E3CalibrationDataHeader.Length;

    // Check the checksum
    Status = CheckChecksum(PdrBuffer, E3CalibrationDataHeader.Length);
    if (EFI_ERROR (Status)) {
      Print(L"The data is not existing in PDR - invalid checksum!!!\n");
    }
  } else {
    Print(L"The E3 data is not existing - invalid signature in PDR!!!\n");
  }
  
  Print(L"The data in PDR (in Hex):");
  if ((Ptr != NULL) && (Length != 0)) {
    for (i=0; i< Length; i++) {
      if ((i%16) == 0) Print(L"\n");
      Print(L"%02X ", Ptr[i]);
    }
  }

  if (PdrBuffer != NULL) {
    gBS->FreePool(PdrBuffer);
  }

  return EFI_SUCCESS;
  
}


/**
  UEFI application entry point which has an interface similar to a
  standard C main function.

  The ShellCEntryLib library instance wrappers the actual UEFI application
  entry point and calls this ShellAppMain function.

  @param  ImageHandle  The image handle of the UEFI Application.
  @param  SystemTable  A pointer to the EFI System Table.

  @retval  0               The application exited normally.
  @retval  Other           An error occurred.

**/
INTN
EFIAPI
ShellAppMain (
  IN UINTN Argc,
  IN CHAR16 **Argv
  )
{
  EFI_STATUS                    Status;
  UINTN                         i2cAddr=0;
  UINTN                         BusNo=0;
  UINTN                         offset=0;
  UINTN                         value=0;  
  UINTN                         size=0;
  CHAR8                         AsciiStr[32]={0};
  CHAR8                         *Buffer=NULL;
  E3_CALIBRATION_DATA_HEADER    *pHeader=NULL;
  
  if (Argc < 2) {
    usage();
    return EFI_INVALID_PARAMETER;
  }
  
  ZeroMem(AsciiStr,sizeof(AsciiStr));
  UnicodeStrToAsciiStrS(Argv[1], AsciiStr, sizeof(AsciiStr));

  if(0 == AsciiStrCmp(AsciiStr, "usbmode")) {
    if(Argc != 2) {
      Print(L"The number of arguments are invalid!\n"); 
      return EFI_INVALID_PARAMETER;
    }
    iRPMMuxSwitch(FALSE);
    Print(L"Switch the E3 mux to USB mode - Done.\n"); 
  } else if(0 == AsciiStrCmp(AsciiStr, "i2cmode"))   {
    iRPMMuxSwitch(TRUE);
    Print(L"Switch the E3 mux to I2C mode - Done.\n"); 
  } else if(0 == AsciiStrCmp(AsciiStr, "writedongle"))   {
    BusNo  = StrHexToUintn (Argv[2]);
    i2cAddr = StrHexToUintn (Argv[3]);
    offset  = StrHexToUintn (Argv[4]);
    value = StrHexToUintn (Argv[5]);
    Status = WriteE3DongleByte(offset, value);
    if (EFI_ERROR(Status)) {
      Print(L"Failed to write the EEPROM on dongle - bus: 0x%X, slaveaddr:0x%X, offset:0x%X, value:0x%X!\n", BusNo, i2cAddr, offset, value); 
      return Status;
    } else {
      Print(L"Sucessful to write the EEPROM on dongle - bus: 0x%X, slaveaddr:0x%X, offset:0x%X, value:0x%X.\n", BusNo, i2cAddr, offset, value);
    }
  } else if(0 == AsciiStrCmp(AsciiStr, "readdongle"))   {
    BusNo  = StrHexToUintn (Argv[2]);
    i2cAddr = StrHexToUintn (Argv[3]);
    offset  = StrHexToUintn (Argv[4]);    
    Status = ReadE3DongleByte(offset, &value);
    if (EFI_ERROR(Status)) {
      Print(L"Failed to read the EEPROM on dongle - bus: 0x%X, slaveaddr:0x%X, offset:0x%X!\n", BusNo, i2cAddr, offset); 
      return Status;
    } else {
      Print(L"Sucessful to write the EEPROM on dongle - bus: 0x%X, slaveaddr:0x%X, offset:0x%X, value:0x%X.\n", BusNo, i2cAddr, offset, value);
    }
  } else if(0 == AsciiStrCmp(AsciiStr, "status"))   {
    //Check the Dongle presence
    if (IsE3DongleAttached()) {
      Print(L"The E3 Dongle is attached and detected!\n"); 
    } else {
      Print(L"The E3 Dongle is NOT detected!\n"); 
    }

    //Dump the data from PDR
    DumpPDRData();
  } else if(0 == AsciiStrCmp(AsciiStr, "dumpdongledata"))   {
    BusNo  = StrHexToUintn (Argv[2]);
    i2cAddr = StrHexToUintn (Argv[3]);
    size  = StrHexToUintn (Argv[4]); 
    
    Buffer = AllocateZeroPool(size);
    Status = ReadE3DongleBuffer(0, size, Buffer);
    if (EFI_ERROR(Status)) {
      Print(L"Failed to dump the EEPROM on dongle - bus: 0x%X, slaveaddr:0x%X, size: 0x%X!\n", BusNo, i2cAddr, size); 
      FreePool(Buffer);
      return Status;
    } else {
      Print(L"Sucessful to dump the EEPROM on dongle - bus: 0x%X, slaveaddr:0x%X, size: 0x%X.\n", BusNo, i2cAddr, size);
    }
    WriteBufferToFile(Buffer,size,Argv[5]);
    FreePool(Buffer);
    return Status;
  } else if(0 == AsciiStrCmp(AsciiStr, "writefiletodongle"))   {
    BusNo  = StrHexToUintn (Argv[2]);
    i2cAddr = StrHexToUintn (Argv[3]);
    GetFileSize (&size, Argv[4]); 
    
    Buffer = AllocateZeroPool(size);
    Status = ReadFileToBuffer(Buffer,&size, Argv[4]);
    if (EFI_ERROR(Status)) {
      FreePool(Buffer);
      return Status;
    }
    Status = WriteE3DongleBuffer(0, size, Buffer);
    if (EFI_ERROR(Status)) {
      Print(L"Failed to write the file %s to the EEPROM on dongle - bus: 0x%X, slaveaddr:0x%X, size: 0x%X!\n",Argv[4], BusNo, i2cAddr, size); 
      FreePool(Buffer);
      return Status;
    } else {
      Print(L"Successful to write the file %s to the EEPROM on dongle - bus: 0x%X, slaveaddr:0x%X, size: 0x%X!\n",Argv[4], BusNo, i2cAddr, size); 
    }

    FreePool(Buffer);
    return Status;
  } else if(0 == AsciiStrCmp(AsciiStr, "genpdr"))   {
    Status= gBS->AllocatePool(EfiBootServicesData, PDR_SIZE, &Buffer);
    if (EFI_ERROR (Status)) {
      Print(L"Cannot AllocatePool!!!\n");
      return Status;
    }
    SetMem(Buffer,PDR_SIZE,0xFF);
    Status= WriteBufferToFile(Buffer,PDR_SIZE,Argv[2]);
    if (Buffer != NULL) {
      gBS->FreePool(Buffer);
    }
    return Status;
  } else if(0 == AsciiStrCmp(AsciiStr, "gentestpdr"))   {
    Status= gBS->AllocatePool(EfiBootServicesData, PDR_SIZE, &Buffer);
    if (EFI_ERROR (Status)) {
      Print(L"Cannot AllocatePool!!!\n");
      return Status;
    }
    pHeader = (E3_CALIBRATION_DATA_HEADER*)Buffer;
    pHeader->Signature = E3_CALIBRATION_DATA_HEADER_SIGNATURE;
    pHeader->Revision = 0x01;
    Status= WriteBufferToFile(Buffer,PDR_SIZE,Argv[2]);
    if (Buffer != NULL) {
      gBS->FreePool(Buffer);
    }
    return Status;
  } else if(0 == AsciiStrCmp(AsciiStr, "externalpower"))   {
    value  = StrHexToUintn (Argv[2]);
    
    Status = ExternalPowerControlForE3Chip((value==0)?FALSE: TRUE);
    if (EFI_ERROR(Status)) {
      Print(L"Failed to control the external power from dongle to E3 chip: value=%d!\n", value); 
    } else {
      Print(L"Successful to control the external power from dongle to E3 chip: value=%d.\n", value); 
    }
    return Status;
  } else if(0 == AsciiStrCmp(AsciiStr, "readioexp"))   {
    offset  = StrHexToUintn (Argv[2]);
    value = 0;
    Status = ReadE3DongleIOExp((UINT8)offset, (UINT8 *)&value);
    if (EFI_ERROR(Status)) {
      Print(L"Failed to read IOExpaner on the dongle: offset=0x%x!\n", offset); 
    } else {
      Print(L"Successful to read IOExpaner on the dongle: offset=0x%x, value=0x%x!\n", offset, value&0xFF); 
    }
    return Status;
  } else if(0 == AsciiStrCmp(AsciiStr, "writeioexp"))   {
    offset  = StrHexToUintn (Argv[2]);
    value  = StrHexToUintn (Argv[3]);
    Status = WriteE3DongleIOExp((UINT8)offset, (UINT8)value);
    if (EFI_ERROR(Status)) {
      Print(L"Failed to write IOExpaner on the dongle: offset=0x%x, value=0x%x!\n", offset, value); 
    } else {
      Print(L"Successful to write IOExpaner on the dongle: offset=0x%x, value=0x%x!\n", offset, value); 
    }
    return Status;
  }

  
  return EFI_SUCCESS;
}