#include <aos/aos.h>
#include <aos/coreboot.h>
#include <spawn/multiboot.h>
#include <elf/elf.h>
#include <string.h>
#include <barrelfish_kpi/arm_core_data.h>
#include <aos/kernel_cap_invocations.h>
#include <aos/cache.h>

#define ARMv8_KERNEL_OFFSET 0xffff000000000000

extern struct bootinfo *bi;

struct mem_info {
    size_t                size;      // Size in bytes of the memory region
    void                  *buf;      // Address where the region is currently mapped
    lpaddr_t              phys_base; // Physical base address   
};

/**
 * Load a ELF image into memory.
 *
 * binary:            Valid pointer to ELF image in current address space
 * mem:               Where the ELF will be loaded
 * entry_point:       Virtual address of the entry point
 * reloc_entry_point: Return the loaded, physical address of the entry_point
 */
__attribute__((__used__))
static errval_t load_elf_binary(genvaddr_t binary, const struct mem_info *mem,
                         genvaddr_t entry_point, genvaddr_t *reloc_entry_point)

{

    struct Elf64_Ehdr *ehdr = (struct Elf64_Ehdr *)binary;

    /* Load the CPU driver from its ELF image. */
    bool found_entry_point= 0;
    bool loaded = 0;

    struct Elf64_Phdr *phdr = (struct Elf64_Phdr *)(binary + ehdr->e_phoff);
    for(size_t i= 0; i < ehdr->e_phnum; i++) {
        if(phdr[i].p_type != PT_LOAD) {
            DEBUG_PRINTF("Segment %d load address 0x% "PRIx64 ", file size %" PRIu64
                  ", memory size 0x%" PRIx64 " SKIP\n", i, phdr[i].p_vaddr,
                  phdr[i].p_filesz, phdr[i].p_memsz);
            continue;
        }

        DEBUG_PRINTF("Segment %d load address 0x% "PRIx64 ", file size %" PRIu64
              ", memory size 0x%" PRIx64 " LOAD\n", i, phdr[i].p_vaddr,
              phdr[i].p_filesz, phdr[i].p_memsz);


        if (loaded) {
            USER_PANIC("Expected one load able segment!\n");
        }
        loaded = 1;

        void *dest = mem->buf;
        lpaddr_t dest_phys = mem->phys_base;

        assert(phdr[i].p_offset + phdr[i].p_memsz <= mem->size);

        /* copy loadable part */
        memcpy(dest, (void *)(binary + phdr[i].p_offset), phdr[i].p_filesz);

        /* zero out BSS section */
        memset(dest + phdr[i].p_filesz, 0, phdr[i].p_memsz - phdr[i].p_filesz);

        if (!found_entry_point) {
            if(entry_point >= phdr[i].p_vaddr
                 && entry_point - phdr[i].p_vaddr < phdr[i].p_memsz) {
               *reloc_entry_point= (dest_phys + (entry_point - phdr[i].p_vaddr));
               found_entry_point= 1;
            }
        }
    }

    if (!found_entry_point) {
        USER_PANIC("No entry point loaded\n");
    }

    return SYS_ERR_OK;
}

/**
 * Relocate an already loaded ELF image. 
 *
 * binary:            Valid pointer to ELF image in current address space
 * mem:               Where the ELF is loaded
 * kernel_:       Virtual address of the entry point
 * reloc_entry_point: Return the loaded, physical address of the entry_point
 */
__attribute__((__used__))
static errval_t
relocate_elf(genvaddr_t binary, struct mem_info *mem, lvaddr_t load_offset)
{
    DEBUG_PRINTF("Relocating image.\n");

    struct Elf64_Ehdr *ehdr = (struct Elf64_Ehdr *)binary;

    size_t shnum  = ehdr->e_shnum;
    struct Elf64_Phdr *phdr = (struct Elf64_Phdr *)(binary + ehdr->e_phoff);
    struct Elf64_Shdr *shead = (struct Elf64_Shdr *)(binary + (uintptr_t)ehdr->e_shoff);

    /* Search for relocaton sections. */
    for(size_t i= 0; i < shnum; i++) {

        struct Elf64_Shdr *shdr=  &shead[i];
        if(shdr->sh_type == SHT_REL || shdr->sh_type == SHT_RELA) {
            if(shdr->sh_info != 0) {
                DEBUG_PRINTF("I expected global relocations, but got"
                              " section-specific ones.\n");
                return ELF_ERR_HEADER;
            }


            uint64_t segment_elf_base= phdr[0].p_vaddr;
            uint64_t segment_load_base=mem->phys_base;
            uint64_t segment_delta= segment_load_base - segment_elf_base;
            uint64_t segment_vdelta= (uintptr_t)mem->buf - segment_elf_base;

            size_t rsize;
            if(shdr->sh_type == SHT_REL){
                rsize= sizeof(struct Elf64_Rel);
            } else {
                rsize= sizeof(struct Elf64_Rela);
            }

            assert(rsize == shdr->sh_entsize);
            size_t nrel= shdr->sh_size / rsize;

            void * reldata = (void*)(binary + shdr->sh_offset);

            /* Iterate through the relocations. */
            for(size_t ii= 0; ii < nrel; ii++) {
                void *reladdr= reldata + ii *rsize;

                switch(shdr->sh_type) {
                    case SHT_REL:
                        DEBUG_PRINTF("SHT_REL unimplemented.\n");
                        return ELF_ERR_PROGHDR;
                    case SHT_RELA:
                    {
                        struct Elf64_Rela *rel= reladdr;

                        uint64_t offset= rel->r_offset;
                        uint64_t sym= ELF64_R_SYM(rel->r_info);
                        uint64_t type= ELF64_R_TYPE(rel->r_info);
                        uint64_t addend= rel->r_addend;

                        uint64_t *rel_target= (void *)offset + segment_vdelta;

                        switch(type) {
                            case R_AARCH64_RELATIVE:
                                if(sym != 0) {
                                    DEBUG_PRINTF("Relocation references a"
                                                 " dynamic symbol, which is"
                                                 " unsupported.\n");
                                    return ELF_ERR_PROGHDR;
                                }

                                /* Delta(S) + A */
                                *rel_target= addend + segment_delta + load_offset;
                                break;

                            default:
                                DEBUG_PRINTF("Unsupported relocation type %d\n",
                                             type);
                                return ELF_ERR_PROGHDR;
                        }
                    }
                    break;
                    default:
                        DEBUG_PRINTF("Unexpected type\n");
                        break;

                }
            }
        }
    }

    return SYS_ERR_OK;
}



errval_t coreboot(coreid_t mpid,
        const char *boot_driver,
        const char *cpu_driver,
        const char *init,
        struct frame_identity urpc_frame_id)
{
    errval_t err;

    // Implement me!
    // - Get a new KCB by retyping a RAM cap to ObjType_KernelControlBlock.
    //   Note that it should at least OBJSIZE_KCB, and it should also be aligned 
    //   to a multiple of 16k.
    struct capref kcb_ram_cap;
    err = ram_alloc_aligned(&kcb_ram_cap, OBJSIZE_KCB, PAGE_SIZE_16K);
    if (err_is_fail(err)) return err;

    struct capref kcb_cap;
    err = slot_alloc(&kcb_cap);
    if (err_is_fail(err)) return err;
    err = cap_retype(kcb_cap, kcb_ram_cap, 0, ObjType_KernelControlBlock, OBJSIZE_KCB, 1);
    if (err_is_fail(err)) return err;

    struct frame_identity kcb_id;
    err = invoke_kcb_identify(kcb_cap, &kcb_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    // - Get and load the CPU and boot driver binary.
    struct mem_region * bootdriver_module = multiboot_find_module(bi, boot_driver);
    if (bootdriver_module == NULL) return LIB_ERR_COREBOOT_FIND_MODULE;

    struct mem_region * cpudriver_module = multiboot_find_module(bi, cpu_driver);
    if (cpudriver_module == NULL) return LIB_ERR_COREBOOT_FIND_MODULE;

    struct mem_region * init_module = multiboot_find_module(bi, init);
    if (init_module == NULL) return LIB_ERR_COREBOOT_FIND_MODULE;

    struct capref init_elf_frame = {
        .cnode = cnode_module,
        .slot = init_module->mrmod_slot,
    };
    struct frame_identity init_elf_frame_id;
    err = frame_identify(init_elf_frame, &init_elf_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    // - Load the ELF binary
    struct capref bootdriver_elf_frame = {
        .cnode = cnode_module,
        .slot = bootdriver_module->mrmod_slot,
    };
    struct frame_identity bootdriver_elf_frame_id;
    err = frame_identify(bootdriver_elf_frame, &bootdriver_elf_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    struct capref remote_bootdriver_elf_frame;
    err = frame_alloc(&remote_bootdriver_elf_frame, bootdriver_elf_frame_id.bytes, NULL);
    if(err_is_fail(err)) return err;

    struct frame_identity remote_bootdriver_elf_frame_id;
    err = frame_identify(remote_bootdriver_elf_frame, &remote_bootdriver_elf_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    char * bootdriver_elf_addr;
    err = paging_map_frame_attr(get_current_paging_state(), (void **) &bootdriver_elf_addr, bootdriver_elf_frame_id.bytes, bootdriver_elf_frame, VREGION_FLAGS_READ);
    if(err_is_fail(err)) return err;

    char * remote_bootdriver_elf_addr;
    err = paging_map_frame(get_current_paging_state(), (void **) &remote_bootdriver_elf_addr, bootdriver_elf_frame_id.bytes, remote_bootdriver_elf_frame);
    if(err_is_fail(err)) return err;

    struct capref cpudriver_elf_frame = {
        .cnode = cnode_module,
        .slot = cpudriver_module->mrmod_slot,
    };
    struct frame_identity cpudriver_elf_frame_id;
    err = frame_identify(cpudriver_elf_frame, &cpudriver_elf_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    struct capref remote_cpudriver_elf_frame;
    err = frame_alloc(&remote_cpudriver_elf_frame, cpudriver_elf_frame_id.bytes, NULL);
    if(err_is_fail(err)) return err;

    struct frame_identity remote_cpudriver_elf_frame_id;
    err = frame_identify(remote_cpudriver_elf_frame, &remote_cpudriver_elf_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    char * cpudriver_elf_addr;
    err = paging_map_frame_attr(get_current_paging_state(), (void **) &cpudriver_elf_addr, cpudriver_elf_frame_id.bytes, cpudriver_elf_frame, VREGION_FLAGS_READ);
    if(err_is_fail(err)) return err;

    char * remote_cpudriver_elf_addr;
    err = paging_map_frame(get_current_paging_state(), (void **) &remote_cpudriver_elf_addr, cpudriver_elf_frame_id.bytes, remote_cpudriver_elf_frame);
    if(err_is_fail(err)) return err;

    // - Find the boot driver entry point. Look for the symbol "boot_entry_psci"
    struct mem_info bootdriver_mem_info = {
        .size = remote_bootdriver_elf_frame_id.bytes,
        .buf = remote_bootdriver_elf_addr,
        .phys_base = remote_bootdriver_elf_frame_id.base,
    };

    struct Elf64_Sym * boot_entry_psci_sym = elf64_find_symbol_by_name((genvaddr_t) bootdriver_elf_addr, bootdriver_elf_frame_id.bytes, "boot_entry_psci", false, STT_FUNC, 0);
    assert(boot_entry_psci_sym != NULL);

    genvaddr_t reloc_bootdriver_entry;
    err = load_elf_binary((genvaddr_t) bootdriver_elf_addr, &bootdriver_mem_info, boot_entry_psci_sym->st_value, &reloc_bootdriver_entry);
    if(err_is_fail(err)) return err;

    // - Find the CPU driver entry point. Look for the symbol "arch_init". Put
    //   the address in the core data struct. 
    struct mem_info cpudriver_mem_info = {
        .size = remote_cpudriver_elf_frame_id.bytes,
        .buf = remote_cpudriver_elf_addr,
        .phys_base = remote_cpudriver_elf_frame_id.base,
    };

    struct Elf64_Sym * arch_init_sym = elf64_find_symbol_by_name((genvaddr_t) cpudriver_elf_addr, cpudriver_elf_frame_id.bytes, "arch_init", false, STT_FUNC, 0);
    assert(arch_init_sym != NULL);

    genvaddr_t reloc_cpudriver_entry;
    err = load_elf_binary((genvaddr_t) cpudriver_elf_addr, &cpudriver_mem_info, arch_init_sym->st_value, &reloc_cpudriver_entry);
    if(err_is_fail(err)) return err;


    // - Relocate the boot and CPU driver. The boot driver runs with a 1:1
    //   VA->PA mapping. The CPU driver is expected to be loaded at the 
    //   high virtual address space, at offset ARMV8_KERNEL_OFFSET.
    err = relocate_elf((genvaddr_t) bootdriver_elf_addr, &bootdriver_mem_info, 0);
    if(err_is_fail(err)) return err;

    err = relocate_elf((genvaddr_t) cpudriver_elf_addr, &cpudriver_mem_info, ARMv8_KERNEL_OFFSET);
    if(err_is_fail(err)) return err;


    // - Allocate a page for the core data struct
    struct capref core_data_frame;
    err = frame_alloc(&core_data_frame, BASE_PAGE_SIZE, NULL);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_ALLOC);
    struct frame_identity core_data_frame_id;
    err = frame_identify(core_data_frame, &core_data_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    // - Allocate stack memory for the new cpu driver (at least 16 pages)
    struct capref cpu_driver_stack_frame;
    err = frame_alloc(&cpu_driver_stack_frame, 16 * BASE_PAGE_SIZE, NULL);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_ALLOC);
    struct frame_identity cpu_driver_stack_frame_id;
    err = frame_identify(cpu_driver_stack_frame, &cpu_driver_stack_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    // Space to load the init process. This should be at least space for init
    // plus ARMV8_CORE_DATA_PAGES × BASE_PAGE_SIZE bytes.
    struct capref init_process_frame;
    err = frame_alloc(&init_process_frame, init_module->mrmod_size + ARMV8_CORE_DATA_PAGES * BASE_PAGE_SIZE, NULL);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_ALLOC);
    struct frame_identity init_process_frame_id;
    err = frame_identify(init_process_frame, &init_process_frame_id);
    if(err_is_fail(err)) return err_push(err, LIB_ERR_FRAME_IDENTIFY);

    // - Fill in the core data struct, for a description, see the definition
    //   in include/target/aarch64/barrelfish_kpi/arm_core_data.h
    struct armv8_core_data *core_data;
    err = paging_map_frame(get_current_paging_state(), (void**) &core_data, PAGE_SIZE, core_data_frame);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_MAP);
    core_data->boot_magic = ARMV8_BOOTMAGIC_PSCI;

    core_data->cpu_driver_stack = cpu_driver_stack_frame_id.base + 16 * BASE_PAGE_SIZE;
    core_data->cpu_driver_stack_limit = cpu_driver_stack_frame_id.base;

    strncpy(core_data->cpu_driver_cmdline, multiboot_module_opts(cpudriver_module), 128);

    core_data->memory.base = init_process_frame_id.base;
    core_data->memory.length = init_process_frame_id.bytes;

    core_data->urpc_frame.base = urpc_frame_id.base;
    core_data->urpc_frame.length = urpc_frame_id.bytes;

    core_data->monitor_binary.base = init_elf_frame_id.base;
    core_data->monitor_binary.length = init_elf_frame_id.bytes;

    core_data->kcb = kcb_id.base;

    core_data->src_core_id = disp_get_core_id();
    core_data->dst_core_id = mpid;
    core_data->src_arch_id = disp_get_core_id();
    core_data->dst_arch_id = mpid;

    // find cpu driver entry point
    // debug_printf("[coreboot]: arch_init_sym->st_value = %p\n", arch_init_sym->st_value);
    // debug_printf("[coreboot]: reloc_cpudriver_entry = %p\n", reloc_cpudriver_entry);
    // debug_printf("[coreboot]: ARMv8_KERNEL_OFFSET = %p\n", ARMv8_KERNEL_OFFSET);
    core_data->cpu_driver_entry = reloc_cpudriver_entry + ARMv8_KERNEL_OFFSET;

    debug_printf("[coreboot] Flush the cache\n");

    // - Flush the cache.
    //memory barrier
    __asm volatile (
        "dmb sy\n"
    );
    //clean and invalidate cache
    // cpu_idcache_wbinv_range(0, VADDR_SIZE);
    cpu_idcache_wbinv_range((genvaddr_t)remote_bootdriver_elf_addr, remote_bootdriver_elf_frame_id.bytes);
    cpu_idcache_wbinv_range((genvaddr_t)remote_cpudriver_elf_addr, remote_cpudriver_elf_frame_id.bytes);
    cpu_idcache_wbinv_range((genvaddr_t)core_data, PAGE_SIZE);

    // - Call the invoke_monitor_spawn_core with the entry point 
    //   of the boot driver and pass the (physical, of course) address of the 
    //   boot struct as argument
    err = invoke_monitor_spawn_core(mpid, CPU_ARM8, reloc_bootdriver_entry, core_data_frame_id.base, 0 /*ignored in aos*/);
    if (err_is_fail(err)) return err_push(err, MON_ERR_SPAWN_CORE);

    // unmap stuff
    paging_unmap(get_current_paging_state(), core_data);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_UNMAP);
    paging_unmap(get_current_paging_state(), bootdriver_elf_addr);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_UNMAP);
    paging_unmap(get_current_paging_state(), remote_bootdriver_elf_addr);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_UNMAP);
    paging_unmap(get_current_paging_state(), cpudriver_elf_addr);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_UNMAP);
    paging_unmap(get_current_paging_state(), remote_cpudriver_elf_addr);
    if (err_is_fail(err)) return err_push(err, LIB_ERR_PMAP_UNMAP);

    return SYS_ERR_OK;

}
