#include <zinnia/archctl.h>
#include <zinnia/channel.h>
#include <zinnia/futex.h>
#include <zinnia/handle.h>
#include <zinnia/mem.h>
#include <zinnia/namespace.h>
#include <zinnia/random.h>
#include <zinnia/status.h>
#include <zinnia/system.h>
#include <zinnia/thread.h>
#include <common/elf.h>
#include <sys/auxv.h>

static uint8_t vdso_symbol_bind(uint8_t info) {
    return ELF64_ST_BIND(info);
}

static uint8_t vdso_symbol_type(uint8_t info) {
    return ELF64_ST_TYPE(info);
}

struct vdso_info {
    uintptr_t base;
    const char* strings;
    const struct elf64_sym* symbols;
    size_t symbol_count;
};

static bool vdso_info_init(struct vdso_info* info) {
    const uintptr_t base = getauxval(AT_SYSINFO_EHDR);
    if (!base)
        return false;

    const struct elf64_ehdr* ehdr = (const void*)base;
    if (ehdr->e_ident[EI_MAG0] != ELF_MAG[0] || ehdr->e_ident[EI_MAG1] != ELF_MAG[1] ||
        ehdr->e_ident[EI_MAG2] != ELF_MAG[2] || ehdr->e_ident[EI_MAG3] != ELF_MAG[3] ||
        ehdr->e_ident[EI_CLASS] != ELFCLASS64 || ehdr->e_ident[EI_DATA] != ELFDATA2LSB ||
        ehdr->e_machine != ELF_ARCH_MACHINE)
        return false;

    const struct elf64_dyn* dynamic = nullptr;
    for (size_t i = 0; i < ehdr->e_phnum; i++) {
        const struct elf64_phdr* phdr = (const void*)(base + ehdr->e_phoff + (i * ehdr->e_phentsize));
        if (phdr->p_type == PT_DYNAMIC) {
            dynamic = (const void*)(base + phdr->p_vaddr);
            break;
        }
    }
    if (!dynamic)
        return false;

    const char* strings = nullptr;
    const struct elf64_sym* symbols = nullptr;
    const struct elf_hash* hash = nullptr;

    for (const struct elf64_dyn* entry = dynamic; entry->d_tag != DT_NULL; entry++) {
        switch (entry->d_tag) {
        case DT_STRTAB:
            strings = (const void*)(base + entry->d_un.d_ptr);
            break;
        case DT_SYMTAB:
            symbols = (const void*)(base + entry->d_un.d_ptr);
            break;
        case DT_HASH:
            hash = (const void*)(base + entry->d_un.d_ptr);
            break;
        default:
            break;
        }
    }

    if (!strings || !symbols || !hash)
        return false;

    *info = (struct vdso_info){
        .base = base,
        .strings = strings,
        .symbols = symbols,
        .symbol_count = hash->n_chain,
    };
    return true;
}

static void* vdso_lookup(const char* name) {
    struct vdso_info info;
    if (!vdso_info_init(&info))
        return nullptr;

    for (size_t i = 0; i < info.symbol_count; i++) {
        const struct elf64_sym* symbol = &info.symbols[i];
        if (vdso_symbol_bind(symbol->st_info) != STB_GLOBAL || vdso_symbol_type(symbol->st_info) != STT_FUNC)
            continue;

        const char* symbol_name = info.strings + symbol->st_name;
        const char* lhs = symbol_name;
        const char* rhs = name;
        while (*lhs && *lhs == *rhs) {
            lhs++;
            rhs++;
        }
        if (*lhs == 0 && *rhs == 0)
            return (void*)(info.base + symbol->st_value);
    }

    return nullptr;
}

#define VDSO_IFUNC(name) \
    static typeof(name)* resolve_##name(void) { \
        return (typeof(name)*)vdso_lookup("__vdso_" #name); \
    } \
    typeof(name) name __attribute__((ifunc("resolve_" #name)));

VDSO_IFUNC(zn_archctl);
VDSO_IFUNC(zn_channel_create);
VDSO_IFUNC(zn_channel_write);
VDSO_IFUNC(zn_channel_read);
VDSO_IFUNC(zn_futex_wait);
VDSO_IFUNC(zn_futex_wake);
VDSO_IFUNC(zn_handle_validate);
VDSO_IFUNC(zn_handle_drop);
VDSO_IFUNC(zn_handle_clone);
VDSO_IFUNC(zn_vmo_create);
VDSO_IFUNC(zn_vmo_create_phys);
VDSO_IFUNC(zn_vmo_map);
VDSO_IFUNC(zn_vas_create);
VDSO_IFUNC(zn_ns_create);
VDSO_IFUNC(zn_ns_move);
VDSO_IFUNC(zn_random_get);
VDSO_IFUNC(zn_random_entropy);
VDSO_IFUNC(zn_status_to_string);
VDSO_IFUNC(zn_log);
VDSO_IFUNC(zn_page_size);
VDSO_IFUNC(zn_powerctl);
VDSO_IFUNC(zn_time);
VDSO_IFUNC(zn_thread_create);
VDSO_IFUNC(zn_thread_start);
VDSO_IFUNC(zn_thread_exit);