#include "hv/arch.h"
#include "hv/config.h"
#include "hv/fdt.h"
#include "hv/guest_loader.h"
#include "hv/log.h"
#include "hv/string.h"

#define ARM64_IMAGE_MAGIC 0x644d5241u
#define ARM64_IMAGE_MAGIC_OFFSET 0x38u
#define ARM64_IMAGE_TEXT_OFFSET 0x08u
#define ARM64_IMAGE_SIZE_OFFSET 0x10u
#define LINUX_BOOTARGS "console=ttyAMA0 earlycon"

static uint32_t read_le32(const uint8_t *p)
{
    return (uint32_t)p[0] | ((uint32_t)p[1] << 8u) |
           ((uint32_t)p[2] << 16u) | ((uint32_t)p[3] << 24u);
}

static uint64_t read_le64(const uint8_t *p)
{
    return (uint64_t)read_le32(p) | ((uint64_t)read_le32(p + 4u) << 32u);
}

bool guest_load_diagnostic(struct guest_image_info *info)
{
    uintptr_t guest_start;
    uintptr_t guest_end;
    uint64_t guest_size;

    if (info == (struct guest_image_info *)0) {
        return false;
    }

    guest_start = (uintptr_t)__guest_image_start;
    guest_end = (uintptr_t)__guest_image_end;
    guest_size = (uint64_t)(guest_end - guest_start);

    if (guest_size == 0u || guest_size > CONFIG_GUEST_RAM_SIZE) {
        return false;
    }
    if (hv_range_overlaps(CONFIG_GUEST_RAM_BASE, CONFIG_GUEST_RAM_SIZE,
                          (uint64_t)(uintptr_t)__image_start,
                          (uint64_t)((uintptr_t)__image_end - (uintptr_t)__image_start))) {
        return false;
    }

    hv_memset((void *)(uintptr_t)CONFIG_GUEST_RAM_BASE, 0, CONFIG_GUEST_RAM_SIZE);
    hv_memcpy((void *)(uintptr_t)CONFIG_GUEST_RAM_BASE, (const void *)guest_start, (size_t)guest_size);

    info->kind = GUEST_PAYLOAD_DIAGNOSTIC;
    info->ipa_entry = CONFIG_GUEST_ENTRY;
    info->ipa_load = CONFIG_GUEST_IPA_BASE;
    info->size = guest_size;
    info->dtb_ipa = 0u;
    info->initrd_ipa = 0u;
    info->initrd_size = 0u;
    return true;
}

bool guest_validate_linux_image_header(const void *image, uint64_t size,
                                       struct guest_image_info *info)
{
    const uint8_t *p = (const uint8_t *)image;
    uint64_t text_offset;
    uint64_t image_size;
    uint64_t entry;

    if (image == (const void *)0 || info == (struct guest_image_info *)0 || size < 0x40u) {
        return false;
    }
    if (read_le32(p + ARM64_IMAGE_MAGIC_OFFSET) != ARM64_IMAGE_MAGIC) {
        return false;
    }

    text_offset = read_le64(p + ARM64_IMAGE_TEXT_OFFSET);
    image_size = read_le64(p + ARM64_IMAGE_SIZE_OFFSET);
    if (image_size == 0u) {
        image_size = size;
    }
    if (image_size > CONFIG_GUEST_RAM_SIZE || text_offset >= CONFIG_GUEST_RAM_SIZE) {
        return false;
    }
    if (hv_add_overflow_u64(CONFIG_GUEST_IPA_BASE, text_offset, &entry)) {
        return false;
    }

    info->kind = GUEST_PAYLOAD_LINUX_IMAGE;
    info->ipa_entry = entry;
    info->ipa_load = entry;
    info->size = image_size;
    info->dtb_ipa = 0u;
    info->initrd_ipa = 0u;
    info->initrd_size = 0u;
    return true;
}

bool guest_load_external_linux(struct guest_image_info *info, paddr_t image_pa,
                               uint64_t image_max_size, paddr_t dtb_pa,
                               uint64_t dtb_max_size, paddr_t initrd_pa,
                               uint64_t initrd_size)
{
    struct fdt_platform_info dtb_info;
    uint64_t initrd_start = 0u;
    uint64_t initrd_end = 0u;

    if (info == (struct guest_image_info *)0 ||
        image_max_size == 0u || image_max_size > CONFIG_GUEST_RAM_SIZE) {
        return false;
    }
    if (!guest_validate_linux_image_header((const void *)(uintptr_t)image_pa,
                                           image_max_size, info)) {
        return false;
    }
    if (!fdt_probe(dtb_pa, &dtb_info) || dtb_info.size > dtb_max_size) {
        return false;
    }
    if (initrd_size != 0u) {
        if (initrd_size > CONFIG_LINUX_INITRD_MAX_SIZE ||
            initrd_pa != CONFIG_LINUX_INITRD_PA ||
            hv_add_overflow_u64(CONFIG_LINUX_INITRD_IPA, initrd_size, &initrd_end) ||
            initrd_end > CONFIG_LINUX_DTB_IPA) {
            return false;
        }
        initrd_start = CONFIG_LINUX_INITRD_IPA;
    }
    if (!fdt_patch_memory(dtb_pa, CONFIG_GUEST_IPA_BASE, CONFIG_GUEST_RAM_SIZE)) {
        return false;
    }
    if (!fdt_patch_linux_guest(dtb_pa, dtb_max_size, LINUX_BOOTARGS,
                               initrd_start, initrd_end)) {
        return false;
    }
    if (!fdt_probe(dtb_pa, &dtb_info)) {
        return false;
    }
    log_simple(LOG_BOOT, 33u, dtb_info.size, dtb_info.memory_size);
    if (dtb_info.initrd_start != 0u || dtb_info.initrd_end != 0u) {
        log_simple(LOG_BOOT, 34u, dtb_info.initrd_start, dtb_info.initrd_end);
    }
    fdt_dump(&dtb_info);

    info->kind = GUEST_PAYLOAD_LINUX_IMAGE;
    info->ipa_load = CONFIG_GUEST_IPA_BASE;
    info->dtb_ipa = CONFIG_LINUX_DTB_IPA;
    if (initrd_start != 0u && initrd_end > initrd_start) {
        info->initrd_ipa = initrd_start;
        info->initrd_size = initrd_end - initrd_start;
    }
    return true;
}