/*******************************************************************************
Copyright (c) 2022, Rockchip Corporation. All rights reserved.

SPDX-License-Identifier: BSD-2-Clause-Patent

*******************************************************************************/
#include <Uefi.h>

#include <Library/BaseLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/ShellCommandLib.h>
#include <Library/ShellLib.h>
#include <Library/UefiLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/PrintLib.h>
#include <Library/ShellCEntryLib.h>
#include <Library/HiiLib.h>
#include <Library/FileHandleLib.h>
#include <Pi/PiFirmwareVolume.h>
#include <Protocol/BlockIo.h>
#include <Protocol/FirmwareVolumeBlock.h>
#include <Protocol/NorFlashProtocol.h>

UNI_NOR_FLASH_PROTOCOL *SpiFlashProtocol;

CONST CHAR16 gShellSpiFlashFileName[] = L"ShellCommand";
EFI_HANDLE gShellSfHiiHandle = NULL;

STATIC CONST SHELL_PARAM_ITEM ParamList[] = {
  {L"read", TypeFlag},
  {L"readfile", TypeFlag},
  {L"write", TypeFlag},
  {L"writefile", TypeFlag},
  {L"erase", TypeFlag},
  {L"update", TypeFlag},
  {L"updatefile", TypeFlag},
  {L"help", TypeFlag},
  {NULL , TypeMax}
  };

typedef enum {
  READ        = 2,
  READ_FILE   = 4,
  WRITE       = 8,
  WRITE_FILE  = 16,
  ERASE       = 32,
  UPDATE      = 64,
  UPDATE_FILE = 128,
} Flags;

/**
  Return the file name of the help text file if not using HII.

  @return The string pointer to the file name.
**/
CONST CHAR16*
EFIAPI
ShellCommandGetManFileNameSpiFlash (
  VOID
  )
{

  return gShellSpiFlashFileName;
}

VOID
SfUsage (
  VOID
  )
{
  Print (L"\nBasic SPI command\n"
         "sf [read | readfile | write | writefile | erase |"
         "update | updatefile]"
         "[<Address> | <FilePath>] <Offset> <Length>\n\n"
         "Address  - Address in RAM to store/load data\n"
         "FilePath - Path to file to read/write data from/to\n"
         "Offset   - Offset from beginning of SPI flash to store/load data\n"
         "Length   - Number of bytes to send\n"
         "Examples:\n"
         "Check if there is response from SPI flash\n"
         "Read 32 bytes from 0xe00000 of SPI flash into RAM at address 0x100000\n"
         "  sf read 0x100000 0xe00000 32\n"
         "Read 0x20 bytes from 0x200000 of SPI flash into RAM at address 0x300000\n"
         "  sf read 0x300000 0x200000 0x20\n"
         "Erase 0x10000 bytes from offset 0x100000 of SPI flash\n"
         "  sf erase 0x100000 0x100000\n"
         "Write 16 bytes from 0x200000 at RAM into SPI flash at address 0x4000000\n"
         "  sf write 0x200000 0x4000000 16\n"
         "Update 100 bytes from 0x100000 at RAM in SPI flash at address 0xe00000\n"
         "  sf update 0x100000 0xe00000 100\n"
         "Read 0x3000 bytes from 0x0 of SPI flash into file fs2:file.bin\n"
         "  sf readfile fs2:file.bin 0x0 0x3000 \n"
         "Update data in SPI flash at 0x3000000 from file Linux.efi\n"
         "  sf updatefile Linux.efi 0x3000000\n"
  );
}

STATIC
EFI_STATUS
OpenAndPrepareFile (
  IN CHAR16 *FilePath,
  SHELL_FILE_HANDLE *FileHandle
  )
{
  EFI_STATUS Status;
  UINT64 OpenMode;

  OpenMode = EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE | EFI_FILE_MODE_CREATE;

  Status = ShellOpenFileByName (FilePath, FileHandle, OpenMode, 0);
  if (EFI_ERROR (Status)) {
    Print (L"sf: Cannot open file\n");
    return Status;
  }

  Status = FileHandleSetPosition(*FileHandle, 0);

  if (EFI_ERROR(Status)) {
    Print (L"sf: Cannot set file position to first byte\n");
    ShellCloseFile (FileHandle);
    return Status;
  }

  return EFI_SUCCESS;
}

SHELL_STATUS
EFIAPI
ShellCommandRunSpiFlash (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
EFI_STATUS              Status;
  LIST_ENTRY            *CheckPackage;
  EFI_PHYSICAL_ADDRESS  Address = 0, Offset = 0;
  SHELL_FILE_HANDLE     FileHandle = NULL;
  UINTN                 ByteCount, I;
  UINT64                FileSize;
  UINT8                 *Buffer = NULL, *FileBuffer = NULL;
  CHAR16                *ProblemParam, *FilePath;
  CONST CHAR16          *AddressStr = NULL, *OffsetStr = NULL;
  CONST CHAR16          *LengthStr = NULL, *FileStr = NULL;
  BOOLEAN               AddrFlag = FALSE, LengthFlag = TRUE, FileFlag = FALSE;
  UINT8                 Flag = 0, CheckFlag = 0;

  Status = gBS->LocateProtocol (
    &gUniNorFlashProtocolGuid,
    NULL,
    (VOID **)&SpiFlashProtocol
  );
  if (EFI_ERROR(Status)) {
    Print (L"sf: Cannot locate SpiFlash protocol\n");
    return SHELL_ABORTED;
  }

  // Parse Shell command line
  Status = ShellInitialize ();
  if (EFI_ERROR (Status)) {
    Print (L"sf: Cannot initialize Shell\n");
    ASSERT_EFI_ERROR (Status);
    return SHELL_ABORTED;
  }

  Status = ShellCommandLineParse (ParamList, &CheckPackage, &ProblemParam, TRUE);
  if (EFI_ERROR (Status)) {
    Print (L"sf: Error while parsing command line\n");
    return SHELL_ABORTED;
  }

  if (ShellCommandLineGetFlag (CheckPackage, L"help")) {
    SfUsage();
    return EFI_SUCCESS;
  }

  // Check flags provided by user
  Flag |= (ShellCommandLineGetFlag (CheckPackage, L"read") << 1);
  Flag |= (ShellCommandLineGetFlag (CheckPackage, L"readfile") << 2);
  Flag |= (ShellCommandLineGetFlag (CheckPackage, L"write") << 3);
  Flag |= (ShellCommandLineGetFlag (CheckPackage, L"writefile") << 4);
  Flag |= (ShellCommandLineGetFlag (CheckPackage, L"erase") << 5);
  Flag |= (ShellCommandLineGetFlag (CheckPackage, L"update") << 6);
  Flag |= (ShellCommandLineGetFlag (CheckPackage, L"updatefile") << 7);

  CheckFlag = Flag;
  for (I = 0; CheckFlag; CheckFlag >>= 1) {
    I += CheckFlag & 1;
    if (I > 1) {
      Print (L"sf: Too many flags\n");
      SfUsage();
      return SHELL_ABORTED;
    }
  }

  switch (Flag) {
  case READ:
  case WRITE:
  case UPDATE:
    AddressStr = ShellCommandLineGetRawValue (CheckPackage, 1);
    OffsetStr = ShellCommandLineGetRawValue (CheckPackage, 2);
    LengthStr = ShellCommandLineGetRawValue (CheckPackage, 3);
    AddrFlag = TRUE;
    break;
  case ERASE:
    OffsetStr = ShellCommandLineGetRawValue (CheckPackage, 1);
    LengthStr = ShellCommandLineGetRawValue (CheckPackage, 2);
    break;
  case READ_FILE:
    FileStr = ShellCommandLineGetRawValue (CheckPackage, 1);
    OffsetStr = ShellCommandLineGetRawValue (CheckPackage, 2);
    LengthStr = ShellCommandLineGetRawValue (CheckPackage, 3);
    FileFlag = TRUE;
    break;
  case WRITE_FILE:
  case UPDATE_FILE:
    FileStr = ShellCommandLineGetRawValue (CheckPackage, 1);
    OffsetStr = ShellCommandLineGetRawValue (CheckPackage, 2);
    LengthFlag = FALSE;
    FileFlag = TRUE;
    break;
  }

  // Read address parameter
  if ((AddressStr == NULL) & AddrFlag) {
    Print (L"sf: No address parameter!\n");
    return SHELL_ABORTED;
  } else if (AddrFlag) {
    Address = ShellHexStrToUintn (AddressStr);
    if (Address == (UINTN)(-1)) {
      Print (L"sf: Wrong address parameter\n");
      return SHELL_ABORTED;
    }
  }

  // Read offset parameter
  if (OffsetStr == NULL) {
    Print (L"sf: No offset Parameter!\n");
    return SHELL_ABORTED;
  } else {
    Offset = ShellHexStrToUintn (OffsetStr);
    if (Offset < 0) {
      Print (L"sf: Wrong offset parameter: %s", OffsetStr);
      return SHELL_ABORTED;
    }
  }

  // Read length parameter
  if ((LengthStr == NULL) & LengthFlag) {
    Print (L"sf: No lenght parameter!\n");
    return SHELL_ABORTED;
  } else if (LengthFlag) {
    ByteCount = ShellStrToUintn (LengthStr);
    if (ByteCount < 0) {
      Print (L"sf: Wrong length parameter %s!\n", LengthStr);
      return SHELL_ABORTED;
    }
  }

  if (FileFlag) {
    // Read FilePath parameter
    if (FileStr == NULL) {
      Print (L"sf: No FilePath parameter!\n");
      return SHELL_ABORTED;
    } else {
      FilePath = (CHAR16 *) FileStr;
      Status = ShellIsFile (FilePath);
      // When read file into flash, file doesn't have to exist
      if (EFI_ERROR (Status) && !(Flag & READ_FILE)) {
        Print (L"sf: Wrong FilePath parameter!\n");
        return SHELL_ABORTED;
      }
    }

    Status = OpenAndPrepareFile (FilePath, &FileHandle);
    if (EFI_ERROR(Status)) {
      Print (L"sf: Error while preparing file\n");
      return SHELL_ABORTED;
    }

    // Get file size in order to check correctness at the end of transfer
    if (Flag & (WRITE_FILE | UPDATE_FILE)) {
      Status = FileHandleGetSize (FileHandle, &FileSize);
      if (EFI_ERROR (Status)) {
        Print (L"sf: Cannot get file size\n");
      }
      ByteCount = (UINTN) FileSize;
    }

    FileBuffer = AllocateZeroPool ((UINTN) ByteCount);
    if (FileBuffer == NULL) {
      Print (L"sf: Cannot allocate memory\n");
      goto Error_Close_File;
    }

    // Read file content and store it in FileBuffer
    if (Flag & (WRITE_FILE | UPDATE_FILE)) {
      Status = FileHandleRead (FileHandle, &ByteCount, FileBuffer);
      if (EFI_ERROR (Status)) {
        Print (L"sf: Read from file error\n");
        goto Error_Free_Buffer;
      } else if (ByteCount != (UINTN) FileSize) {
        Print (L"sf: Not whole file read. Abort\n");
        goto Error_Free_Buffer;
      }
    }
  }

  Buffer = (UINT8 *)(UINTN)Address;
  if (FileFlag) {
    Buffer = FileBuffer;
  }

  switch (Flag) {
  case READ:
  case READ_FILE:
    Status = SpiFlashProtocol->Read (SpiFlashProtocol, Offset, Buffer, ByteCount);
    break;
  case ERASE:
    Status = SpiFlashProtocol->Erase (SpiFlashProtocol, Offset, ByteCount);
    break;
  case WRITE:
  case WRITE_FILE:
    Status = SpiFlashProtocol->Write (SpiFlashProtocol, Offset, Buffer, ByteCount);
    break;
  case UPDATE:
  case UPDATE_FILE:
    Status = SpiFlashProtocol->Update (SpiFlashProtocol, Offset, Buffer, ByteCount);
    break;
  }

  if (EFI_ERROR (Status)) {
    Print (L"sf: Error while performing spi transfer\n");
    return SHELL_ABORTED;
  }

  switch (Flag) {
  case ERASE:
    Print (L"sf: %d bytes succesfully erased at offset 0x%x\n", ByteCount,
      Offset);
    break;
  case WRITE:
  case WRITE_FILE:
    Print (L"sf: Write %d bytes at offset 0x%x\n", ByteCount, Offset);
    break;
  case UPDATE:
  case UPDATE_FILE:
    Print (L"sf: Update %d bytes at offset 0x%x\n", ByteCount, Offset);
    break;
  case READ:
    Print (L"sf: Read %d bytes from offset 0x%x\n", ByteCount, Offset);
    break;
  case READ_FILE:
    Status = FileHandleWrite (FileHandle, &ByteCount, FileBuffer);
    if (EFI_ERROR(Status)) {
      Print (L"sf: Error while writing into file\n");
      goto Error_Free_Buffer;
    }
    break;
  }

  if (FileFlag) {
    FreePool (FileBuffer);

    if (FileHandle != NULL) {
      ShellCloseFile (&FileHandle);
    }
  }

  return EFI_SUCCESS;

Error_Free_Buffer:
  FreePool (FileBuffer);
Error_Close_File:
  ShellCloseFile (&FileHandle);
  return SHELL_ABORTED;
}

EFI_STATUS
EFIAPI
ShellSpiFlashLibConstructor (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  gShellSfHiiHandle = NULL;

  gShellSfHiiHandle = HiiAddPackages (
                        &gShellSfHiiGuid, gImageHandle,
                        UefiShellSpiFlashLibStrings, NULL
                        );
  if (gShellSfHiiHandle == NULL) {
    return EFI_DEVICE_ERROR;
  }

  ShellCommandRegisterCommandName (
     L"sf", ShellCommandRunSpiFlash, ShellCommandGetManFileNameSpiFlash, 0,
     L"sf", TRUE , gShellSfHiiHandle, STRING_TOKEN (STR_GET_HELP_SF)
     );

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
ShellSpiFlashLibDestructor (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{

  if (gShellSfHiiHandle != NULL) {
    HiiRemovePackages (gShellSfHiiHandle);
  }
  return EFI_SUCCESS;
}