/**
 * \file
 * \brief init process for child spawning
 */

/*
 * Copyright (c) 2007, 2008, 2009, 2010, 2016, ETH Zurich.
 * All rights reserved.
 *
 * This file is distributed under the terms in the attached LICENSE file.
 * If you do not find this file, copies can be found by writing to:
 * ETH Zurich D-INFK, Haldeneggsteig 4, CH-8092 Zurich. Attn: Systems Group.
 */

#include <stdio.h>
#include <stdlib.h>

#include <aos/aos.h>
#include <aos/morecore.h>
#include <aos/paging.h>
#include <aos/waitset.h>
#include <aos/aos_rpc.h>
#include <aos/aos_urpc.h>
#include <mm/mm.h>
#include <spawn/spawn.h>
#include <grading.h>

#include "mem_alloc.h"
#include <aos/coreboot.h>
#include <aos/cache.h>
#include <aos/kernel_cap_invocations.h>

#include <barrelfish_kpi/startup_arm.h>
#include <aos/performance.h>
#include <aos/deferred.h>
#include <fs/fs.h>
#include <aos/waitset_chan.h>

struct bootinfo *bi;

coreid_t my_core_id;
struct platform_info platform_info;

// only valid on app core
struct aos_urpc urpc_to_bsp;

// only valid on app core
static struct aos_urpc_server urpc_to_app_server;

// only valid on bsp core
struct aos_urpc urpc_to_app;
struct waitset urpc_to_app_ws;


static errval_t ram_alloc_remote_core(struct capref *ret, size_t size, size_t alignment)
{
    errval_t err;

    err = do_aos_urpc(
        &urpc_to_bsp, RPC_MTYPE_GET_RAM_CAP,
        NULL_CAP, 0, size, alignment,
        ret, NULL, NULL, NULL
    );
    if (err_is_fail(err)) return err_push(err, LIB_ERR_RAM_ALLOC);

    return SYS_ERR_OK;
}

static errval_t
initialize_self_rpc(void) {
    errval_t err;
    // we establish an rpc channel to ourselves so init can use UMP
    // if we had plenty of time we could have done some nice abstraction over the
    // rpc interface to do this but here we are

    struct capref rpc_shared_frame;
    err = frame_alloc(&rpc_shared_frame, RPC_SHARED_SIZE, NULL);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_FRAME_ALLOC);
    }

    void *rpc_shared_memory;
    err = paging_map_frame(
        get_current_paging_state(), &rpc_shared_memory,
        RPC_SHARED_SIZE, rpc_shared_frame);
    if (err_is_fail(err)) {
        cap_delete(rpc_shared_frame);
        return err;
    }

    // server side
    struct lmp_chan *self_chan_server = malloc(sizeof(struct lmp_chan));
    if (self_chan_server == NULL) return LIB_ERR_MALLOC_FAIL;
    lmp_chan_init(self_chan_server);
    err = endpoint_create(DEFAULT_LMP_BUF_WORDS, &self_chan_server->local_cap, &self_chan_server->endpoint);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_ENDPOINT_CREATE);
    err = lmp_chan_alloc_recv_slot(self_chan_server);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_LMP_ALLOC_RECV_SLOT);

    // client side

    struct lmp_chan *self_chan_client = malloc(sizeof(struct lmp_chan));
    if (self_chan_client == NULL) return LIB_ERR_MALLOC_FAIL;
    lmp_chan_init(self_chan_client);
    err = endpoint_create(DEFAULT_LMP_BUF_WORDS, &self_chan_client->local_cap, &self_chan_client->endpoint);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_ENDPOINT_CREATE);
    err = lmp_chan_alloc_recv_slot(self_chan_server);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_LMP_ALLOC_RECV_SLOT);

    self_chan_server->remote_cap = self_chan_client->local_cap;
    self_chan_client->remote_cap = self_chan_server->local_cap;
    struct aos_rpc_server *rpc_server = malloc(sizeof(struct aos_rpc_server));
    if (rpc_server == NULL) return LIB_ERR_MALLOC_FAIL;
    rpc_server_init(rpc_server, self_chan_server, rpc_shared_memory);
    rpc_server_register_recv(rpc_server);

    // set init rpc
    struct aos_rpc *init_rpc = malloc(sizeof(struct aos_rpc));
    err = aos_rpc_init(init_rpc, self_chan_client, rpc_shared_memory);
    if (err_is_fail(err)){
        return err_push(err, ERR_NOTIMP);
    }
    /* set init RPC client in our program state */
    set_init_rpc(init_rpc);

    return SYS_ERR_OK;
}

static bool event_dispatching_active = false;
static int dispatch_on_waitset(void *arg) {
    struct waitset *ws = arg;
    while (event_dispatching_active) {
        errval_t err = event_dispatch(ws);
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "in event_dispatch");
            abort();
        }
    }
    return 0;
}

static void noop_callback(void *arg) {
}

__attribute__((__unused__))
static struct waitset_chanstate self_rpc_waitset_chan;
static errval_t init_fs(void) {
    errval_t err;

    // debug_printf("[init_fs]\n");

    // create a thread dispatching so RPC works
    event_dispatching_active = true;
    struct thread *init_fs_dispatcher = thread_create_varstack_alloc(dispatch_on_waitset, get_default_waitset(), THREADS_DEFAULT_STACK_BYTES);
    if (init_fs_dispatcher == NULL) return LIB_ERR_THREAD_CREATE;

    // initialize the filesystem
    err = filesystem_init();
    if (err_is_fail(err)) return err;

    // debug_printf("[init_fs] filesystem_init done\n");

    // wait for the dispatcher to finish
    event_dispatching_active = false;
    // give the dispatcher a skip
    waitset_chanstate_init(&self_rpc_waitset_chan, CHANTYPE_OTHER);
    err = waitset_chan_trigger_closure(
        get_default_waitset(),
        &self_rpc_waitset_chan,
        MKCLOSURE(noop_callback, NULL)
    );
    if (err_is_fail(err)) return err;
    err = thread_join(init_fs_dispatcher, NULL);
    if (err_is_fail(err)) return err;

    // debug_printf("[init_fs] done\n");

    return SYS_ERR_OK;
}

static int
bsp_main(int argc, char *argv[]) {
    errval_t err;

    // Grading
    grading_setup_bsp_init(argc, argv);

    // First argument contains the bootinfo location, if it's not set
    bi = (struct bootinfo*)strtol(argv[1], NULL, 10);
    assert(bi);

    err = initialize_ram_alloc();
    if(err_is_fail(err)){
        DEBUG_ERR(err, "initialize_ram_alloc");
    }

    // TODO: initialize mem allocator, vspace management here

    err = cap_retype(cap_selfep, cap_dispatcher, 0, ObjType_EndPointLMP, 0, 1);
    if (err_is_fail(err)) return err;

    // Grading
    grading_test_early();


    // TODO: Spawn system processes, boot second core etc. here

    // A URPC frame, to hold the cross-core communication channels that you
    // will implement in Section 7.16.
    struct capref urpc_frame;
    err = frame_alloc(&urpc_frame, MON_URPC_SIZE, NULL);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_ALLOC);

    struct frame_identity urpc_frame_id;
    err = frame_identify(urpc_frame, &urpc_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_CAP_IDENTIFY);

    // boot second core
    char *cpu_driver_name;
    switch (platform_info.platform) {
        case PI_PLATFORM_QEMU:
            cpu_driver_name = "cpu_a57_qemu";
            break;
        case PI_PLATFORM_IMX8X:
            cpu_driver_name = "cpu_imx8x";
            break;
        default:
            USER_PANIC("Platform not implemented");
    }

    void *urpc_addr;
    err = paging_map_frame(get_current_paging_state(), &urpc_addr, MON_URPC_SIZE, urpc_frame);
    if (err_is_fail(err)) return err;
    memset(urpc_addr, 0, MON_URPC_SIZE);

    // memory barrier
    __asm volatile (
        "dmb sy\n"
    );
    cpu_idcache_wbinv_range((uintptr_t)urpc_addr, MON_URPC_SIZE);

    err = coreboot(1, "boot_armv8_generic", cpu_driver_name, "init", urpc_frame_id);
    if (err_is_fail(err)) return err_push(err, MON_ERR_SPAWN_CORE);

    struct aos_urpc_frame *urpc = urpc_addr;

    struct aos_urpc_server *urpc_to_bsp_server = malloc(sizeof(struct aos_urpc_server));
    urpc_to_bsp_server->g.shared_mem = &urpc->shared_mem_to_bsp;
    urpc_to_bsp_server->g.async_reply = urpc_server_async_reply;
    urpc_to_bsp_server->meta = &urpc->meta_to_bsp;
    thread_create(urpc_server, urpc_to_bsp_server);

    urpc_to_app.shared_mem = &urpc->shared_mem_to_app;
    urpc_to_app.meta = &urpc->meta_to_app;

    waitset_init(&urpc_to_app_ws);
    thread_create(urpc_client_loop, &urpc_to_app_ws);

    // spawn the block driver server
    struct spawninfo block_driver_si;
    domainid_t block_driver_pid;
    err = spawn_load_by_name("block_driver_server", &block_driver_si, &block_driver_pid);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "when spawning block_driver_server");
    }

    //spawn shell
    struct spawninfo shelly_si;
    domainid_t shelly_pid;
    err = spawn_load_by_name("shelly", &shelly_si, &shelly_pid);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "when spawning shelly");
    }

    // Spawn enet driver
    struct spawninfo enet_si;
    domainid_t enet_pid;
    err = spawn_load_by_name("enet", &enet_si, &enet_pid);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "when spawning enet");
    }

    // requried to connect to UMP channels from init
    err = initialize_self_rpc();
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "in initialize_self_rpc");
        abort();
    }

    // initialize the file system so it is available for spawning processes
    err = init_fs();
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "when initializing the file system");
    }

    // Grading
    grading_test_late();

    DEBUG_PRINTF("Message handler loop\n");
    // Hang around
    struct waitset *default_ws = get_default_waitset();
    while (true) {
        err = event_dispatch(default_ws);
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "in event_dispatch");
            abort();
        }
    }

    return EXIT_SUCCESS;
}

static errval_t deserialize_bootinfo(struct bootinfo_serialized * serialized, struct bootinfo ** ret) {
    errval_t err;
    assert(serialized != NULL);

    size_t bootinfo_size = sizeof(struct bootinfo) + serialized->regions_length * sizeof(struct mem_region);
    err = cnode_create_foreign_l2(cap_root, ROOTCN_SLOT_MODULECN, NULL);

    // struct capref ret_frame;
    // err = frame_alloc(&ret_frame, bootinfo_size, NULL);
    // if (err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_ALLOC);

    // err = paging_map_frame(get_current_paging_state(), (void **) ret, bootinfo_size, ret_frame);
    // if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_MAP);
    *ret = malloc(bootinfo_size);
    if(*ret == NULL) return LIB_ERR_MALLOC_FAIL;

    (*ret)->host_msg = serialized->host_msg;
    (*ret)->host_msg_bits = serialized->host_msg_bits;
    (*ret)->mem_spawn_core = serialized->mem_spawn_core;
    (*ret)->regions_length = serialized->regions_length;

    //forge the slot 0 cap
    struct capref slot_0_forge_destination = {
        .cnode = cnode_module,
        .slot = 0,
    };
    err = devframe_forge(slot_0_forge_destination, serialized->slot_0_frame_base, serialized->slot_0_frame_bytes, my_core_id);
    if (err_is_fail(err)) return err;

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

        (*ret)->regions[i].mrmod_size = serialized->regions[i].mrmod_size;
        (*ret)->regions[i].mrmod_data = serialized->regions[i].mrmod_data;
        (*ret)->regions[i].mrmod_slot = serialized->regions[i].mrmod_slot;

        if (serialized->regions[i].mr_type != RegionType_Module) {
            continue;
        }

        //forge cap
        struct capref forge_destination = {
            .cnode = cnode_module,
            .slot = serialized->regions[i].mrmod_slot,
        };
        err = devframe_forge(forge_destination, serialized->regions[i].mrmod_frame_base, serialized->regions[i].mrmod_frame_bytes, my_core_id);
        if (err_is_fail(err)) return err;
    }

    return SYS_ERR_OK;
}


static int
app_main(int argc, char *argv[]) {
    errval_t err;

    DEBUG_PRINTF("[app_main]\n");

    // Implement me in Milestone 5
    // Remember to call
    // - grading_setup_app_init(..);
    // - grading_test_early();
    // - grading_test_late();


    // TODO: initialize mem allocator, vspace management here
    err = cap_retype(cap_selfep, cap_dispatcher, 0, ObjType_EndPointLMP, 0, 1);
    if (err_is_fail(err)) return err;

    struct aos_urpc_frame *urpc = (void*)MON_URPC_VBASE;

    urpc_to_bsp.shared_mem = &urpc->shared_mem_to_bsp;
    urpc_to_bsp.meta = &urpc->meta_to_bsp;

    urpc_to_app_server.g.shared_mem = &urpc->shared_mem_to_app;
    urpc_to_app_server.g.async_reply = urpc_server_async_reply;
    urpc_to_app_server.meta = &urpc->meta_to_app;

    ram_alloc_set(ram_alloc_remote_core);

#ifdef PERFORMANCE_PERFORMANCE
    // wait for other stuff to complete
    barrelfish_usleep(2000000);

    struct performance_context p;
    // measure performance measurement performance
    for(size_t i = 0; i < 100; ++i) {
        perf_init(&p, "aos_performance");
        perf_add_now(&p, "start");
        perf_add_now(&p, "done");
        perf_print(&p);
        barrelfish_usleep(10000);
    }
    abort();
#endif

#ifdef URPC_PERFORMANCE
    // wait for other stuff to complete
    barrelfish_usleep(2000000);

    struct performance_context urpc_pcontext;
    // measure URPC performance
    for(size_t i = 0; i < 100; ++i) {
        perf_init(&urpc_pcontext, "aos_urpc_nop");
        perf_add_now(&urpc_pcontext, "start");
        do_aos_urpc(&urpc_to_bsp, RPC_MTYPE_NOP, NULL_CAP, 0, 0, 0, NULL, NULL, NULL, NULL);
        perf_add_now(&urpc_pcontext, "done");
        perf_print(&urpc_pcontext);
        // make sure we give the server some time to print measurements
        barrelfish_usleep(10000);
    }
    abort();
#endif

    // Allocate all pages of the thread stack now.
    // We can't have page faults while the thread is running,
    // because URPC calls may only be done from the main thread.
    thread_create_varstack_alloc(urpc_server, &urpc_to_app_server, THREADS_DEFAULT_STACK_BYTES);


    // TODO: Spawn system processes etc. here

    struct bootinfo_serialized * bi_ser;
    err = aos_urpc_get_bootinfo(&urpc_to_bsp, &bi_ser);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "in aos_urpc_get_bootinfo");
        abort();
    }
    err = deserialize_bootinfo(bi_ser, &bi);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "in deserialize_bootinfo");
        abort();
    }
    free(bi_ser);

    DEBUG_PRINTF("[app_main]: received bootinfo with %d regions\n", bi->regions_length);

    // Grading
    grading_setup_app_init(bi);

    // Grading
    grading_test_early();

    // Spawn echoserver
    struct spawninfo echoserver_si;
    domainid_t echoserver_pid;
    err = spawn_load_by_name("echoserver", &echoserver_si, &echoserver_pid);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "when spawning echoserver");
    }

    err = initialize_self_rpc();
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "in initialize_self_rpc");
        abort();
    }

    err = init_fs();
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "when initializing the file system");
    }

    // Grading
    grading_test_late();

    DEBUG_PRINTF("Message handler loop\n");
    // Hang around
    struct waitset *default_ws = get_default_waitset();
    while (true) {
        err = event_dispatch(default_ws);
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "in event_dispatch");
            abort();
        }
    }

    return EXIT_SUCCESS;
}

int main(int argc, char *argv[])
{
    errval_t err;

    /* obtain the core information from the kernel*/
    err = invoke_kernel_get_core_id(cap_kernel, &my_core_id);
    if (err_is_fail(err)) {
        USER_PANIC_ERR(err, "failed to obtain the core id from the kernel\n");
    }

    /* Set the core id in the disp_priv struct */
    disp_set_core_id(my_core_id);

    /* obtain the platform information */
    err = invoke_kernel_get_platform_info(cap_kernel, &platform_info);
    if (err_is_fail(err)) {
        USER_PANIC_ERR(err, "failed to obtain the platform info from the kernel\n");
    }

    char *platform;
    switch (platform_info.platform) {
        case PI_PLATFORM_QEMU:
            platform = "QEMU";
            break;
        case PI_PLATFORM_IMX8X:
            platform = "IMX8X";
            break;
        default:
            platform = "UNKNOWN";
    }

    DEBUG_PRINTF("init domain starting on core %" PRIuCOREID " (%s), invoked as:", my_core_id, platform);
    for (int i = 0; i < argc; i++) {
       printf(" %s", argv[i]);
    }
    printf("\n");



    fflush(stdout);


    if(my_core_id == 0) return bsp_main(argc, argv);
    else                return app_main(argc, argv);
}