#include "pa_tz_api.h"
#include "config.h"
#include "driver_log.h"
#include "dr_qsee.h"
#include "entry.h"

#include "object.h"
#include "qsee_log.h"

#include <tee_internal_api.h>
#include <tees_sys.h>


RegisteredClientBuffer g_client_buffer[2];

/**
 * @brief Handler for QSEE IPC calls
 * @param [in] h Client context
 * @param [in] op Operation (command id)
 * @param [in,out] a Arguments
 * @param [in] k Count of arguments and its type
 * @return Result (Object_OK, Object_ERROR)
 */
static int32_t QseeIpcInvoke(ObjectCxt h, ObjectOp op, ObjectArg *a, ObjectCounts k);

/**
 * @brief Register existing translation Client Virt Address to Driver Virt Address
 * @param [in] client_addr Address in client space
 * @param [in] driver_addr Address in driver space
 * @param [in] size Size of memory
 * @param [in] idx Place to hold translation
 */
static inline void RegisterClientBuffer(DrvClientAddress client_addr, void *driver_addr, size_t size, size_t idx);

/**
 * @brief Get caller's application name
 * @param [in] credentials Object
 * @param [out] app_name Application name
 * @param [in,out] app_name_size Size of application name
 * @return Result (Object_OK, Object_ERROR)
 */
static int32_t GetAppName(Object credentials, char *app_name, size_t *app_name_size);

typedef union {
   TEE_UUID uuid;
   uint8_t raw[32]; // SM8150 requires at least 32 bytes
} UuidParam;

static UuidParam g_current_caller;

TEE_Result PlatformGetCallerUuid(TEE_UUID *uuid) {
  if (!uuid) {
    LOG_E("Invalid argument.\n");
    return TEE_ERROR_BAD_PARAMETERS;
  }

  *uuid = g_current_caller.uuid;

  return TEE_SUCCESS;
}

TEE_Result TA_CreateEntryPoint(void) {
  qsee_log_set_mask(0xFF);
  return TEE_SUCCESS;
}

void TA_DestroyEntryPoint(void) {
}

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

  if (TEE_PARAM_TYPE_GET(paramTypes, 0) != TEE_PARAM_TYPE_MEMREF_INOUT) {
    LOG_E("Type of input parameters have other type.\n");
    return TEE_ERROR_BAD_PARAMETERS;
  }

  TEE_Result result = Entry();
  if (result != TEE_SUCCESS) {
    LOG_E("Failed Entry.\n");
    LOG_D("Received result: 0x%x\n", result);
    return result;
  }

  return TEE_SUCCESS;
}

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

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

  if (TEE_PARAM_TYPE_GET(paramTypes, 0) != TEE_PARAM_TYPE_MEMREF_INOUT ||
      TEE_PARAM_TYPE_GET(paramTypes, 1) != TEE_PARAM_TYPE_VALUE_OUTPUT) {
    LOG_E("Type of input parameters have other type\n");
    return TEE_ERROR_BAD_PARAMETERS;
  }

  TEE_Result status = TEES_HandleDriverCommand(commandID,
      (void *)params[0].memref.buffer, params[0].memref.size,
      (void *)params[0].memref.buffer, &(params[0].memref.size));
  if (status != TEE_SUCCESS) {
    LOG_E("TEES_HandleDriverCommand return error.\n");
    LOG_D("Received result: 0x%x\n", status);
  }

  params[1].value.a = status;

  return TEE_SUCCESS;
}

int32_t tz_module_open(uint32_t uid, Object cred, Object *objOut) {
  if (uid == DRIVER_ID) {
    *objOut = (Object){QseeIpcInvoke, NULL};

    size_t ta_name_size = sizeof(g_current_caller);
    TEE_MemFill(&g_current_caller, 0, sizeof(g_current_caller));

    int32_t ret = GetAppName(cred, (char *)&g_current_caller.uuid, &ta_name_size);
    if (ret != Object_OK) {
      LOG_E("GetAppName() is failed.\n");
      LOG_D("Received result: %d\n", ret);
      return Object_ERROR;
    }

    LOG_D("Client name: %s.\n", &g_current_caller.uuid);

    return Object_OK;
  }

  LOG_E("Incorrect driver ID.\n");
  LOG_D("Driver ID: %x.\n", uid);

  *objOut = (Object){NULL, NULL};
  return Object_ERROR_INVALID;
}

static int32_t QseeIpcInvoke(ObjectCxt h, ObjectOp op, ObjectArg *a, ObjectCounts k) {
  switch (ObjectOp_methodID(op)) {
    case Object_OP_release: {
      return Object_OK;
    }
    case Object_OP_retain: {
      return Object_OK;
    }
    default: {
      /*
       * Content of buffers:
       * 0 - main input buffer with serialized command
       * 1 - additional input buffer (for TaskWriteFromTrustlet command)
       * 2 - registered memories input buffer (array with client's addresses of additional buffers)
       * 3 - main output buffer for command response
       * 4 - additional output buffer (for TaskReadToTrustlet command)
       */
      if (k != ObjectCounts_pack(3, 2, 0, 0)) {
        LOG_E("Unsupported object counts.\n");
        break;
      }

      // register additional input buffer
      RegisterClientBuffer(((DrvClientBuffers *)(a[2].b.ptr))->in_addr, a[1].b.ptr, a[1].b.size, 0);
      // register additional output buffer
      RegisterClientBuffer(((DrvClientBuffers *)(a[2].b.ptr))->out_addr, a[4].b.ptr, a[4].b.size, 1);

      const void *input = a[0].b.ptr;
      uint32_t input_size = (uint32_t)a[0].b.size;

      void *output = a[3].b.ptr;
      uint32_t *p_output_size = (uint32_t *)&a[3].b.size;

      TEE_Result status = TEES_HandleDriverCommand(0,
          input, input_size, output, p_output_size);

      if (CleanInternalMappings() != PA_TZ_SUCCESS) {
        LOG_E("Trustlet has mapped memory before exit.\n");
      }

      if (status != TEE_SUCCESS) {
        LOG_E("TEES_HandleDriverCommand return error.\n");
        LOG_D("Received result: 0x%x\n", status);
        return Object_ERROR;
      }

      return Object_OK;
    }
  }

  return Object_ERROR_INVALID;
}

static inline void RegisterClientBuffer(DrvClientAddress client_addr, void *driver_addr, size_t size, size_t idx) {
  g_client_buffer[idx] = (RegisteredClientBuffer){client_addr, driver_addr, size};
}

static int32_t GetAppName(Object credentials, char *app_name, size_t *app_name_size) {
  ObjectArg args[2];
  args[0].bi = (ObjectBufIn){"n", 1};
  args[1].b = (ObjectBuf){app_name, *app_name_size};

  int32_t result = Object_invoke(credentials, 1, args, ObjectCounts_pack(1, 1, 0, 0));
  if (result != Object_OK) {
    LOG_E("Object_invoke() is failed.\n");
    LOG_D("Received result: %d\n", result);
    return result;
  }

  *app_name_size = args[1].b.size;

  return Object_OK;
}