#include <aos/aos.h>
#include <aos/aos_rpc.h>
#include <spawn/rpc_server.h>
#include <grading.h>
#include <spawn/spawn.h>

typedef errval_t (*rpc_handler_t)(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
);

static rpc_handler_t rpc_handlers[RPC_MTYPE_COUNT];

static void rpc_server_handle_recv(void *arg);
static void rpc_server_send_reply(void *arg);

void rpc_server_init(struct aos_rpc_server *rpc, struct lmp_chan *chan, void *shared_mem) {
    rpc->chan = chan;
    rpc->shared_mem = shared_mem;
}

void rpc_server_register_recv(struct aos_rpc_server *rpc) {
    errval_t err = lmp_chan_register_recv(rpc->chan, get_default_waitset(), MKCLOSURE(rpc_server_handle_recv, rpc));
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "Could not register receive handler");
    }
}

static void rpc_server_handle_recv(void *arg) {
    errval_t err;
    struct aos_rpc_server *rpc = (struct aos_rpc_server *)arg;

    struct lmp_recv_msg msg = LMP_RECV_MSG_INIT;
    struct capref arg_cap;

    err = lmp_chan_recv(rpc->chan, &msg, &arg_cap);
    assert(err_is_ok(err));

    rpc->ret_cap = NULL_CAP;
    rpc->ret_size = 0;
    rpc->ret0 = 0;
    rpc->ret1 = 0;
    if (msg.words[0] >= RPC_MTYPE_COUNT) {
        err = AOS_ERR_RPC_UNKNOWN_MSG_TYPE;
    } else if (msg.words[1] > RPC_SHARED_SIZE) {
        err = AOS_ERR_RPC_ARG_TOO_BIG;
    } else {
        rpc_handler_t handler = rpc_handlers[msg.words[0]];
        if (handler == NULL) {
            err = AOS_ERR_RPC_UNKNOWN_MSG_TYPE;
        } else {
            err = handler(
                rpc,
                arg_cap, msg.words[1], msg.words[2], msg.words[3],
                &rpc->ret_cap, &rpc->ret_size, &rpc->ret0, &rpc->ret1
            );
        }
    }
    rpc->ret_err = err;

    rpc_server_send_reply(rpc);
}

static void rpc_server_send_reply(void *arg) {
    struct aos_rpc_server *rpc = (struct aos_rpc_server *)arg;
    errval_t err;

    err = lmp_chan_send4(
        rpc->chan, LMP_FLAG_YIELD | LMP_FLAG_SYNC,
        rpc->ret_cap, rpc->ret_err, rpc->ret_size, rpc->ret0, rpc->ret1
    );
    if (err_is_fail(err)) {
        if (!lmp_err_is_transient(err)) {
            DEBUG_ERR(err, "Could not send RPC reply");
            return;
        }
        // Cannot send right now, try again later
        err = lmp_chan_register_send(rpc->chan, get_default_waitset(), MKCLOSURE(rpc_server_send_reply, arg));
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "Could not register send handler");
            return;
        }
    }

    // Sent reply, wait for next RPC call
    rpc_server_register_recv(rpc);
}

// RPC handlers

static errval_t handle_rpc_send_number(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    uintptr_t num = arg0;

    grading_rpc_handle_number(num);

    return SYS_ERR_OK;
}

static errval_t handle_rpc_send_string(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    const char* string = rpc->shared_mem;
    if (arg_size == 0 || string[arg_size - 1] != '\0') {
        return ERR_INVALID_ARGS;
    }

    grading_rpc_handler_string(string);

    return SYS_ERR_OK;
}

static errval_t handle_rpc_get_ram_cap(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    size_t bytes = arg0;
    size_t alignment = arg1;

    grading_rpc_handler_ram_cap(bytes, alignment);

    // TODO
    *ret0 = bytes;

    return LIB_ERR_NOT_IMPLEMENTED;
}

static errval_t handle_rpc_serial_getchar(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {

    grading_rpc_handler_serial_getchar();

    // MARKER SHELL: Replace with UART call
    *ret0 = getchar();

    return LIB_ERR_NOT_IMPLEMENTED;
}

static errval_t handle_rpc_serial_putchar(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    char c = arg0;

    grading_rpc_handler_serial_putchar(c);

    // MARKER SHELL: Replace with UART call
    putchar(c);

    return SYS_ERR_OK;
}

static errval_t handle_rpc_serial_write(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    // MARKER SHELL: Replace with UART call
    for (size_t i = 0; i < arg_size; ++i){
        putchar(((char *)rpc->shared_mem)[i]);
    }

    // return the number of characters written. Important for UART for early return when the buffer is full
    *ret0 = arg_size;

    return SYS_ERR_OK;
}

static errval_t handle_rpc_serial_read(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    // MARKER SHELL: Replace with UART call
    for (size_t i = 0; i < arg_size; ++i){
        ((char *)rpc->shared_mem)[i] = getchar();
    }

    // return the number of characters read. Important for UART for early return when the buffer is empty
    *ret0 = arg_size;

    return SYS_ERR_OK;
}

static errval_t handle_rpc_process_spawn(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    char *cmdline = rpc->shared_mem;
    coreid_t core = arg0;
    if (arg_size == 0 || cmdline[arg_size - 1] != '\0') {
        return ERR_INVALID_ARGS;
    }

    grading_rpc_handler_process_spawn(cmdline, core);

    struct spawninfo *si = malloc(sizeof(struct spawninfo));

    char *argv[MAX_CMDLINE_ARGS];
    int argc;
    spawn_parse_cmd(cmdline, &argc, argv);

    return spawn_load_argv(argc, argv, si, (domainid_t *)ret0);
}

static errval_t handle_rpc_process_get_name(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    domainid_t pid = arg0;

    grading_rpc_handler_process_get_name(pid);

    // search for the pid
    char *name = NULL;
    for (struct spawninfo *si = spawn_get_process_list(); si != NULL; si = si->next) {
        if (si->pid == pid) {
            name = si->binary_name;
            break;
        }
    }

    if (name == NULL) {
        return AOS_ERR_RPC_PROCESS_UNKNOWN_PID;
    }

    size_t name_len = strlen(name) + 1;
    if(name_len > RPC_SHARED_SIZE) {
        return AOS_ERR_RPC_RET_TOO_BIG;
    }

    memcpy(rpc->shared_mem, name, name_len);
    *ret_size = name_len;

    return SYS_ERR_OK;
}

static errval_t handle_rpc_process_get_all_pids(
    struct aos_rpc_server *rpc,
    struct capref arg_cap, size_t arg_size, uintptr_t arg0, uintptr_t arg1,
    struct capref *ret_cap, size_t *ret_size, uintptr_t *ret0, uintptr_t *ret1
) {
    grading_rpc_handler_process_get_all_pids();

    // size in bytes
    size_t list_size = 0;
    domainid_t *buf = rpc->shared_mem;

    // get all PIDs that fit into the given space
    for (struct spawninfo *si = spawn_get_process_list(); si != NULL; si = si->next) {
        if (list_size + sizeof(domainid_t) > RPC_SHARED_SIZE) {
            return AOS_ERR_RPC_RET_TOO_BIG;
        }

        *buf = si->pid;
        ++buf;
        list_size += sizeof(domainid_t);
    }

    *ret_size = list_size;

    return SYS_ERR_OK;
}


static rpc_handler_t rpc_handlers[RPC_MTYPE_COUNT] = {
    [RPC_MTYPE_SEND_NUMBER] = handle_rpc_send_number,
    [RPC_MTYPE_SEND_STRING] = handle_rpc_send_string,
    [RPC_MTYPE_GET_RAM_CAP] = handle_rpc_get_ram_cap,
    [RPC_MTYPE_SERIAL_GETCHAR] = handle_rpc_serial_getchar,
    [RPC_MTYPE_SERIAL_PUTCHAR] = handle_rpc_serial_putchar,
    [RPC_MTYPE_SERIAL_WRITE] = handle_rpc_serial_write,
    [RPC_MTYPE_SERIAL_READ] = handle_rpc_serial_read,
    [RPC_MTYPE_PROCESS_SPAWN] = handle_rpc_process_spawn,
    [RPC_MTYPE_PROCESS_GET_NAME] = handle_rpc_process_get_name,
    [RPC_MTYPE_PROCESS_GET_ALL_PIDS] = handle_rpc_process_get_all_pids,
};