#include "libseL4.h"

#include <pthread.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

const int TYPE_UNTYPED = 0;
const int TYPE_CNODE = 1;
const int TYPE_UNTYPED_OLD = 2;
const int TYPE_TCB = 3;
const int TYPE_ENDPOINT = 4;

const int RET_OK = 0;
const int RET_ERR = 1;

cnode_info_t * initialize(int cnode_size_bits, int untyped_memory_size_bits) {
    int cnode_size = 16 * int_pow(2, cnode_size_bits);
    int untyped_memory_size = int_pow(2, untyped_memory_size_bits);
    cnode_info_t * root_cnode_info = malloc(16 + cnode_size + untyped_memory_size);
    void * root = (void *)((char *)root_cnode_info + 16);
    void * untyped_memory = (void *)((char *)root + cnode_size);

    untyped_object_info_t * mem_info = (untyped_object_info_t *)root;
    mem_info->size = untyped_memory_size;
    mem_info->address = untyped_memory;
    mem_info->type = TYPE_UNTYPED;

    root_cnode_info->radix = cnode_size_bits;
    root_cnode_info->address = root;
    root_cnode_info->type = TYPE_CNODE;

    return root_cnode_info;
}

info_t * look_up_cap_internal(cnode_info_t * root, uint32_t cap, int depth_limit, int bits_translated) {
    uint32_t cap_knocked_off_bits = (cap << bits_translated) >> bits_translated;
    uint32_t index = (cap_knocked_off_bits >> (32 - bits_translated - root->radix));
    info_t * info = root->address + index;

    bits_translated += root->radix;
    if (info->type != TYPE_CNODE || bits_translated >= depth_limit) {
        return info;
    }
    else {
        return look_up_cap_internal((cnode_info_t *)info, cap, depth_limit, bits_translated);
    }
}

info_t * look_up_cap(cnode_info_t * root, uint32_t cap, int depth_limit) {
    if (depth_limit == 0) {
        return (info_t *)root;
    }

    return look_up_cap_internal(root, cap, depth_limit, 0);
}

int seL4_Untyped_Retype(cnode_info_t * root, uint32_t cap, int type,
        int num_objects, int size_bits, uint32_t node_index, int node_depth,
        int node_offset) {

    info_t * candidate_untyped_info = look_up_cap(root, cap, 32);
    if (candidate_untyped_info->type != TYPE_UNTYPED) {
        return RET_ERR;
    }
    untyped_object_info_t * untyped_info = (untyped_object_info_t *)candidate_untyped_info;

    // Find the CNode in which to put the new entries
    info_t * candidate_cnode = look_up_cap(root, node_index, node_depth);
    if (candidate_cnode->type != TYPE_CNODE) {
        return RET_ERR;
    }
    cnode_info_t * cnode_info = (cnode_info_t *)candidate_cnode;

    int new_obj_size;
    if (type == TYPE_UNTYPED) {
        new_obj_size = int_pow(2, size_bits);
    }
    else if (type == TYPE_CNODE) {
        new_obj_size = 16 * int_pow(2, size_bits);
    }
    else if (type == TYPE_TCB) {
        new_obj_size = sizeof(tcb_t);
    }
    else if (type == TYPE_ENDPOINT) {
        new_obj_size = sizeof(endpoint_t);
    }
    else {
        return RET_ERR;
    }

    int required_size = num_objects * new_obj_size;
    if (required_size > untyped_info->size) {
        return RET_ERR;
    }

    // Mark the original untyped memory capability as old so that it cannot be
    // retyped again
    untyped_info->type = TYPE_UNTYPED_OLD;

    // For each new object, compute its address in memory and write a new info
    // struct (CNode slot entry) describing it
    for (int i = 0; i < num_objects; i++) {
        int offset = i * new_obj_size;
        info_t * new_info = (info_t *)(cnode_info->address + node_offset + i);
        if (type == TYPE_UNTYPED) {
            untyped_object_info_t * new_untyped_info = (untyped_object_info_t *)new_info;
            new_untyped_info->type = TYPE_UNTYPED;
            new_untyped_info->size = new_obj_size;
            new_untyped_info->address = (void *)((char *)untyped_info->address + offset);
        }
        else if (type == TYPE_CNODE) {
            cnode_info_t * new_cnode_info = (cnode_info_t *)new_info;
            new_cnode_info->type = TYPE_CNODE;
            new_cnode_info->radix = size_bits;
            new_cnode_info->address = (info_t *)((char *)untyped_info->address + offset);
        }
        else if (type == TYPE_TCB) {
            tcb_info_t * new_tcb_info = (tcb_info_t *)new_info;
            new_tcb_info->type = TYPE_TCB;
            new_tcb_info->address = (tcb_t *)((char *)untyped_info->address + offset);
        }
        else if (type == TYPE_ENDPOINT) {
            endpoint_info_t * new_endpoint_info = (endpoint_info_t *)new_info;
            new_endpoint_info->type = TYPE_ENDPOINT;
            new_endpoint_info->address = (endpoint_t *)((char *)untyped_info->address + offset);

            endpoint_t * endpoint = new_endpoint_info->address;
            pthread_mutex_init(&endpoint->mutex, NULL);
            pthread_cond_init(&endpoint->sender_wait_cond, NULL);
            pthread_cond_init(&endpoint->receiver_wait_cond, NULL);
            endpoint->num_waiting_receivers = 0;
            endpoint->transfer_in_progress = false;
        }
        else {
            return RET_ERR;
        }
    }
    return RET_OK;
}

int seL4_CNode_Copy(cnode_info_t * root, uint32_t cap, int cap_depth,
        uint32_t cnode_cap, int cnode_depth, int cnode_offset) {

    info_t * cap_info = look_up_cap(root, cap, cap_depth);
    info_t * candidate_cnode_info = look_up_cap(root, cnode_cap, cnode_depth);
    if (candidate_cnode_info->type != TYPE_CNODE) {
        return RET_ERR;
    }
    cnode_info_t * cnode_info = (cnode_info_t *)candidate_cnode_info;

    *(cnode_info->address + cnode_offset) = *cap_info;

    return RET_OK;
}

int seL4_TCB_SetSpace(cnode_info_t * root, uint32_t tcb_cap, uint32_t cnode_cap, int depth_limit) {
    info_t * candidate_tcb_info = look_up_cap(root, tcb_cap, 32);
    if (candidate_tcb_info->type != TYPE_TCB) {
        return RET_ERR;
    }

    info_t * candidate_cnode_info = look_up_cap(root, cnode_cap, depth_limit);
    if (candidate_cnode_info->type != TYPE_CNODE) {
        return RET_ERR;
    }

    tcb_info_t * tcb_info = (tcb_info_t *)candidate_tcb_info;
    cnode_info_t * cnode_info = (cnode_info_t *)candidate_cnode_info;

    tcb_t * tcb = tcb_info->address;
    tcb->cspace_root = cnode_info;

    return RET_OK;
}

int seL4_TCB_SetEntryPoint(cnode_info_t * root, uint32_t tcb_cap, void (* entry_point)(cnode_info_t *)) {
    info_t * candidate_tcb_info = look_up_cap(root, tcb_cap, 32);
    if (candidate_tcb_info->type != TYPE_TCB) {
        return RET_ERR;
    }

    tcb_info_t * tcb_info = (tcb_info_t *)candidate_tcb_info;
    tcb_t * tcb = tcb_info->address;
    tcb->entry_point = entry_point;

    return RET_OK;
}

void * pthread_entry_func(void * arg) {
    tcb_t * tcb = (tcb_t *)arg;
    tcb->entry_point(tcb->cspace_root);
    pthread_exit(0);
}

int seL4_TCB_Resume(cnode_info_t * root, uint32_t tcb_cap) {
    info_t * candidate_tcb_info = look_up_cap(root, tcb_cap, 32);
    if (candidate_tcb_info->type != TYPE_TCB) {
        return RET_ERR;
    }

    tcb_info_t * tcb_info = (tcb_info_t *)candidate_tcb_info;
    tcb_t * tcb = tcb_info->address;

    int ret = pthread_create(&(tcb->pthread), NULL, &pthread_entry_func, (void *)tcb);

    if (ret == 0) {
        return RET_OK;
    }
    else {
        return RET_ERR;
    }
}

int seL4_Send(cnode_info_t * root, uint32_t endpoint_cap, uint32_t * buf, int len) {
    if (len > ENDPOINT_MSG_LEN) {
        return RET_ERR;
    }

    info_t * slot = look_up_cap(root, endpoint_cap, 32);
    if (slot->type != TYPE_ENDPOINT) {
        return RET_ERR;
    }

    endpoint_info_t * endpoint_slot = (endpoint_info_t *)slot;
    endpoint_t * endpoint = endpoint_slot->address;

    pthread_mutex_lock(&endpoint->mutex);

    while (endpoint->num_waiting_receivers < 1) {
        pthread_cond_wait(&endpoint->sender_wait_cond, &endpoint->mutex);
    }

    endpoint->transfer_in_progress = true;
    for (int i = 0; i < len; i++) {
        endpoint->buf[i] = buf[i];
    }
    for (int i = len; i < ENDPOINT_MSG_LEN; i++) {
        endpoint->buf[i] = 0;
    }
    endpoint->msg_len = len;

    pthread_cond_signal(&endpoint->receiver_wait_cond);

    while (!endpoint->transfer_completed) {
        pthread_cond_wait(&endpoint->sender_wait_cond, &endpoint->mutex);
    }
    endpoint->transfer_completed = false;

    pthread_mutex_unlock(&endpoint->mutex);
}

int seL4_Recv(cnode_info_t * root, uint32_t endpoint_cap, uint32_t * buf, int buf_len, int * recv_len) {
    if (buf_len < ENDPOINT_MSG_LEN) {
        return RET_ERR;
    }

    info_t * slot = look_up_cap(root, endpoint_cap, 32);
    if (slot->type != TYPE_ENDPOINT) {
        return RET_ERR;
    }

    endpoint_info_t * endpoint_slot = (endpoint_info_t *)slot;
    endpoint_t * endpoint = endpoint_slot->address;

    pthread_mutex_lock(&endpoint->mutex);

    endpoint->num_waiting_receivers++;

    while (!endpoint->transfer_in_progress) {
        pthread_cond_signal(&endpoint->sender_wait_cond);
        pthread_cond_wait(&endpoint->receiver_wait_cond, &endpoint->mutex);
    }

    for (int i = 0; i < endpoint->msg_len; i++) {
        buf[i] = endpoint->buf[i];
    }
    *recv_len = endpoint->msg_len;

    endpoint->num_waiting_receivers--;
    endpoint->transfer_in_progress = false;
    endpoint->transfer_completed = true;

    pthread_cond_signal(&endpoint->sender_wait_cond);
    pthread_mutex_unlock(&endpoint->mutex);
}

int int_pow(int base, int exponent) {
    int result = 1;
    for (int i = 0; i < exponent; i++) {
        result = result * base;
    }
    return result;
}