#include "test_TA_spi_driver.h"
#include <stdio.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <tee_spi.h>
#include <tee_internal_api.h>

enum cmd {
    CMD_SPI_CONFIG = 0,
    CMD_SPI_WRITE,
    CMD_SPI_READ,
    CMD_SPI_WRITE_READ,
    CMD_SPI_WRITE_READ_NOVERIFY
};

#define BUF_SIZE (32 * 1024)
#define SPI_MAGIC_CHAR 0x55

static unsigned char tx_buf[BUF_SIZE] = { 0 };
static unsigned char rx_buf[BUF_SIZE] = { 0 };

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;

    return;
}

static TEE_Result cmd_spi_config(unsigned int port)
{
    printf("TA: CMD_SPI_CONFIG received\n");

    TEES_SPIHandler handler;
    TEES_SPIConfig cfg = {
        .speedHz = 5000000,
        .CPOL = 0,
        .CPHA = 0,
        .bitsPerWord = SPI_TYPE_HWORD,
        .isDMAMode = false
    };

    TEE_Result res = TEES_SPIInit(port, NULL, &handler);
    if (res != TEE_SUCCESS) {
        printf("TEST: Failed to initalize SPI, tee_err: %#x\n", res);
        return res;
    }

    res = TEES_SPISetConfig(handler, &cfg);
    if (res != TEE_SUCCESS) {
        printf("TEST: Failed to configure SPI, tee_err: %#x\n", res);
        TEES_SPIExit(handler);
        return res;
    }

    if (TEES_SPIExit(handler)) {
        printf("Faled to close SPI device\n");
        TEE_Panic(0);
    }

    if (res != TEE_SUCCESS) {
        printf("Failed to set SPI speed. err: %#x\n", res);
    }

    return res;
}

static TEE_Result cmd_spi_read(unsigned int port)
{
    printf("TA: CMD_SPI_READ received\n");

    TEES_SPIHandler handler;
    TEES_SPIConfig cfg = {
        .speedHz = 5000000,
        .CPOL = 0,
        .CPHA = 0,
        .bitsPerWord = SPI_TYPE_HWORD,
        .isDMAMode = false
    };

    TEE_Result res = TEES_SPIInit(port, &cfg, &handler);
    if (res != TEE_SUCCESS) {
        printf("TEST: Failed to initalize SPI, tee_err: %#x\n", res);
        return res;
    }

    TEES_SPITransfer rx = {
        .bufAddr = rx_buf,
        .bufLen = sizeof(rx_buf),
        .transferredLen = 0
    };

    res = TEES_SPIRead(handler, &rx);
    if (TEES_SPIExit(handler)) {
        TEE_Panic(0);
    }

    if (rx.transferredLen == 0) {
        printf("TEST: nothing was transferred over SPI, err: %#x\n", res);
        return TEE_ERROR_GENERIC;
    }

    printf("RX: len: %zu,  %.50s\n", rx.transferredLen, rx_buf);
    return res;
}

static TEE_Result cmd_spi_write(unsigned int port)
{
    printf("TA: CMD_SPI_WRITE transmitted\n");

    TEES_SPIHandler handler;
    TEES_SPIConfig cfg = {
        .speedHz = 5000000,
        .CPOL = 0,
        .CPHA = 0,
        .bitsPerWord = SPI_TYPE_HWORD,
        .isDMAMode = false
    };

    TEE_Result res = TEES_SPIInit(port, &cfg, &handler);
    if (res != TEE_SUCCESS) {
        printf("TEST: Failed to initalize SPI, tee_err: %#x\n", res);
        return res;
    }

    TEES_SPITransfer tx = {
        .bufAddr = tx_buf,
        .bufLen = sizeof(tx_buf),
        .transferredLen = 0
    };

    TEE_MemFill(tx_buf, SPI_MAGIC_CHAR, sizeof(tx_buf));
    res = TEES_SPIWrite(handler, &tx);
    if (TEES_SPIExit(handler)) {
        TEE_Panic(0);
    }

    if (tx.transferredLen == 0) {
        printf("TEST: nothing was transferred over SPI, err: %#x\n", res);
        return TEE_ERROR_GENERIC;
    }

    printf("TX: len: %zu,  %.50s\n", tx.transferredLen, tx_buf);
    return res;
}

static TEE_Result cmd_spi_write_read(unsigned int port, bool check_validity)
{
    printf("TA: CMD_SPI_WRITE_READ received\n");

    TEES_SPIHandler handler;
    TEES_SPIConfig cfg = {
        .speedHz = 5000000,
        .CPOL = 0,
        .CPHA = 0,
        .bitsPerWord = SPI_TYPE_HWORD,
        .isDMAMode = false
    };

    TEE_Result res = TEES_SPIInit(port, &cfg, &handler);
    if (res != TEE_SUCCESS) {
        printf("TEST: Failed to initalize SPI, tee_err: %#x\n", res);
        return res;
    }

    TEES_SPITransfer tx = {
        .bufAddr = tx_buf,
        .bufLen = sizeof(tx_buf),
        .transferredLen = 0
    };

    TEES_SPITransfer rx = {
        .bufAddr = rx_buf,
        .bufLen = sizeof(rx_buf),
        .transferredLen = 0
    };

    TEE_MemFill(tx_buf, SPI_MAGIC_CHAR, sizeof(tx_buf));
    res = TEES_SPIWriteRead(handler, &tx, &rx);
    if (TEES_SPIExit(handler)) {
        TEE_Panic(0);
    }

    if (check_validity) {
        if (rx.transferredLen == 0 || tx.transferredLen == 0) {
            printf("TEST: nothing was transferred over SPI, err: %#x\n", res);
            return TEE_ERROR_GENERIC;
        }

        if (TEE_MemCompare(tx_buf, rx_buf, BUF_SIZE) != 0) {
            printf("TST: SPI_WRITE_READ error: incorrect data received\n");
            return TEE_ERROR_GENERIC;
        }
    }

    printf("Tx transferred: %zu, Rx transferred: %zu\n", tx.transferredLen, rx.transferredLen);
    printf("TX: %.50s\n", tx_buf);
    printf("RX: %.50s\n", rx_buf);

    return res;
}

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

    TEE_Result ret = 0;
    const unsigned int port = params[0].value.a;

    switch (commandID) {
        case CMD_SPI_CONFIG: {
            ret = cmd_spi_config(port);
            break;
        }
        case CMD_SPI_READ: {
            ret = cmd_spi_read(port);
            break;
        }
        case CMD_SPI_WRITE: {
            ret = cmd_spi_write(port);
            break;
        }
        case CMD_SPI_WRITE_READ: {
            ret = cmd_spi_write_read(port, true);
            break;
        }
        case CMD_SPI_WRITE_READ_NOVERIFY: {
            ret = cmd_spi_write_read(port, false);
            break;
        }

        default:
            return TEE_ERROR_GENERIC;
    }

    return ret;
}
