#include "hv/fdt.h"
#include "hv/string.h"
#include "hv/uart.h"

#define FDT_MAGIC 0xd00dfeedu
#define FDT_BEGIN_NODE 1u
#define FDT_END_NODE 2u
#define FDT_PROP 3u
#define FDT_NOP 4u
#define FDT_END 9u
#define FDT_MAX_SIZE 0x00200000u
#define FDT_MAX_DEPTH 16u

struct fdt_header_view {
    uint32_t totalsize;
    uint32_t off_dt_struct;
    uint32_t off_dt_strings;
    uint32_t size_dt_strings;
    uint32_t size_dt_struct;
};

struct fdt_property_ref {
    uint8_t *value;
    uint32_t len;
};

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

static uint64_t read_cells(const uint32_t *cells, uint32_t count)
{
    uint64_t value = 0;
    for (uint32_t i = 0; i < count; i++) {
        value = (value << 32u) | (uint64_t)be32(&cells[i]);
    }
    return value;
}

static uint32_t align4_u32(uint32_t value)
{
    return (value + 3u) & ~3u;
}

static const uint8_t *align4(const uint8_t *p)
{
    uintptr_t v = (uintptr_t)p;
    v = (v + 3u) & ~(uintptr_t)3u;
    return (const uint8_t *)v;
}

static bool read_header(const uint8_t *fdt, struct fdt_header_view *h)
{
    if (fdt == (const uint8_t *)0 || h == (struct fdt_header_view *)0) {
        return false;
    }
    if (be32(fdt) != FDT_MAGIC) {
        return false;
    }

    h->totalsize = be32(fdt + 4u);
    h->off_dt_struct = be32(fdt + 8u);
    h->off_dt_strings = be32(fdt + 12u);
    h->size_dt_strings = be32(fdt + 32u);
    h->size_dt_struct = be32(fdt + 36u);
    if (h->totalsize < 40u || h->totalsize > FDT_MAX_SIZE) {
        return false;
    }
    if (h->off_dt_struct >= h->totalsize || h->off_dt_strings >= h->totalsize) {
        return false;
    }
    if (h->size_dt_struct > h->totalsize - h->off_dt_struct ||
        h->size_dt_strings > h->totalsize - h->off_dt_strings) {
        return false;
    }
    return true;
}

static bool string_in_block(const char *s, const uint8_t *start, const uint8_t *end)
{
    const uint8_t *p = (const uint8_t *)s;
    if (p < start || p >= end) {
        return false;
    }
    while (p < end) {
        if (*p == 0u) {
            return true;
        }
        p++;
    }
    return false;
}

static bool node_name_is_memory(const char *name)
{
    return hv_strncmp(name, "memory", 6u) == 0;
}

static bool find_property(uint64_t fdt_pa, const char *node_name,
                          const char *prop_name, struct fdt_property_ref *ref)
{
    uint8_t *fdt = (uint8_t *)(uintptr_t)fdt_pa;
    struct fdt_header_view h;
    uint8_t *struct_p;
    uint8_t *struct_end;
    const uint8_t *strings;
    const uint8_t *strings_end;
    bool wanted_node[FDT_MAX_DEPTH];
    uint32_t depth = 0;

    if (node_name == (const char *)0 || prop_name == (const char *)0 ||
        ref == (struct fdt_property_ref *)0 || !read_header(fdt, &h)) {
        return false;
    }
    hv_memset(wanted_node, 0, sizeof(wanted_node));
    struct_p = fdt + h.off_dt_struct;
    struct_end = struct_p + h.size_dt_struct;
    strings = fdt + h.off_dt_strings;
    strings_end = strings + h.size_dt_strings;

    while (struct_p + 4u <= struct_end) {
        uint32_t token = be32(struct_p);
        struct_p += 4u;
        if (token == FDT_BEGIN_NODE) {
            const char *name = (const char *)struct_p;
            if (!string_in_block(name, struct_p, struct_end) || depth + 1u >= FDT_MAX_DEPTH) {
                return false;
            }
            wanted_node[depth] = hv_strcmp(name, node_name) == 0;
            while (struct_p < struct_end && *struct_p != 0u) {
                struct_p++;
            }
            if (struct_p >= struct_end) {
                return false;
            }
            struct_p = (uint8_t *)align4(struct_p + 1u);
            depth++;
        } else if (token == FDT_END_NODE) {
            if (depth == 0u) {
                return false;
            }
            depth--;
            wanted_node[depth] = false;
        } else if (token == FDT_PROP) {
            uint32_t len;
            uint32_t nameoff;
            const char *name;
            if (struct_p + 8u > struct_end) {
                return false;
            }
            len = be32(struct_p);
            nameoff = be32(struct_p + 4u);
            struct_p += 8u;
            if (nameoff >= h.size_dt_strings || len > (uint32_t)(struct_end - struct_p)) {
                return false;
            }
            name = (const char *)(strings + nameoff);
            if (!string_in_block(name, strings, strings_end)) {
                return false;
            }
            if (depth != 0u && wanted_node[depth - 1u] && hv_strcmp(name, prop_name) == 0) {
                ref->value = struct_p;
                ref->len = len;
                return true;
            }
            struct_p = (uint8_t *)align4(struct_p + len);
        } else if (token == FDT_NOP) {
            continue;
        } else if (token == FDT_END) {
            return false;
        } else {
            return false;
        }
    }
    return false;
}

static bool find_memory_property(uint64_t fdt_pa, const char *prop_name,
                                 struct fdt_property_ref *ref)
{
    uint8_t *fdt = (uint8_t *)(uintptr_t)fdt_pa;
    struct fdt_header_view h;
    uint8_t *struct_p;
    uint8_t *struct_end;
    const uint8_t *strings;
    const uint8_t *strings_end;
    bool memory_node[FDT_MAX_DEPTH];
    uint32_t depth = 0;

    if (prop_name == (const char *)0 || ref == (struct fdt_property_ref *)0 ||
        !read_header(fdt, &h)) {
        return false;
    }
    hv_memset(memory_node, 0, sizeof(memory_node));
    struct_p = fdt + h.off_dt_struct;
    struct_end = struct_p + h.size_dt_struct;
    strings = fdt + h.off_dt_strings;
    strings_end = strings + h.size_dt_strings;

    while (struct_p + 4u <= struct_end) {
        uint32_t token = be32(struct_p);
        struct_p += 4u;
        if (token == FDT_BEGIN_NODE) {
            const char *name = (const char *)struct_p;
            if (!string_in_block(name, struct_p, struct_end) || depth + 1u >= FDT_MAX_DEPTH) {
                return false;
            }
            memory_node[depth] = node_name_is_memory(name);
            while (struct_p < struct_end && *struct_p != 0u) {
                struct_p++;
            }
            if (struct_p >= struct_end) {
                return false;
            }
            struct_p = (uint8_t *)align4(struct_p + 1u);
            depth++;
        } else if (token == FDT_END_NODE) {
            if (depth == 0u) {
                return false;
            }
            depth--;
            memory_node[depth] = false;
        } else if (token == FDT_PROP) {
            uint32_t len;
            uint32_t nameoff;
            const char *name;
            if (struct_p + 8u > struct_end) {
                return false;
            }
            len = be32(struct_p);
            nameoff = be32(struct_p + 4u);
            struct_p += 8u;
            if (nameoff >= h.size_dt_strings || len > (uint32_t)(struct_end - struct_p)) {
                return false;
            }
            name = (const char *)(strings + nameoff);
            if (!string_in_block(name, strings, strings_end)) {
                return false;
            }
            if (depth != 0u && memory_node[depth - 1u] &&
                hv_strcmp(name, prop_name) == 0) {
                ref->value = struct_p;
                ref->len = len;
                return true;
            }
            struct_p = (uint8_t *)align4(struct_p + len);
        } else if (token == FDT_NOP) {
            continue;
        } else if (token == FDT_END) {
            return false;
        } else {
            return false;
        }
    }
    return false;
}

static void write_be64(uint8_t *p, uint64_t value)
{
    for (uint32_t i = 0; i < 8u; i++) {
        p[i] = (uint8_t)(value >> (56u - (i * 8u)));
    }
}

static void write_be32(uint8_t *p, uint32_t value)
{
    p[0] = (uint8_t)(value >> 24u);
    p[1] = (uint8_t)(value >> 16u);
    p[2] = (uint8_t)(value >> 8u);
    p[3] = (uint8_t)value;
}

static bool patch_string_prop(uint64_t fdt_pa, const char *node, const char *prop,
                              const char *value)
{
    struct fdt_property_ref ref;
    size_t len;
    if (value == (const char *)0 || !find_property(fdt_pa, node, prop, &ref)) {
        return false;
    }
    len = hv_strlen(value) + 1u;
    if (len > ref.len) {
        return false;
    }
    hv_memset(ref.value, 0, ref.len);
    hv_memcpy(ref.value, value, len);
    return true;
}

static bool repurpose_string_prop(uint64_t fdt_pa, const char *node_name,
                                  const char *prop_name, const char *value)
{
    uint8_t *fdt = (uint8_t *)(uintptr_t)fdt_pa;
    struct fdt_header_view h;
    uint8_t *struct_p;
    uint8_t *struct_end;
    uint8_t *strings;
    const uint8_t *strings_end;
    bool wanted_node[FDT_MAX_DEPTH];
    uint32_t depth = 0;
    size_t value_len;
    size_t prop_len;

    if (node_name == (const char *)0 || prop_name == (const char *)0 ||
        value == (const char *)0 || !read_header(fdt, &h)) {
        return false;
    }

    value_len = hv_strlen(value) + 1u;
    prop_len = hv_strlen(prop_name) + 1u;
    hv_memset(wanted_node, 0, sizeof(wanted_node));
    struct_p = fdt + h.off_dt_struct;
    struct_end = struct_p + h.size_dt_struct;
    strings = fdt + h.off_dt_strings;
    strings_end = strings + h.size_dt_strings;

    while (struct_p + 4u <= struct_end) {
        uint32_t token = be32(struct_p);
        struct_p += 4u;
        if (token == FDT_BEGIN_NODE) {
            const char *name = (const char *)struct_p;
            if (!string_in_block(name, struct_p, struct_end) || depth + 1u >= FDT_MAX_DEPTH) {
                return false;
            }
            wanted_node[depth] = hv_strcmp(name, node_name) == 0;
            while (struct_p < struct_end && *struct_p != 0u) {
                struct_p++;
            }
            if (struct_p >= struct_end) {
                return false;
            }
            struct_p = (uint8_t *)align4(struct_p + 1u);
            depth++;
        } else if (token == FDT_END_NODE) {
            if (depth == 0u) {
                return false;
            }
            depth--;
            wanted_node[depth] = false;
        } else if (token == FDT_PROP) {
            uint32_t len;
            uint32_t nameoff;
            char *name;
            if (struct_p + 8u > struct_end) {
                return false;
            }
            len = be32(struct_p);
            nameoff = be32(struct_p + 4u);
            struct_p += 8u;
            if (nameoff >= h.size_dt_strings || len > (uint32_t)(struct_end - struct_p)) {
                return false;
            }
            name = (char *)(strings + nameoff);
            if (!string_in_block(name, strings, strings_end)) {
                return false;
            }
            if (depth != 0u && wanted_node[depth - 1u] &&
                len >= value_len && hv_strlen(name) + 1u >= prop_len) {
                hv_memset(name, 0, hv_strlen(name) + 1u);
                hv_memcpy(name, prop_name, prop_len);
                hv_memset(struct_p, 0, len);
                hv_memcpy(struct_p, value, value_len);
                return true;
            }
            struct_p = (uint8_t *)align4(struct_p + len);
        } else if (token == FDT_NOP) {
            continue;
        } else if (token == FDT_END) {
            return false;
        } else {
            return false;
        }
    }
    return false;
}

static bool patch_u64_prop(uint64_t fdt_pa, const char *node, const char *prop,
                           uint64_t value)
{
    struct fdt_property_ref ref;
    if (!find_property(fdt_pa, node, prop, &ref) || ref.len < 8u) {
        return false;
    }
    write_be64(ref.value, value);
    return true;
}

static bool find_node_end(uint64_t fdt_pa, const char *node_name, uint8_t **out)
{
    uint8_t *fdt = (uint8_t *)(uintptr_t)fdt_pa;
    struct fdt_header_view h;
    uint8_t *struct_p;
    uint8_t *struct_end;
    bool wanted_node[FDT_MAX_DEPTH];
    uint32_t depth = 0;

    if (node_name == (const char *)0 || out == (uint8_t **)0 || !read_header(fdt, &h)) {
        return false;
    }
    hv_memset(wanted_node, 0, sizeof(wanted_node));
    struct_p = fdt + h.off_dt_struct;
    struct_end = struct_p + h.size_dt_struct;

    while (struct_p + 4u <= struct_end) {
        uint8_t *token_p = struct_p;
        uint32_t token = be32(struct_p);
        struct_p += 4u;
        if (token == FDT_BEGIN_NODE) {
            const char *name = (const char *)struct_p;
            if (!string_in_block(name, struct_p, struct_end) || depth + 1u >= FDT_MAX_DEPTH) {
                return false;
            }
            wanted_node[depth] = hv_strcmp(name, node_name) == 0;
            while (struct_p < struct_end && *struct_p != 0u) {
                struct_p++;
            }
            if (struct_p >= struct_end) {
                return false;
            }
            struct_p = (uint8_t *)align4(struct_p + 1u);
            depth++;
        } else if (token == FDT_END_NODE) {
            if (depth == 0u) {
                return false;
            }
            if (wanted_node[depth - 1u]) {
                *out = token_p;
                return true;
            }
            depth--;
            wanted_node[depth] = false;
        } else if (token == FDT_PROP) {
            uint32_t len;
            if (struct_p + 8u > struct_end) {
                return false;
            }
            len = be32(struct_p);
            struct_p += 8u;
            if (len > (uint32_t)(struct_end - struct_p)) {
                return false;
            }
            struct_p = (uint8_t *)align4(struct_p + len);
        } else if (token == FDT_NOP) {
            continue;
        } else if (token == FDT_END) {
            return false;
        } else {
            return false;
        }
    }
    return false;
}

static bool append_chosen_prop(uint64_t fdt_pa, uint64_t max_size, const char *prop,
                               const uint8_t *value, uint32_t len)
{
    uint8_t *fdt = (uint8_t *)(uintptr_t)fdt_pa;
    struct fdt_header_view h;
    uint8_t *insert;
    uint8_t *strings;
    uint32_t prop_name_len;
    uint32_t value_len;
    uint32_t record_size;
    uint32_t growth;
    uint32_t nameoff;

    if (prop == (const char *)0 || value == (const uint8_t *)0 ||
        max_size > FDT_MAX_SIZE || !read_header(fdt, &h) ||
        !find_node_end(fdt_pa, "chosen", &insert)) {
        return false;
    }

    prop_name_len = (uint32_t)hv_strlen(prop) + 1u;
    value_len = align4_u32(len);
    record_size = 12u + value_len;
    growth = record_size + prop_name_len;
    if (max_size < h.totalsize || growth > max_size - h.totalsize) {
        return false;
    }

    nameoff = h.size_dt_strings;
    hv_memmove(insert + record_size, insert, h.totalsize - (uint32_t)(insert - fdt));
    write_be32(insert, FDT_PROP);
    write_be32(insert + 4u, len);
    write_be32(insert + 8u, nameoff);
    hv_memset(insert + 12u, 0, value_len);
    hv_memcpy(insert + 12u, value, len);

    strings = fdt + h.off_dt_strings + record_size;
    hv_memcpy(strings + h.size_dt_strings, prop, prop_name_len);
    write_be32(fdt + 4u, h.totalsize + growth);
    write_be32(fdt + 12u, h.off_dt_strings + record_size);
    write_be32(fdt + 32u, h.size_dt_strings + prop_name_len);
    write_be32(fdt + 36u, h.size_dt_struct + record_size);
    return true;
}

static bool ensure_chosen_u64(uint64_t fdt_pa, uint64_t max_size, const char *prop,
                              uint64_t value)
{
    uint8_t buf[8];
    if (patch_u64_prop(fdt_pa, "chosen", prop, value)) {
        return true;
    }
    write_be64(buf, value);
    return append_chosen_prop(fdt_pa, max_size, prop, buf, sizeof(buf));
}

bool fdt_patch_memory(uint64_t fdt_pa, uint64_t base, uint64_t size)
{
    struct fdt_property_ref ref;

    if (size == 0u || !find_memory_property(fdt_pa, "reg", &ref) || ref.len < 16u) {
        return false;
    }

    write_be64(ref.value, base);
    write_be64(ref.value + 8u, size);
    if (ref.len > 16u) {
        hv_memset(ref.value + 16u, 0, ref.len - 16u);
    }
    return true;
}

bool fdt_bootargs_valid(const char *bootargs)
{
    size_t len = 0;
    if (bootargs == (const char *)0) {
        return true;
    }
    while (bootargs[len] != '\0') {
        if ((unsigned char)bootargs[len] < 0x20u || (unsigned char)bootargs[len] > 0x7eu) {
            return false;
        }
        len++;
        if (len >= FDT_BOOTARGS_MAX) {
            return false;
        }
    }
    return len != 0u;
}

bool fdt_probe(uint64_t fdt_pa, struct fdt_platform_info *info)
{
    const uint8_t *fdt = (const uint8_t *)(uintptr_t)fdt_pa;
    struct fdt_header_view h;
    const uint8_t *struct_p;
    const uint8_t *struct_end;
    const uint8_t *strings;
    const uint8_t *strings_end;
    bool memory_node[FDT_MAX_DEPTH];
    bool chosen_node[FDT_MAX_DEPTH];
    uint32_t depth = 0;
    uint32_t address_cells = 2u;
    uint32_t size_cells = 2u;

    if (info == (struct fdt_platform_info *)0) {
        return false;
    }
    hv_memset(info, 0, sizeof(*info));
    if (fdt_pa == 0u || !hv_is_aligned_u64(fdt_pa, 8u)) {
        return false;
    }
    if (!read_header(fdt, &h)) {
        return false;
    }

    hv_memset(memory_node, 0, sizeof(memory_node));
    hv_memset(chosen_node, 0, sizeof(chosen_node));
    struct_p = fdt + h.off_dt_struct;
    struct_end = struct_p + h.size_dt_struct;
    strings = fdt + h.off_dt_strings;
    strings_end = strings + h.size_dt_strings;

    while (struct_p + 4u <= struct_end) {
        uint32_t token = be32(struct_p);
        struct_p += 4u;

        if (token == FDT_BEGIN_NODE) {
            const char *name = (const char *)struct_p;
            if (!string_in_block(name, struct_p, struct_end) || depth + 1u >= FDT_MAX_DEPTH) {
                return false;
            }
            memory_node[depth] = node_name_is_memory(name);
            chosen_node[depth] = hv_strcmp(name, "chosen") == 0;
            while (struct_p < struct_end && *struct_p != 0u) {
                struct_p++;
            }
            if (struct_p >= struct_end) {
                return false;
            }
            struct_p = align4(struct_p + 1u);
            depth++;
        } else if (token == FDT_END_NODE) {
            if (depth == 0u) {
                return false;
            }
            depth--;
            memory_node[depth] = false;
            chosen_node[depth] = false;
        } else if (token == FDT_PROP) {
            uint32_t len;
            uint32_t nameoff;
            const char *prop_name;
            const uint8_t *value;

            if (struct_p + 8u > struct_end) {
                return false;
            }
            len = be32(struct_p);
            nameoff = be32(struct_p + 4u);
            struct_p += 8u;
            if (nameoff >= h.size_dt_strings || len > (uint32_t)(struct_end - struct_p)) {
                return false;
            }
            prop_name = (const char *)(strings + nameoff);
            if (!string_in_block(prop_name, strings, strings_end)) {
                return false;
            }
            value = struct_p;

            if (depth == 1u && hv_strcmp(prop_name, "#address-cells") == 0 && len >= 4u) {
                address_cells = be32(value);
            } else if (depth == 1u && hv_strcmp(prop_name, "#size-cells") == 0 && len >= 4u) {
                size_cells = be32(value);
            } else if (depth != 0u && hv_strcmp(prop_name, "device_type") == 0 &&
                       len >= 7u && hv_strcmp((const char *)value, "memory") == 0) {
                memory_node[depth - 1u] = true;
            } else if (depth != 0u && memory_node[depth - 1u] &&
                       hv_strcmp(prop_name, "reg") == 0 &&
                       address_cells <= 2u && size_cells <= 2u &&
                       len >= (address_cells + size_cells) * 4u) {
                const uint32_t *cells = (const uint32_t *)value;
                info->memory_base = read_cells(cells, address_cells);
                info->memory_size = read_cells(cells + address_cells, size_cells);
            } else if (depth != 0u && chosen_node[depth - 1u]) {
                if (hv_strcmp(prop_name, "bootargs") == 0 && len != 0u) {
                    info->bootargs = (const char *)value;
                } else if (hv_strcmp(prop_name, "linux,initrd-start") == 0 && len >= 4u) {
                    info->initrd_start = read_cells((const uint32_t *)value, len >= 8u ? 2u : 1u);
                } else if (hv_strcmp(prop_name, "linux,initrd-end") == 0 && len >= 4u) {
                    info->initrd_end = read_cells((const uint32_t *)value, len >= 8u ? 2u : 1u);
                }
            }

            struct_p = align4(struct_p + len);
        } else if (token == FDT_NOP) {
            continue;
        } else if (token == FDT_END) {
            info->valid = true;
            info->fdt_pa = fdt_pa;
            info->size = h.totalsize;
            return true;
        } else {
            return false;
        }
    }

    return false;
}

bool fdt_patch_linux_guest(uint64_t fdt_pa, uint64_t max_size, const char *bootargs,
                           uint64_t initrd_start, uint64_t initrd_end)
{
    struct fdt_platform_info info;
    bool ok = true;

    if (!fdt_bootargs_valid(bootargs) || !fdt_probe(fdt_pa, &info)) {
        return false;
    }
    if (bootargs != (const char *)0) {
        ok = patch_string_prop(fdt_pa, "chosen", "bootargs", bootargs) ||
             repurpose_string_prop(fdt_pa, "chosen", "bootargs", bootargs) ||
             append_chosen_prop(fdt_pa, max_size, "bootargs",
                                (const uint8_t *)bootargs,
                                (uint32_t)hv_strlen(bootargs) + 1u);
    }
    if (initrd_start != 0u && initrd_end > initrd_start) {
        ok = ok && ensure_chosen_u64(fdt_pa, max_size, "linux,initrd-start", initrd_start);
        ok = ok && ensure_chosen_u64(fdt_pa, max_size, "linux,initrd-end", initrd_end);
    }
    return ok;
}

void fdt_dump(const struct fdt_platform_info *info)
{
    if (info == (const struct fdt_platform_info *)0 || !info->valid) {
        uart_puts("fdt unavailable\n");
        return;
    }
    uart_puts("fdt pa=");
    uart_put_hex64(info->fdt_pa);
    uart_puts(" size=");
    uart_put_dec64(info->size);
    uart_puts(" memory=");
    uart_put_hex64(info->memory_base);
    uart_puts("+");
    uart_put_hex64(info->memory_size);
    if (info->initrd_start != 0u || info->initrd_end != 0u) {
        uart_puts(" initrd=");
        uart_put_hex64(info->initrd_start);
        uart_puts("-");
        uart_put_hex64(info->initrd_end);
    }
    uart_puts("\n");
    if (info->bootargs != (const char *)0) {
        uart_puts("bootargs ");
        uart_puts(info->bootargs);
        uart_puts("\n");
    }
}