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

extern coreid_t my_core_id;
rpc_handler_t rpc_handlers[RPC_MTYPE_COUNT];
extern struct bootinfo *bi;

extern struct aos_urpc urpc_to_bsp;
extern struct aos_urpc urpc_to_app;
extern struct waitset urpc_to_app_ws;

static void rpc_server_handle_recv(void *arg);
static void rpc_server_send_reply(void *arg);
static void rpc_server_async_reply(struct generic_rpc_server *g_rpc, errval_t ret_err);

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

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 {
            if (my_core_id == 0 || 
                (msg.words[0] == RPC_MTYPE_PROCESS_SPAWN && msg.words[1] == my_core_id)/* dont forward spawn rpcs */) {
                err = handler(
                    &rpc->g,
                    arg_cap, msg.words[1], msg.words[2], msg.words[3],
                    &rpc->ret_cap, &rpc->ret_size, &rpc->ret0, &rpc->ret1
                );
            } else {
                //forward rpc to urpc
                if (msg.words[1] > 0) {
                    assert(msg.words[1] <= RPC_SHARED_SIZE);
                    memcpy(urpc_to_bsp.shared_mem, rpc->g.shared_mem, msg.words[1]);
                }

                err = do_aos_urpc(
                    &urpc_to_bsp, msg.words[0],
                    arg_cap, msg.words[1], msg.words[2], msg.words[3],
                    &rpc->ret_cap, &rpc->ret_size, &rpc->ret0, &rpc->ret1
                );

                if (rpc->ret_size > 0) {
                    assert(rpc->ret_size <= RPC_SHARED_SIZE);
                    memcpy(rpc->g.shared_mem, urpc_to_bsp.shared_mem, rpc->ret_size);
                }
            }
        }
    }

    if (err == AOS_ERR_RPC_ASYNC_REPLY) return;
    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);
}

static void rpc_server_async_reply(struct generic_rpc_server *g_rpc, errval_t ret_err) {
    struct aos_rpc_server *rpc = (struct aos_rpc_server *)g_rpc;
    rpc->ret_err = ret_err;
    rpc_server_send_reply(rpc);
}

// RPC handlers

static errval_t handle_rpc_nop(
    struct generic_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
) {
    __asm volatile("nop");

    return SYS_ERR_OK;
}

static errval_t handle_rpc_send_number(
    struct generic_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 generic_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 generic_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
) {
    errval_t err;

    size_t bytes = arg0;
    size_t alignment = arg1;

    err = ram_alloc_aligned(ret_cap, bytes, alignment);
    if (err_is_fail(err)) return err;
    *ret0 = bytes;

    grading_rpc_handler_ram_cap(bytes, alignment);

    return SYS_ERR_OK;
}

static errval_t handle_rpc_serial_getchar(
    struct generic_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 SYS_ERR_OK;
}

static errval_t handle_rpc_serial_putchar(
    struct generic_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 generic_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
    // return the number of characters written. Important for UART for early return when the buffer is full
    *ret0 = dummy_terminal_write((char *)rpc->shared_mem, arg_size);

    return SYS_ERR_OK;
}

static errval_t handle_rpc_serial_read(
    struct generic_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
    // return the number of characters read. Important for UART for early return when the buffer is empty
    *ret_size = dummy_terminal_read((char *)rpc->shared_mem, arg0);

    return SYS_ERR_OK;
}

struct rpc_to_urpc_arg {
    struct waitset_chanstate chan;
    struct generic_rpc_server *rpc;
    uintptr_t msg_type;
    size_t arg_size;
    uintptr_t arg0;
    uintptr_t arg1;
    size_t *ret_size;
    uintptr_t *ret0;
    uintptr_t *ret1;
};

// forward an RPC call over URPC, asynchronously (meaning without blocking the default waitset)
static void rpc_to_urpc(void *arg) {
    errval_t err;
    struct rpc_to_urpc_arg *urpc_arg = arg;

    if (urpc_arg->arg_size > 0) {
        assert(urpc_arg->arg_size <= RPC_SHARED_SIZE);
        memcpy(urpc_to_app.shared_mem, urpc_arg->rpc->shared_mem, urpc_arg->arg_size);
    }

    err = do_aos_urpc(
        &urpc_to_app, urpc_arg->msg_type,
        NULL_CAP, urpc_arg->arg_size, urpc_arg->arg0, urpc_arg->arg1,
        NULL, urpc_arg->ret_size, urpc_arg->ret0, urpc_arg->ret1
    );

    if (*urpc_arg->ret_size > 0) {
        assert(*urpc_arg->ret_size <= RPC_SHARED_SIZE);
        memcpy(urpc_arg->rpc->shared_mem, urpc_to_app.shared_mem, *urpc_arg->ret_size);
    }

    urpc_arg->rpc->async_reply(urpc_arg->rpc, err);
    free(urpc_arg);
}

static errval_t handle_rpc_process_spawn(
    struct generic_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);

    if (core == my_core_id) {
        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);
    }

    assert(my_core_id == 0); // Spawning from core 1 on core 0 is not supported.
    assert(rpc->async_reply != NULL);

    struct rpc_to_urpc_arg *urpc_arg = malloc(sizeof(struct rpc_to_urpc_arg));
    if (urpc_arg == NULL) return LIB_ERR_MALLOC_FAIL;
    waitset_chanstate_init(&urpc_arg->chan, CHANTYPE_OTHER);
    urpc_arg->msg_type = RPC_MTYPE_PROCESS_SPAWN;
    urpc_arg->rpc = rpc;
    urpc_arg->arg_size = arg_size;
    urpc_arg->arg0 = arg0;
    urpc_arg->arg1 = arg1;
    urpc_arg->ret_size = ret_size;
    urpc_arg->ret0 = ret0;
    urpc_arg->ret1 = ret1;

    waitset_chan_trigger_closure(&urpc_to_app_ws, &urpc_arg->chan,
        MKCLOSURE(rpc_to_urpc, urpc_arg));

    return AOS_ERR_RPC_ASYNC_REPLY;
}

static errval_t handle_rpc_process_get_name(
    struct generic_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);

    *ret_size = RPC_SHARED_SIZE;
    return get_process_name(pid, rpc->shared_mem, ret_size);
}

static errval_t handle_rpc_process_get_all_pids(
    struct generic_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();

    domainid_t *buf = rpc->shared_mem;
    size_t list_size = RPC_SHARED_SIZE / sizeof(domainid_t);
    errval_t err = get_all_pids(buf, &list_size);
    if (err_is_fail(err)) return err;
    *ret_size = list_size * sizeof(domainid_t);
    return SYS_ERR_OK;
}

static errval_t handle_rpc_allocate_pid(
    struct generic_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* process_name = rpc->shared_mem;
    if (arg_size == 0 || process_name[arg_size - 1] != '\0') {
        return ERR_INVALID_ARGS;
    }

    domainid_t pid;
    errval_t err = allocate_pid(process_name, &pid);
    if (err_is_fail(err)) return err;
    *ret0 = pid;
    return SYS_ERR_OK;
}

static errval_t serialize_bootinfo(struct bootinfo * bootinfo, struct bootinfo_serialized * ret) {
    errval_t err;
    assert(bootinfo != NULL);
    // struct capref serialized_bootinfo_frame;

    //allocate space for serialized
    // size_t bootinfo_size = sizeof(struct bootinfo_serialized) + bootinfo->regions_length * sizeof(struct mem_region_serialized);
    // err = frame_alloc(&serialized_bootinfo_frame, bootinfo_size, NULL);
    // if (err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_ALLOC);

    struct bootinfo_serialized * serialized_bootinfo = ret;
    // struct bootinfo_serialized * serialized_bootinfo;
    // err = paging_map_frame(get_current_paging_state(), (void **) &serialized_bootinfo, bootinfo_size, serialized_bootinfo_frame);
    // if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_MAP);

    serialized_bootinfo->host_msg = bootinfo->host_msg;
    serialized_bootinfo->host_msg_bits = bootinfo->host_msg_bits;
    serialized_bootinfo->mem_spawn_core = bootinfo->mem_spawn_core;
    serialized_bootinfo->regions_length = bootinfo->regions_length;

    struct capref slot_0_cap = {
        .cnode = cnode_module,
        .slot = 0,
    };
    struct capability slot_0_cap_id;
    err = cap_direct_identify(slot_0_cap, &slot_0_cap_id);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_CAP_IDENTIFY);

    serialized_bootinfo->slot_0_frame_base = slot_0_cap_id.u.frame.base;
    serialized_bootinfo->slot_0_frame_bytes = slot_0_cap_id.u.frame.bytes;

    for (int i = 0; i < bootinfo->regions_length; ++i) {
        serialized_bootinfo->regions[i].mr_base = bootinfo->regions[i].mr_base;
        serialized_bootinfo->regions[i].mr_bytes = bootinfo->regions[i].mr_bytes;
        serialized_bootinfo->regions[i].mr_type = bootinfo->regions[i].mr_type;
        serialized_bootinfo->regions[i].mr_consumed = bootinfo->regions[i].mr_consumed;

        serialized_bootinfo->regions[i].mrmod_slot = bootinfo->regions[i].mrmod_slot;
        serialized_bootinfo->regions[i].mrmod_data = bootinfo->regions[i].mrmod_data;
        serialized_bootinfo->regions[i].mrmod_size = bootinfo->regions[i].mrmod_size;

        if (bootinfo->regions[i].mr_type != RegionType_Module) {
            serialized_bootinfo->regions[i].mrmod_frame_base = 0;
            serialized_bootinfo->regions[i].mrmod_frame_bytes = 0;
            continue;
        }

        struct capref elf_frame_cap = {
            .cnode = cnode_module,
            .slot = bootinfo->regions[i].mrmod_slot,
        };

        struct capability elf_frame_id;
        err = cap_direct_identify(elf_frame_cap, &elf_frame_id);
        if (err_is_fail(err)) return err_push(err, LIB_ERR_CAP_IDENTIFY);

        // debug_printf("[serialize_bootinfo] elf_frame_id.type = %d\n", elf_frame_id.type);
        assert(elf_frame_id.type == ObjType_DevFrame);

        serialized_bootinfo->regions[i].mrmod_frame_base = elf_frame_id.u.devframe.base;
        serialized_bootinfo->regions[i].mrmod_frame_bytes = elf_frame_id.u.devframe.bytes;
    }

    return SYS_ERR_OK;
}


static errval_t handle_rpc_get_bootinfo(
    struct generic_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
) {
    errval_t err;
    size_t bootinfo_serialized_size = sizeof(struct bootinfo_serialized) + bi->regions_length * sizeof(struct mem_region_serialized);
    assert(bootinfo_serialized_size <= RPC_SHARED_SIZE);

    err = serialize_bootinfo(bi, (struct bootinfo_serialized *) rpc->shared_mem);
    if (err_is_fail(err)) return err;

    *ret_size = bootinfo_serialized_size;

    return SYS_ERR_OK;
}

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,
    [RPC_MTYPE_ALLOCATE_PID] = handle_rpc_allocate_pid,
    [RPC_MTYPE_NOP] = handle_rpc_nop,
    [RPC_MTYPE_GET_BOOTINFO] = handle_rpc_get_bootinfo,
};