#include "command.h"
#include "task_nwd_access.h"

#include <string.h>

#include <tee_internal_api.h>
#include <tees_log.h>
#include <tees_mem.h>

#ifdef CONFIG_TEEGRIS40
#include <tees_extension.h>
#endif

#ifndef PAGE_SIZE
#define PAGE_SIZE (4096)
#endif

TEE_Result TA_CreateEntryPoint(void) {
  return TEE_SUCCESS;
}

void TA_DestroyEntryPoint(void) {
}

TEE_Result TA_OpenSessionEntryPoint(uint32_t paramTypes, TEE_Param params[4],
                                    void **sessionContext) {
  (void)paramTypes;
  (void)params;
  (void)sessionContext;

  return TEE_SUCCESS;
}

void TA_CloseSessionEntryPoint(void *sessionContext) {
  (void)sessionContext;
}

static PaTzResult Test_SimpleAuthenticate(TciCommand *command) {
  PaTzResult tz_result = PaTzAuthenticate(command->handler);

  command->tz_result = tz_result;

  return tz_result;
}

static PaTzResult Test_PerformanceAuthentication(TciCommand *command) {
  TEE_Time time_start, time_end;

  TEE_GetSystemTime(&time_start);

  PaTzResult tz_result = PaTzAuthenticate(command->handler);

  TEE_GetSystemTime(&time_end);

  command->tz_result = tz_result;
  command->authentication_exec_time =
      (time_end.seconds - time_start.seconds) * 1000 + time_end.millis - time_start.millis;

  return tz_result;
}

static PaTzResult Test_PerformanceRead(TciCommand *command) {
  uint8_t nwd_buffer[MAX_BUFFER_SIZE] = {0};
  TEE_Time time_start, time_end;

  TEE_GetSystemTime(&time_start);

  PaTzResult tz_result = PaTzReadFromNwdTask(command->handler,
      command->address, command->size, nwd_buffer);

  TEE_GetSystemTime(&time_end);

  command->tz_result = tz_result;

  command->authentication_exec_time =
      (time_end.seconds - time_start.seconds) * 1000 + time_end.millis - time_start.millis;

  for (int i = 0; i < command->size; ++i)
    if (nwd_buffer[i] != kStartInitValue) {
      tz_result = PA_TZ_GENERAL_ERROR;
      break;
    }

  return tz_result;
}

static PaTzResult Test_PerformanceWrite(TciCommand *command) {
  uint8_t swd_buffer[MAX_BUFFER_SIZE];
  TEE_Time time_start, time_end;

  // Fill SWd buffer by SWd pattern
  TEE_MemFill(swd_buffer, kUpdateInitValue, sizeof(swd_buffer));

  TEE_GetSystemTime(&time_start);

  PaTzResult tz_result = PaTzWriteToNwdTask(command->handler,
      (void *)swd_buffer, command->size,command->address);

  TEE_GetSystemTime(&time_end);

  command->tz_result = tz_result;

  command->authentication_exec_time =
      (time_end.seconds - time_start.seconds) * 1000 + time_end.millis - time_start.millis;

  return tz_result;
}

static PaTzResult Test_AuthenticateWithProvidedAppName(TciCommand *command) {
  PaTzRules rules = PaTzRulesCreate();
  PaTzResult tz_result;
  size_t len = 0;

  while (strlen(command->app_name + len)) {
    tz_result = PaTzRulesAddProcessName(rules, command->app_name + len);
    if (tz_result != PA_SUCCESS) {
      MB_LOGE("Can not add name to rules\n");
      return tz_result;
    }

    len += strlen(command->app_name + len) + 1;
  }

  tz_result = PaTzAuthenticateWithRules(command->handler, rules, &command->process_info);

  PaTzRulesDestroy(rules);

  return tz_result;
}

static PaTzResult Test_AuthenticateByPidAndAppName(TciCommand *command) {
  PaTzRules rules = PaTzRulesCreate();
  PaHandler handler = {0};
  PaTzResult tz_result;
  int pid;

  tz_result = PaTzRulesAddProcessName(rules, command->app_name);
  if (tz_result != PA_SUCCESS) {
    MB_LOGE("Can not add name to rules\n");
    return tz_result;
  }

#ifdef CONFIG_TEEGRIS40
  TEE_Result res;
  struct TEES_ClientCredentials creds;

  if ((res = TEES_GetClientCredentials(&creds)) != TEE_SUCCESS) {
    MB_LOGE("OpenSession: TEES_GetClientCredentials() failed with %#x\n", res);
    PaTzRulesDestroy(rules);

    return res;
  }

  pid = creds.pid;
#else
  pid = command->pid;
#endif

  PaTzResult result = PaTzHandlerCreateFromPid(pid, &handler);
  if (result != PA_SUCCESS) {
    MB_LOGE("PaTzHandlerCreateFromPid failed %d.\n", result);
    PaTzRulesDestroy(rules);

    return result;
  }

  tz_result = PaTzAuthenticateWithRules(handler, rules, &command->process_info);

  PaTzRulesDestroy(rules);
  PaTzHandlerDestroy(&handler);

  return tz_result;
}

static PaTzResult Test_AuthenticateByAppName(TciCommand *command) {
  PaHandler handler = {0};

  PaTzResult result = PaTzHandlerCreateFromProcessName(command->app_name, &handler);
  if (result != PA_SUCCESS) {
    MB_LOGE("PaTzHandlerCreateFromProcessName failed %d.\n", result);
    return result;
  }

  result = PaTzAuthenticate(handler);

  PaTzHandlerDestroy(&handler);

  return result;
}

static PaTzResult Test_StressAuthentication(TciCommand *command) {
  for (size_t i = 0; i < command->handler_num; ++i) {
    for (size_t j = 0; j < command->authentication_attempts_number; ++j) {
      MB_LOGD("Try #%u: Start authentication with {%u}-handler ...\n", j, i);
      PaTzResult handle_result = PaTzAuthenticate(command->handler_set[i]);
      if (handle_result != PA_TZ_SUCCESS) {
        MB_LOGE("CA authentication is FAILED [0x%08x].\n", handle_result);
        return handle_result;
      }
    }
  }

  return PA_TZ_SUCCESS;
}

TEE_Result TA_InvokeCommandEntryPoint(void *sessionContext, uint32_t commandID,
                                      uint32_t paramTypes, TEE_Param params[4]) {
  (void)sessionContext;

  PaTzResult handle_result = PA_TZ_GENERAL_ERROR;
  uint8_t nwd_buffer[MAX_BUFFER_SIZE] = {0};

  TciCommand tci_command = {0};
  size_t tci_command_len = params[0].memref.size;

  if (TEE_PARAM_TYPE_GET(paramTypes, 0) != TEE_PARAM_TYPE_MEMREF_INOUT
      || tci_command_len < sizeof(TciCommand)) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  TEE_MemMove(&tci_command, params[0].memref.buffer, sizeof(tci_command));

  tci_command.handler.flags = kPaHandlerTidUnusedMark;

  MB_LOGE("Invoking command %d in test trustlet.\n", commandID);

  switch (commandID) {
    case kCommandIdAuthenticate: {
      void *command_buffer = NULL;
      size_t command_buffer_size = 0;

      TEE_Result status_command_buffer = TEES_GetSharedCommandBuffer(
          &command_buffer, &command_buffer_size);

      if (status_command_buffer == TEE_ERROR_NOT_SUPPORTED) {
        MB_LOGW("TEES_GetSharedCommandBuffer isn't supported\n");
      } else if (status_command_buffer != TEE_SUCCESS) {
        MB_LOGE("Can't get address of command buffer\n");
        tci_command.resp = RESP_FAILED;
        break;
      }

      ProcessInfo process_info = {0};

      handle_result = PaTzAuthenticateWithCommandBuffer(tci_command.handler,
                                                        command_buffer,
                                                        command_buffer_size,
                                                        &process_info);
      break;
    }
    case kCommandReadFromNwdTask: {
      handle_result = ReadFromNwdTask(tci_command.handler, tci_command.address,
                                      tci_command.size, nwd_buffer,
                                      &tci_command.checksum_big_buffer);
      break;
    }
    case kCommandWriteToNwdTask: {
      handle_result = PaTzWriteToNwdTask(tci_command.handler,
                                         (const void *) MESSAGE_FROM_SWD,
                                         sizeof(MESSAGE_FROM_SWD),
                                         tci_command.address);
      break;
    }
    case kCommandReadBigBufferNwdTask: {
      handle_result = ReadBigBufferToNwdTask(tci_command.handler,
                                             tci_command.address_big_buffer,
                                             tci_command.size_big_buffer,
                                             &tci_command.checksum_big_buffer);

      break;
    }
    case kCommandWriteBigBufferNwdTask: {
      handle_result = WriteBigBufferToNwdTask(tci_command.handler,
                                              tci_command.address_big_buffer,
                                              tci_command.size_big_buffer,
                                              &tci_command.checksum_big_buffer);
      break;
    }
    case kPerformanceNativeClient: {
      handle_result = Test_PerformanceAuthentication(&tci_command);
      break;
    }
    case kPerformanceNativeClientRead: {
      handle_result = Test_PerformanceRead(&tci_command);
      break;
    }
    case kPerformanceNativeClientWrite: {
      handle_result = Test_PerformanceWrite(&tci_command);
      break;
    }
    case kTestSimpleAuthenticate: {
      handle_result = Test_SimpleAuthenticate(&tci_command);
      break;
    }
    case KAuthenticateWithProcessInfo: {
      handle_result = PaTzAuthenticateWithRules(tci_command.handler, NULL, &tci_command.process_info);
      break;
    }
    case KAuthenticateWithProvidedAppName: {
      handle_result = Test_AuthenticateWithProvidedAppName(&tci_command);
      break;
    }
    case kStressAuthentication: {
      handle_result = Test_StressAuthentication(&tci_command);
      break;
    }
    case kAuthenticateByAppName: {
      handle_result = Test_AuthenticateByAppName(&tci_command);
      break;
    }
    case kAuthenticateByPidAndAppName: {
      handle_result = Test_AuthenticateByPidAndAppName(&tci_command);
      break;
    }
    default: {
      MB_LOGE("Unknown command Id!\n");
      return TEE_ERROR_BAD_PARAMETERS;
    }
  }

  tci_command.tz_result = handle_result;

  if (handle_result != PA_TZ_SUCCESS) {
    tci_command.resp = RESP_FAILED;
  } else {
    tci_command.resp = RESP_SUCCESS;
  }

  MB_LOGE("Command: %d, handle_result: %x\n", commandID, handle_result);

  TEE_MemMove(params[0].memref.buffer, &tci_command, sizeof(tci_command));

  return TEE_SUCCESS;
}