#include "libseL4.h"

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

/*
 * Send and receive messages on an endpoint capability in a loop. Arguments
 * control what message is sent and whether to send a message before receiving.
 * This method assumes that the provided CSpace root has an endpoint capability
 * in slot 0.
 */
void ping_pong(cnode_info_t * cspace_root, bool send_first, char * msg);

/*
 * These methods are used to start the ping-pong threads, since only one
 * argument is allowed for a function used as a TCB entry point.
 */
void ping(cnode_info_t * cspace_root);
void pong(cnode_info_t * cspace_root);

/*
 * Print a dump of the specified memory range interpreted as integers.
 */
void print_memory_snapshot(void * mem, int length);

/*
 * Make sure all the info structs have sizes of 16 bytes.
 */
void check_struct_sizes();

int main(int argc, char ** argv) {
    bool print_memory = false;
    if (argc >= 2 && strcmp(argv[1], "-p") == 0) {
        print_memory = true;
    }

    // Check to make sure that all of the structs are the right size (i.e.,
    // that all of the info structs are 16 bytes)
    check_struct_sizes();

    // Create an initial CNode with 16 slots and an untyped memory region
    // with 1024 bytes
    cnode_info_t * root = initialize(4, 10);

    // Retype the entire 1024-byte untyped memory region into four 256-byte
    // untyped memory regions, placing their capabilities inside the root
    // CNode starting at slot 1.
    int ret = seL4_Untyped_Retype(root, 0x00000000, TYPE_UNTYPED, 4, 8, 0x00000000, 0, 1);
    if (ret != RET_OK) {
        printf("Error retyping (location A)\n");
        exit(1);
    }

    // Retype the first 256-byte untyped memory region into two CNodes with 8
    // slots each, placing their capabilities inside the root CNode starting
    // at slot 5.
    ret = seL4_Untyped_Retype(root, 0x10000000, TYPE_CNODE, 2, 3, 0x00000000, 0, 5);
    if (ret != RET_OK) {
        printf("Error retyping (location B)\n");
        exit(1);
    }

    // Retype the second 256-byte untyped memory region into two TCBs, placing
    // their capabilities inside the root CNode starting at slot 7
    ret = seL4_Untyped_Retype(root, 0x20000000, TYPE_TCB, 2, 1, 0x00000000, 0, 7);
    if (ret != RET_OK) {
        printf("Error retyping (location C)\n");
        exit(1);
    }

    // Retype the third 256-byte untyped memory region into an endpoint,
    // placing its capability inside the root CNode starting at slot 9
    ret = seL4_Untyped_Retype(root, 0x30000000, TYPE_ENDPOINT, 1, 1, 0x00000000, 0, 9);
    if (ret != RET_OK) {
        printf("Error retyping (location D)\n");
        exit(1);
    }

    // Copy the endpoint capability into the first child CNode
    ret = seL4_CNode_Copy(root, 0x90000000, 4, 0x50000000, 4, 0);
    if (ret != RET_OK) {
        printf("Error copying endpoint into first child CNode\n");
        exit(1);
    }

    // Copy the endpoint capability into the second child CNode
    ret = seL4_CNode_Copy(root, 0x90000000, 4, 0x60000000, 4, 0);
    if (ret != RET_OK) {
        printf("Error copying endpoint into second child CNode\n");
        exit(1);
    }

    // Initialize ping TCB
    ret = seL4_TCB_SetSpace(root, 0x70000000, 0x50000000, 4);
    if (ret != RET_OK) {
        printf("Error setting CSpace for first thread\n");
        exit(1);
    }
    ret = seL4_TCB_SetEntryPoint(root, 0x70000000, &ping);
    if (ret != RET_OK) {
        printf("Error setting entry point for first thread\n");
    }

    // Initialize pong TCB 
    ret = seL4_TCB_SetSpace(root, 0x80000000, 0x60000000, 4);
    if (ret != RET_OK) {
        printf("Error setting CSpace for second thread\n");
        exit(1);
    }
    ret = seL4_TCB_SetEntryPoint(root, 0x80000000, &pong);
    if (ret != RET_OK) {
        printf("Error setting entry point for second thread\n");
    }

    // Either print a snapshot of memory or start the threads
    if (print_memory) {
        print_memory_snapshot(root, 16 + 256 + 1024);
    }
    else {
        ret = seL4_TCB_Resume(root, 0x70000000);
        if (ret != RET_OK) {
            printf("Error starting first thread\n");
        }
        ret = seL4_TCB_Resume(root, 0x80000000);
        if (ret != RET_OK) {
            printf("Error starting second thread\n");
        }
    }

    pthread_exit(NULL);
}

void ping_pong(cnode_info_t * cspace_root, bool send_first, char * msg) {
    uint32_t recv_buf[ENDPOINT_MSG_LEN];
    int recv_len;
    uint32_t send_buf[ENDPOINT_MSG_LEN];
    int send_len;
    int ret;

    int send_buf_index = 0;
    char * current_char = msg;
    while (*current_char != '\0' && send_buf_index < ENDPOINT_MSG_LEN) {
        send_buf[send_buf_index] = (uint32_t)*current_char;
        send_buf_index++;
        current_char++;
    }
    send_len = send_buf_index;

    if (send_first) {
        int ret = seL4_Send(cspace_root, 0x00000000, send_buf, send_len);
        if (ret != RET_OK) {
            printf("Error sending message\n");
            exit(1);
        }

        sleep(1);
    }

    while (true) {
        int ret = seL4_Recv(cspace_root, 0x00000000, recv_buf, ENDPOINT_MSG_LEN, &recv_len);
        if (ret != RET_OK) {
            printf("Error receiving message\n");
            exit(1);
        }

        char recv_msg[recv_len + 1];
        for (int j = 0; j < recv_len; j++) {
            recv_msg[j] = (char)recv_buf[j];
        }
        recv_msg[recv_len] = '\0';

        printf("Received message: %s\n", recv_msg);

        ret = seL4_Send(cspace_root, 0x00000000, send_buf, send_len);
        if (ret != RET_OK) {
            printf("Error sending message\n");
            exit(1);
        }

        sleep(1);
    }
}

void ping(cnode_info_t * cspace_root) {
    ping_pong(cspace_root, true, "ping");
}

void pong(cnode_info_t * cspace_root) {
    ping_pong(cspace_root, false, "pong");
}

void print_memory_snapshot(void * mem, int length) {
    for (int i = 0; i < length / 4; i++) {
        int * curMem = (int *)((char *)mem + 4 * i);
        printf("%lu (%d): %d\n", (unsigned long)(mem + 4 * i), 4 * i, *curMem);
    }
}

void check_struct_sizes() {
    if (sizeof(info_t) != 16) {
        printf("Error: info_t struct is not 16 bytes long\n");
        exit(1);
    }

    if (sizeof(untyped_object_info_t) != 16) {
        printf("Error: untyped_object_info_t struct is not 16 bytes long\n");
        exit(1);
    }

    if (sizeof(cnode_info_t) != 16) {
        printf("Error: cnode_info_t struct is not 16 bytes long\n");
        exit(1);
    }

    if (sizeof(tcb_info_t) != 16) {
        printf("Error: tcb_info_t struct is not 16 bytes long\n");
        exit(1);
    }

    if (sizeof(endpoint_info_t) != 16) {
        printf("Error: endpoint_info_t struct is not 16 bytes long\n");
        exit(1);
    }
}