/** @file
  Functions to deal with Disk buffer.

  Copyright (c) 2005 - 2018, Intel Corporation. All rights reserved. <BR>
  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include "HexEditor.h"
#include <Protocol/BlockIo.h>

extern EFI_HANDLE                 HImageHandleBackup;
extern HEFI_EDITOR_BUFFER_IMAGE   HBufferImage;

extern BOOLEAN                    HBufferImageNeedRefresh;
extern BOOLEAN                    HBufferImageOnlyLineNeedRefresh;
extern BOOLEAN                    HBufferImageMouseNeedRefresh;

extern HEFI_EDITOR_GLOBAL_EDITOR  HMainEditor;

HEFI_EDITOR_DISK_IMAGE            HDiskImage;
HEFI_EDITOR_DISK_IMAGE            HDiskImageBackupVar;

//
// for basic initialization of HDiskImage
//
HEFI_EDITOR_DISK_IMAGE            HDiskImageConst = {
  NULL,
  0,
  0,
  0
};

/**
  Initialization function for HDiskImage.

  @retval EFI_SUCCESS     The operation was successful.
  @retval EFI_LOAD_ERROR  A load error occurred.
**/
EFI_STATUS
HDiskImageInit (
  VOID
  )
{
  //
  // basically initialize the HDiskImage
  //
  CopyMem (&HDiskImage, &HDiskImageConst, sizeof (HDiskImage));

  CopyMem (&HDiskImageBackupVar, &HDiskImageConst, sizeof (HDiskImageBackupVar));

  return EFI_SUCCESS;
}

/**
  Backup function for HDiskImage. Only a few fields need to be backup.
  This is for making the Disk buffer refresh as few as possible.

  @retval EFI_SUCCESS           The operation was successful.
  @retval EFI_OUT_OF_RESOURCES  gST->ConOut of resources.
**/
EFI_STATUS
HDiskImageBackup (
  VOID
  )
{
  //
  // backup the disk name, offset and size
  //
  //
  SHELL_FREE_NON_NULL (HDiskImageBackupVar.Name);

  HDiskImageBackupVar.Name = CatSPrint(NULL, L"%s", HDiskImage.Name);
  if (HDiskImageBackupVar.Name == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  HDiskImageBackupVar.Offset  = HDiskImage.Offset;
  HDiskImageBackupVar.Size    = HDiskImage.Size;

  return EFI_SUCCESS;
}

/**
  Cleanup function for HDiskImage.

  @retval EFI_SUCCESS           The operation was successful.
**/
EFI_STATUS
HDiskImageCleanup (
  VOID
  )
{
  SHELL_FREE_NON_NULL (HDiskImage.Name);
  SHELL_FREE_NON_NULL (HDiskImageBackupVar.Name);

  return EFI_SUCCESS;
}

/**
  Set FileName field in HFileImage.

  @param[in] Str      File name to set.
  @param[in] Offset   The offset.
  @param[in] Size     The size.

  @retval EFI_SUCCESS           The operation was successful.
  @retval EFI_OUT_OF_RESOURCES  A memory allocation failed.
**/
EFI_STATUS
HDiskImageSetDiskNameOffsetSize (
  IN CONST CHAR16   *Str,
  IN UINTN    Offset,
  IN UINTN    Size
  )
{
  if (Str == HDiskImage.Name) {
    //
    // This function might be called using HDiskImage.FileName as Str.
    // Directly return without updating HDiskImage.FileName.
    //
    return EFI_SUCCESS;
  }

  //
  // free the old file name
  //
  SHELL_FREE_NON_NULL (HDiskImage.Name);
  HDiskImage.Name = AllocateCopyPool (StrSize (Str), Str);
  if (HDiskImage.Name == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  HDiskImage.Offset     = Offset;
  HDiskImage.Size       = Size;

  return EFI_SUCCESS;
}

/**
  Read a disk from disk into HBufferImage.

  @param[in] DeviceName   filename to read.
  @param[in] Offset       The offset.
  @param[in] Size         The size.
  @param[in] Recover      if is for recover, no information print.

  @retval EFI_SUCCESS           The operation was successful.
  @retval EFI_OUT_OF_RESOURCES  A memory allocation failed.
  @retval EFI_LOAD_ERROR        A load error occurred.
  @retval EFI_INVALID_PARAMETER A parameter was invalid.
**/
EFI_STATUS
HDiskImageRead (
  IN CONST CHAR16   *DeviceName,
  IN UINTN    Offset,
  IN UINTN    Size,
  IN BOOLEAN  Recover
  )
{
  CONST EFI_DEVICE_PATH_PROTOCOL  *DevicePath;
  EFI_DEVICE_PATH_PROTOCOL        *DupDevicePath;
  EFI_DEVICE_PATH_PROTOCOL        *DupDevicePathForFree;
  EFI_HANDLE                      Handle;
  EFI_BLOCK_IO_PROTOCOL           *BlkIo;
  EFI_STATUS                      Status;

  VOID                            *Buffer;
  CHAR16                          *Str;
  UINTN                           Bytes;

  HEFI_EDITOR_LINE                *Line;

  HBufferImage.BufferType = FileTypeDiskBuffer;

  DevicePath              = gEfiShellProtocol->GetDevicePathFromMap(DeviceName);
  if (DevicePath == NULL) {
    StatusBarSetStatusString (L"Cannot Find Device");
    return EFI_INVALID_PARAMETER;
  }
  DupDevicePath = DuplicateDevicePath(DevicePath);
  DupDevicePathForFree = DupDevicePath;
  //
  // get blkio interface
  //
  Status = gBS->LocateDevicePath(&gEfiBlockIoProtocolGuid,&DupDevicePath,&Handle);
  FreePool(DupDevicePathForFree);
  if (EFI_ERROR (Status)) {
    StatusBarSetStatusString (L"Read Disk Failed");
    return Status;
  }
  Status = gBS->OpenProtocol(Handle, &gEfiBlockIoProtocolGuid, (VOID**)&BlkIo, gImageHandle, NULL, EFI_OPEN_PROTOCOL_GET_PROTOCOL);
  if (EFI_ERROR (Status)) {
    StatusBarSetStatusString (L"Read Disk Failed");
    return Status;
  }
  //
  // if Offset exceeds LastBlock,
  //   return error
  //
  if (Offset > BlkIo->Media->LastBlock || Offset + Size > BlkIo->Media->LastBlock) {
    StatusBarSetStatusString (L"Invalid Offset + Size");
    return EFI_LOAD_ERROR;
  }

  Bytes   = BlkIo->Media->BlockSize * Size;
  Buffer  = AllocateZeroPool (Bytes);

  if (Buffer == NULL) {
    StatusBarSetStatusString (L"Read Disk Failed");
    return EFI_OUT_OF_RESOURCES;
  }

  //
  // read from disk
  //
  Status = BlkIo->ReadBlocks (
                    BlkIo,
                    BlkIo->Media->MediaId,
                    Offset,
                    Bytes,
                    Buffer
                    );

  if (EFI_ERROR (Status)) {
    FreePool (Buffer);
    StatusBarSetStatusString (L"Read Disk Failed");
    return EFI_LOAD_ERROR;
  }

  HBufferImageFree ();

  //
  // convert buffer to line list
  //
  Status = HBufferImageBufferToList (Buffer, Bytes);
  FreePool (Buffer);

  if (EFI_ERROR (Status)) {
    StatusBarSetStatusString (L"Read Disk Failed");
    return Status;
  }

  Status = HDiskImageSetDiskNameOffsetSize (DeviceName, Offset, Size);
  if (EFI_ERROR (Status)) {
    StatusBarSetStatusString (L"Read Disk Failed");
    return EFI_OUT_OF_RESOURCES;
  }
  //
  // initialize some variables
  //
  HDiskImage.BlockSize                = BlkIo->Media->BlockSize;

  HBufferImage.DisplayPosition.Row    = 2;
  HBufferImage.DisplayPosition.Column = 10;

  HBufferImage.MousePosition.Row      = 2;
  HBufferImage.MousePosition.Column   = 10;

  HBufferImage.LowVisibleRow          = 1;
  HBufferImage.HighBits               = TRUE;

  HBufferImage.BufferPosition.Row     = 1;
  HBufferImage.BufferPosition.Column  = 1;

  if (!Recover) {
    Str = CatSPrint(NULL, L"%d Lines Read", HBufferImage.NumLines);
    if (Str == NULL) {
      StatusBarSetStatusString (L"Read Disk Failed");
      return EFI_OUT_OF_RESOURCES;
    }

    StatusBarSetStatusString (Str);
    SHELL_FREE_NON_NULL (Str);

    HMainEditor.SelectStart = 0;
    HMainEditor.SelectEnd   = 0;

  }

  //
  // has line
  //
  if (HBufferImage.Lines != NULL) {
    HBufferImage.CurrentLine = CR (
                                HBufferImage.ListHead->ForwardLink,
                                HEFI_EDITOR_LINE,
                                Link,
                                EFI_EDITOR_LINE_LIST
                                );
  } else {
    //
    // create a dummy line
    //
    Line = HBufferImageCreateLine ();
    if (Line == NULL) {
      StatusBarSetStatusString (L"Read Disk Failed");
      return EFI_OUT_OF_RESOURCES;
    }

    HBufferImage.CurrentLine = Line;
  }

  HBufferImage.Modified           = FALSE;
  HBufferImageNeedRefresh         = TRUE;
  HBufferImageOnlyLineNeedRefresh = FALSE;
  HBufferImageMouseNeedRefresh    = TRUE;

  return EFI_SUCCESS;
}

/**
  Save lines in HBufferImage to disk.
  NOT ALLOW TO WRITE TO ANOTHER DISK!!!!!!!!!

  @param[in] DeviceName   The device name.
  @param[in] Offset       The offset.
  @param[in] Size         The size.

  @retval EFI_SUCCESS           The operation was successful.
  @retval EFI_OUT_OF_RESOURCES  A memory allocation failed.
  @retval EFI_LOAD_ERROR        A load error occurred.
  @retval EFI_INVALID_PARAMETER A parameter was invalid.
**/
EFI_STATUS
HDiskImageSave (
  IN CHAR16 *DeviceName,
  IN UINTN  Offset,
  IN UINTN  Size
  )
{

  CONST EFI_DEVICE_PATH_PROTOCOL  *DevicePath;
  EFI_DEVICE_PATH_PROTOCOL        *DupDevicePath;
  EFI_DEVICE_PATH_PROTOCOL        *DupDevicePathForFree;
  EFI_BLOCK_IO_PROTOCOL           *BlkIo;
  EFI_STATUS                      Status;
  EFI_HANDLE                      Handle;
  VOID                            *Buffer;
  UINTN                           Bytes;

  //
  // if not modified, directly return
  //
  if (HBufferImage.Modified == FALSE) {
    return EFI_SUCCESS;
  }

  HBufferImage.BufferType = FileTypeDiskBuffer;

  DevicePath              = gEfiShellProtocol->GetDevicePathFromMap(DeviceName);
  if (DevicePath == NULL) {
//    StatusBarSetStatusString (L"Cannot Find Device");
    return EFI_INVALID_PARAMETER;
  }
  DupDevicePath = DuplicateDevicePath(DevicePath);
  DupDevicePathForFree = DupDevicePath;

  //
  // get blkio interface
  //
  Status = gBS->LocateDevicePath(&gEfiBlockIoProtocolGuid,&DupDevicePath,&Handle);
  FreePool(DupDevicePathForFree);
  if (EFI_ERROR (Status)) {
//    StatusBarSetStatusString (L"Read Disk Failed");
    return Status;
  }
  Status = gBS->OpenProtocol(Handle, &gEfiBlockIoProtocolGuid, (VOID**)&BlkIo, gImageHandle, NULL, EFI_OPEN_PROTOCOL_GET_PROTOCOL);
  if (EFI_ERROR (Status)) {
//    StatusBarSetStatusString (L"Read Disk Failed");
    return Status;
  }

  Bytes   = BlkIo->Media->BlockSize * Size;
  Buffer  = AllocateZeroPool (Bytes);

  if (Buffer == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  //
  // concatenate the line list to a buffer
  //
  Status = HBufferImageListToBuffer (Buffer, Bytes);
  if (EFI_ERROR (Status)) {
    FreePool (Buffer);
    return Status;
  }

  //
  // write the buffer to disk
  //
  Status = BlkIo->WriteBlocks (
                    BlkIo,
                    BlkIo->Media->MediaId,
                    Offset,
                    Bytes,
                    Buffer
                    );

  FreePool (Buffer);

  if (EFI_ERROR (Status)) {
    return EFI_LOAD_ERROR;
  }
  //
  // now not modified
  //
  HBufferImage.Modified = FALSE;

  return EFI_SUCCESS;
}