#include <tees_log.h>  /* This is Multibuild's logging API header */

#include <tees_rpmb.h>

#include <rpmb.h>

#define D_HANDLE_INITIALIZED  0xA7A6A5A4
#define P_HANDLE_INITIALIZED  0xA4A5A6A7

typedef struct __TEES_RpmbDeviceHandle {
  TEES_RpmbDeviceInfo deviceInfo;
  uint32_t initialized;
} TEES_RpmbDeviceHandle_st;

typedef struct __TEES_RpmbPartitionHandle {
  TEES_RpmbDeviceHandle deviceHandle;
  TEES_RpmbPartitionInfo partitionInfo;
  uint32_t initialized;
} TEES_RpmbPartitionHandle_st;

TEE_Result TEES_RpmbDeviceInit(TEES_RpmbDeviceType deviceType,
                               TEES_RpmbDeviceHandle *device) {
  if (device == NULL) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  if (deviceType != TEES_RPMB_DEVICE_DEFAULT) {
    MB_LOGE("deviceType is not supported: deviceType = %d\n", deviceType);
    return TEE_ERROR_NOT_SUPPORTED;
  }

  *device = (struct __TEES_RpmbDeviceHandle *)
            TEE_Malloc(sizeof(struct __TEES_RpmbDeviceHandle), 0);
  if (NULL == *device) {
    MB_LOGE("Can not allocate memory.\n");
    return TEE_ERROR_GENERIC;
  }

  TEES_RpmbDeviceHandle_st *s = *device;
  s->deviceInfo.deviceId = 0;
  s->deviceInfo.totalSectors = RPMB_MAX_ADDRESS_NUM * RPMB_MAX_PARTITION_ID;
  s->deviceInfo.bytesPerSector = RPMB_BLOCK_SIZE;
  s->deviceInfo.freeSectors = 0;
  s->initialized = D_HANDLE_INITIALIZED;

  return TEE_SUCCESS;
}

TEE_Result TEES_RpmbOpenPartition(TEES_RpmbDeviceHandle device,
                                  uint32_t partitionId,
                                  TEES_RpmbPartitionHandle *partition) {
  if (device == NULL || device->initialized != D_HANDLE_INITIALIZED ||
      partition == NULL) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  if (partitionId > RPMB_MAX_PARTITION_ID) {
    MB_LOGE("Not supported partitionId: %u\n", partitionId);
    return TEE_ERROR_ITEM_NOT_FOUND;
  }

  *partition = (struct __TEES_RpmbPartitionHandle *)
               TEE_Malloc(sizeof(struct __TEES_RpmbPartitionHandle),
                          HINT_FILL_WITH_ZEROS);
  if (NULL == *partition) {
    MB_LOGE("Can not allocate memory.\n");
    return TEE_ERROR_GENERIC;
  }

  TEES_RpmbPartitionHandle_st *s = *partition;
  s->deviceHandle = device;

  s->partitionInfo.partitionId = partitionId;
  s->partitionInfo.totalSectors = RPMB_MAX_ADDRESS_NUM;
  s->partitionInfo.bytesPerSector = RPMB_BLOCK_SIZE;
  s->initialized = P_HANDLE_INITIALIZED;


  return TEE_SUCCESS;
}

TEE_Result TEES_RpmbCreatePartition(TEES_RpmbDeviceHandle device,
                                    uint32_t partitionId,
                                    uint32_t partitionSize,
                                    uint32_t flags) {
  (void)device;
  (void)partitionId;
  (void)partitionSize;
  (void)flags;

  if (flags != 0 || device == NULL ||
      device->initialized != D_HANDLE_INITIALIZED) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  return TEE_ERROR_NOT_SUPPORTED;
}

TEE_Result TEES_RpmbRemovePartition(TEES_RpmbDeviceHandle device,
                                    uint32_t partitionId,
                                    uint32_t flags) {
  (void)device;
  (void)partitionId;
  (void)flags;

  if (flags != 0 || device == NULL ||
      device->initialized != D_HANDLE_INITIALIZED) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  return TEE_ERROR_NOT_SUPPORTED;
}

void TEES_RpmbClosePartition(TEES_RpmbPartitionHandle partition) {
  TEE_MemFill(partition, 0, sizeof(struct __TEES_RpmbPartitionHandle));
  TEE_Free(partition);
}

void TEES_RpmbDeviceClose(TEES_RpmbDeviceHandle device) {
  TEE_MemFill(device, 0, sizeof(struct __TEES_RpmbDeviceHandle));
  TEE_Free(device);
}

TEE_Result TEES_RpmbGetDeviceInfo(TEES_RpmbDeviceHandle device,
                                  TEES_RpmbDeviceInfo *deviceInfo) {
  if (device == NULL || deviceInfo == NULL ||
      device->initialized != D_HANDLE_INITIALIZED) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  TEE_MemMove(deviceInfo, &device->deviceInfo, sizeof(TEES_RpmbDeviceInfo));

  return TEE_SUCCESS;
}

TEE_Result TEES_RpmbGetPartitionInfo(TEES_RpmbPartitionHandle partition,
                                     TEES_RpmbPartitionInfo *partitionInfo) {
  if (partition == NULL || partitionInfo == NULL ||
      partition->initialized != P_HANDLE_INITIALIZED) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  TEE_MemMove(partitionInfo, &partition->partitionInfo,
              sizeof(TEES_RpmbPartitionInfo));

  return TEE_SUCCESS;
}

TEE_Result TEES_RpmbWriteData(TEES_RpmbPartitionHandle partition,
                              uint8_t *dataBuffer,
                              uint32_t bufferLen,
                              uint32_t startSector,
                              uint32_t *sectors) {
  int ret;

  if (dataBuffer == NULL || bufferLen == 0 ||
      partition == NULL ||
      partition->initialized != P_HANDLE_INITIALIZED ||
      startSector >= partition->partitionInfo.totalSectors || sectors ==
      NULL) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  if (bufferLen > (partition->partitionInfo.totalSectors - startSector) *
      partition->partitionInfo.bytesPerSector) {
    return TEE_ERROR_STORAGE_NO_SPACE;
  }

  *sectors = (bufferLen / partition->partitionInfo.bytesPerSector);
  if (bufferLen % partition->partitionInfo.bytesPerSector != 0) {
    *sectors += 1;
  }

  if (*sectors <= RPMB_MAX_BLOCK_NUM) {
    ret = TEES_RPMBWrite(partition->partitionInfo.partitionId, startSector,
                         dataBuffer, *sectors, RPMB_TYPE_BLOCK);
    if (ret != 0) {
      MB_LOGE("tlApiRpmbWrite failure: 0x%x\n", ret);
      *sectors = 0;
      return TEE_ERROR_GENERIC;
    }
  } else {
    uint32_t sectorsLeft = *sectors;
    uint32_t position = startSector;
    while (sectorsLeft > RPMB_MAX_BLOCK_NUM) {
      ret = TEES_RPMBWrite(partition->partitionInfo.partitionId, position,
                           dataBuffer + (position - startSector) *
                           partition->partitionInfo.bytesPerSector,
                           RPMB_MAX_BLOCK_NUM, RPMB_TYPE_BLOCK);
      if (ret != 0) {
        MB_LOGE("tlApiRpmbWrite failure: 0x%x\n", ret);
        *sectors -= sectorsLeft;
        return TEE_ERROR_GENERIC;
      }
      position += RPMB_MAX_BLOCK_NUM;
      sectorsLeft -= RPMB_MAX_BLOCK_NUM;
    }

    ret = TEES_RPMBWrite(partition->partitionInfo.partitionId, position,
                         dataBuffer + (position - startSector) *
                         partition->partitionInfo.bytesPerSector, sectorsLeft, RPMB_TYPE_BLOCK);
    if (ret != 0) {
      MB_LOGE("tlApiRpmbWrite failure: 0x%x\n", ret);
      *sectors -= sectorsLeft;
      return TEE_ERROR_GENERIC;
    }
  }

  return TEE_SUCCESS;
}

TEE_Result TEES_RpmbReadData(TEES_RpmbPartitionHandle partition,
                             uint8_t *dataBuffer,
                             uint32_t *bufferLen,
                             uint32_t startSector,
                             uint32_t sectors) {
  int ret;

  if (dataBuffer == NULL || bufferLen == NULL ||
      partition == NULL ||
      partition->initialized != P_HANDLE_INITIALIZED ||
      startSector >= partition->partitionInfo.totalSectors ||
      sectors > partition->partitionInfo.totalSectors) {
    return TEE_ERROR_BAD_PARAMETERS;
  }

  if (*bufferLen < sectors * partition->partitionInfo.bytesPerSector) {
    return TEE_ERROR_SHORT_BUFFER;
  }

  *bufferLen = sectors * partition->partitionInfo.bytesPerSector;

  if (sectors <= RPMB_MAX_BLOCK_NUM) {
    ret = TEES_RPMBRead(partition->partitionInfo.partitionId, startSector,
                        dataBuffer, sectors, RPMB_TYPE_BLOCK);

    if (ret != 0) {
      MB_LOGE("tlApiRpmbRead failure: 0x%x\n", ret);
      *bufferLen = 0;
      return TEE_ERROR_GENERIC;
    }
  } else {
    uint32_t sectorsLeft = sectors;
    uint32_t position = startSector;

    while (sectorsLeft > RPMB_MAX_BLOCK_NUM) {
      ret = TEES_RPMBRead(partition->partitionInfo.partitionId, position,
                          dataBuffer + (position - startSector) *
                          partition->partitionInfo.bytesPerSector,
                          RPMB_MAX_BLOCK_NUM,
                          RPMB_TYPE_BLOCK);
      if (ret != 0) {
        MB_LOGE("tlApiRpmbRead failure: 0x%x\n", ret);
        *bufferLen = (sectors - sectorsLeft) *
                     partition->partitionInfo.bytesPerSector;
        return TEE_ERROR_GENERIC;
      }
      position += RPMB_MAX_BLOCK_NUM;
      sectorsLeft -= RPMB_MAX_BLOCK_NUM;
    }

    ret = TEES_RPMBRead(partition->partitionInfo.partitionId, position,
                        dataBuffer + (position - startSector) *
                        partition->partitionInfo.bytesPerSector,
                        sectorsLeft,
                        RPMB_TYPE_BLOCK);
    if (ret != 0) {
      MB_LOGE("tlApiRpmbRead failure: 0x%x\n", ret);
      *bufferLen = (sectors - sectorsLeft) *
                   partition->partitionInfo.bytesPerSector;
      return TEE_ERROR_GENERIC;
    }
  }

  return TEE_SUCCESS;
}