#include "tee_pthread_matrix_mult_TA.h"
#include <errno.h>
#include <pthread.h>
#include <stdio.h>
#include <tee_internal_api.h>
#include <string.h>

/* Use stringification for compute number of digits in INT_MAX value */
#define _XSTR(s) #s
#define XSTR(s) _XSTR(s)

enum { MULTIPLY_MATRIX_SIMPLE };

static TEE_Result test_pthread_math_multiply_matrix_simple(void);

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;
}

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

    switch (commandID) {
        case MULTIPLY_MATRIX_SIMPLE:
            return test_pthread_math_multiply_matrix_simple();

        default:
            printf("Wrong command id %u\n", commandID);
            return TEE_ERROR_BAD_PARAMETERS;
    }

}

#define M       2
#define N       3
#define P       2

int A[M][N] = {
    {1, 2, 3},
    {4, 5, 6}
};
int B[N][P] = {
    {6, 5},
    {4, 3},
    {2, 1}
};
/* Result-matrix */
int C[M][P];

int C_precomputed[M][P] = {
    {20, 14},
    {56, 41},
};

static void __print_matrix(int *arr, size_t rows, size_t cols)
{
/*  __INT_MAX__ is defined in secure-os toolchain */
#define INT_MAX_STRING XSTR(__INT_MAX__)
    size_t i;
    size_t j;

    /* Digits amount in representation of INT_MAX + minus + space */
    size_t row_elem_size = strlen(INT_MAX_STRING) + 2;
    size_t buf_size = row_elem_size * cols + 1;
    int written;

    char *buffer = TEE_Malloc(buf_size, 0);
    if (buffer == NULL) {
        printf("%s(): failed to allocate buffer for row of matrix\n", __func__);
        return;
    }

    for (i = 0; i < rows; ++i) {
        size_t sz = buf_size;
        char *buf = buffer;

        for (j = 0; j < cols; ++j) {
            written = snprintf(buf, sz, "%d ", *(arr + (i * cols) + j));
            if (written < 0 || (size_t) written >= sz) {
                printf("%s(): snprintf() failed. ret = %d\n",
                       __func__, written);
                return;
            }
            buf += written;
            sz -= written;
        }
        printf("%s\n", buffer);
    }

    TEE_Free(buffer);
}

static void *mat_mult_by_row(void *arg)
{
    int row = (* (int*) arg);
    int col;
    int i;
    int sum;

    for (col = 0; col < P; ++col) {
        sum = 0;
        for (i = 0; i < N; ++i) {
            sum += A[row][i] * B[i][col];
        }
        C[row][col] = sum;
    }
    return NULL;
}

static TEE_Result test_pthread_math_multiply_matrix_simple(void)
{
    int i;
    int threads_num = 0;
    int err = 0;
    TEE_Result res = TEE_SUCCESS;
    /* 'M' is the number of rows in the result-matrix 'C'.
     * So each thread computes each row of result-matrix. */
    pthread_t th[M];

    /* Create computational threads */
    for (i = 0; i < M; ++i, ++threads_num) {
        err = pthread_create(&th[i], NULL, mat_mult_by_row, (void *) &i);
        if (err) {
            printf("pthread_create() failed, errno %d, iter %d\n", err, i);
            res = TEE_ERROR_GENERIC;
            break;
        }
    }

    /* Wait for workers to finish their jobs */
    for (i = 0; i < threads_num; ++i) {
        err = pthread_join(th[i], NULL);
        if (err) {
            printf("pthread_join() failed, errno %d, iter %d\n", err, i);
            res |= TEE_ERROR_GENERIC;
        }
    }

    /* Hint: memcmp is fine. because there can't be any padding
     * between array elements or inside elements, because it is
     * array of int, not array of structs */
    if (!TEE_MemCompare(C, C_precomputed, sizeof(C))) {
        printf("%s(): Test passed\n", __func__);
        res = TEE_SUCCESS;
    } else {
        printf("%s(): Test failed\n", __func__);
        res = TEE_ERROR_GENERIC;
    }

    /* Display result-matrix */
    printf("Calculated result-matrix:\n");
    __print_matrix(&C[0][0], M, P);

    return res;
}
