/**
*  @Package     : BeniPkg
*  @FileName    : TestTcp.c
*  @Date        : 20231228
*  @Author      : Jiangwei
*  @Version     : 1.0
*  @Description :
*    MNP test.
*
*  @History:
*    20231228: Initialize.
*
*  This program and the accompanying materials
*  are licensed and made available under the terms and conditions of the BSD License
*  which accompanies this distribution. The full text of the license may be found at
*  http://opensource.org/licenses/bsd-license.php
*
*  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
*  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
**/

#include "Test.h"

TCP_TRANSPORT_PROTOCOL  *mTransport = NULL;
//
// Event notify functions, from which gBS->Exit shouldn't be called, can signal
// this event when the application should exit
//
STATIC EFI_EVENT  mFinishedEvent = NULL;

/*
  This is the NotifyFunction passed to CreateEvent in the FastbootAppEntryPoint
  It will be called by the UEFI event framework when the transport protocol
  implementation signals that data has been received from the Fastboot host.
  The parameters are ignored.

  @param[in]  Event                 The event.
  @param[in]  Context               The context of event.

  @retval  NA

*/
STATIC
VOID
DataReady (
  IN  EFI_EVENT                     Event,
  IN  VOID                          *Context
  )
{
  UINTN       Size;
  VOID        *Data;
  EFI_STATUS  Status;

  do {
    Status = mTransport->Receive (&Size, &Data);
    if (!EFI_ERROR (Status)) {
      BeniDumpHex (2, 0, Size, Data);
      FreePool (Data);
    }
  } while (!EFI_ERROR (Status));

  //
  // Quit if there was a fatal error.
  //
  if (Status != EFI_NOT_READY) {
    //
    // (Put a newline at the beginning as we are probably in the data phase,
    //  so the download progress line, with no '\n' is probably on the console)
    //
    Print (L"Fatal error receiving data. Exiting.\r\n");
    gBS->SignalEvent (mFinishedEvent);
  }
}

/**
  Test TCP.

  @param  NA

  @retval  NA

**/
VOID
TestTcp (
  VOID
  )
{
  EFI_STATUS                      Status = EFI_ABORTED;
  EFI_EVENT                       ReceiveEvent = NULL;
  EFI_EVENT                       WaitEventArray[2] = {NULL, NULL};
  UINTN                           EventIndex = 0;
  EFI_SIMPLE_TEXT_INPUT_PROTOCOL  *TextIn = NULL;
  EFI_INPUT_KEY                   Key = {0, 0};


  Status = gBS->LocateProtocol (
                  &gTcpTransportProtocolGuid,
                  NULL,
                  (VOID **)&mTransport
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "[%a][%d] Failed. - %r\n", __FUNCTION__, __LINE__, Status));
    return;
  }

  Status = gBS->LocateProtocol (
                  &gEfiSimpleTextInProtocolGuid,
                  NULL,
                  (VOID **)&TextIn
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "[%a][%d] Failed. - %r\n", __FUNCTION__, __LINE__, Status));
    return;
  }

  //
  // Disable watchdog.
  //
  Status = gBS->SetWatchdogTimer (0, 0x10000, 0, NULL);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "[%a][%d] Failed. - %r\n", __FUNCTION__, __LINE__, Status));
    return;
  }

  //
  // Create event for receipt of data from the host.
  //
  Status = gBS->CreateEvent (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  DataReady,
                  NULL,
                  &ReceiveEvent
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "[%a][%d] Failed. - %r\n", __FUNCTION__, __LINE__, Status));
    return;
  }

  //
  // Start listening for data.
  //
  Status = mTransport->Start (ReceiveEvent);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "[%a][%d] Failed. - %r\n", __FUNCTION__, __LINE__, Status));
    return;
  }

  Print (L"Press RETURN or SPACE key to quit.\r\n");

  //
  // Quit when the user presses any key, or mFinishedEvent is signalled.
  //
  WaitEventArray[0] = mFinishedEvent;
  WaitEventArray[1] = TextIn->WaitForKey;
  while (1) {
    gBS->WaitForEvent (2, WaitEventArray, &EventIndex);
    Status = TextIn->ReadKeyStroke (gST->ConIn, &Key);
    if (Key.ScanCode == SCAN_NULL) {
      if ((Key.UnicodeChar == CHAR_CARRIAGE_RETURN) ||
          (Key.UnicodeChar == L' ')) {
        break;
      }
    }
  }

  mTransport->Stop ();

  return;
}