/*
    * libloader.cpp
    * Dynamic library loading implementation for MontaukOS
    * Copyright (c) 2026 Daniel Hammer
*/

#include <libloader/libloader.h>
#include <montauk/syscall.h>
#include <libc/string.h>
#include <stdlib.h>

namespace libloader {

// Table of currently loaded libraries
static LibHandle g_loadedLibs[MaxLoadedLibs];

namespace {

struct Elf64Header {
    uint8_t  e_ident[16];
    uint16_t e_type;
    uint16_t e_machine;
    uint32_t e_version;
    uint64_t e_entry;
    uint64_t e_phoff;
    uint64_t e_shoff;
    uint32_t e_flags;
    uint16_t e_ehsize;
    uint16_t e_phentsize;
    uint16_t e_phnum;
    uint16_t e_shentsize;
    uint16_t e_shnum;
    uint16_t e_shstrndx;
};

struct Elf64SectionHeader {
    uint32_t sh_name;
    uint32_t sh_type;
    uint64_t sh_flags;
    uint64_t sh_addr;
    uint64_t sh_offset;
    uint64_t sh_size;
    uint32_t sh_link;
    uint32_t sh_info;
    uint64_t sh_addralign;
    uint64_t sh_entsize;
};

struct Elf64Sym {
    uint32_t st_name;
    uint8_t  st_info;
    uint8_t  st_other;
    uint16_t st_shndx;
    uint64_t st_value;
    uint64_t st_size;
};

static constexpr uint16_t ET_EXEC = 2;
static constexpr uint16_t ET_DYN = 3;
static constexpr uint16_t EM_X86_64 = 62;
static constexpr uint32_t SHT_SYMTAB = 2;
static constexpr uint32_t SHT_STRTAB = 3;
static constexpr uint32_t SHT_DYNSYM = 11;
static constexpr uint16_t SHN_UNDEF = 0;

static bool readFile(const char* path, uint8_t** outBuf, uint64_t* outSize) {
    *outBuf = nullptr;
    *outSize = 0;

    int fd = montauk::open(path);
    if (fd < 0) return false;

    uint64_t size = montauk::getsize(fd);
    if (size < sizeof(Elf64Header) || size > (1 << 20)) {
        montauk::close(fd);
        return false;
    }

    auto* buf = (uint8_t*)montauk::alloc(size);
    if (!buf) {
        montauk::close(fd);
        return false;
    }

    int read = montauk::read(fd, buf, 0, size);
    montauk::close(fd);
    if (read != (int)size) {
        montauk::free(buf);
        return false;
    }

    *outBuf = buf;
    *outSize = size;
    return true;
}

static bool isValidElf(const Elf64Header* hdr, uint64_t size) {
    if (size < sizeof(Elf64Header)) return false;
    if (hdr->e_ident[0] != 0x7f || hdr->e_ident[1] != 'E' ||
        hdr->e_ident[2] != 'L' || hdr->e_ident[3] != 'F') {
        return false;
    }
    if (hdr->e_ident[4] != 2 || hdr->e_ident[5] != 1) return false;
    if (hdr->e_machine != EM_X86_64) return false;
    if (hdr->e_type != ET_EXEC && hdr->e_type != ET_DYN) return false;
    if (hdr->e_shentsize != sizeof(Elf64SectionHeader)) return false;
    if (hdr->e_shoff == 0 || hdr->e_shnum == 0) return false;
    if (hdr->e_shoff + uint64_t(hdr->e_shnum) * hdr->e_shentsize > size) return false;
    if (hdr->e_shstrndx >= hdr->e_shnum) return false;
    return true;
}

static const Elf64SectionHeader* getSectionHeaders(const Elf64Header* hdr) {
    return (const Elf64SectionHeader*)((const uint8_t*)hdr + hdr->e_shoff);
}

static const char* getSectionNameTable(const Elf64Header* hdr, uint64_t fileSize) {
    const Elf64SectionHeader* shdrs = getSectionHeaders(hdr);
    const Elf64SectionHeader& strhdr = shdrs[hdr->e_shstrndx];
    if (strhdr.sh_type != SHT_STRTAB) return nullptr;
    if (strhdr.sh_offset + strhdr.sh_size > fileSize) return nullptr;
    return (const char*)hdr + strhdr.sh_offset;
}

static const Elf64SectionHeader* findSection(const Elf64Header* hdr, uint64_t fileSize,
                                             const char* name, uint32_t expectedType) {
    const Elf64SectionHeader* shdrs = getSectionHeaders(hdr);
    const char* shstrtab = getSectionNameTable(hdr, fileSize);
    if (!shstrtab) return nullptr;

    for (uint16_t i = 0; i < hdr->e_shnum; i++) {
        const Elf64SectionHeader* sh = &shdrs[i];
        if (sh->sh_name >= shdrs[hdr->e_shstrndx].sh_size) continue;
        if (sh->sh_type != expectedType) continue;
        const char* secName = shstrtab + sh->sh_name;
        if (strcmp(secName, name) == 0) return sh;
    }
    return nullptr;
}

static bool lookupSymbolValue(const char* path, const char* symbolName, uint64_t* outValue) {
    uint8_t* fileData = nullptr;
    uint64_t fileSize = 0;
    if (!readFile(path, &fileData, &fileSize)) return false;

    const Elf64Header* hdr = (const Elf64Header*)fileData;
    if (!isValidElf(hdr, fileSize)) {
        montauk::free(fileData);
        return false;
    }

    const Elf64SectionHeader* symtab = findSection(hdr, fileSize, ".dynsym", SHT_DYNSYM);
    const Elf64SectionHeader* strtab = findSection(hdr, fileSize, ".dynstr", SHT_STRTAB);
    if (!symtab || !strtab) {
        symtab = findSection(hdr, fileSize, ".symtab", SHT_SYMTAB);
        strtab = findSection(hdr, fileSize, ".strtab", SHT_STRTAB);
    }

    if (!symtab || !strtab || symtab->sh_entsize != sizeof(Elf64Sym) ||
        symtab->sh_offset + symtab->sh_size > fileSize ||
        strtab->sh_offset + strtab->sh_size > fileSize) {
        montauk::free(fileData);
        return false;
    }

    const Elf64Sym* symbols = (const Elf64Sym*)(fileData + symtab->sh_offset);
    const char* strings = (const char*)(fileData + strtab->sh_offset);
    uint64_t count = symtab->sh_size / sizeof(Elf64Sym);

    bool found = false;
    for (uint64_t i = 0; i < count; i++) {
        const Elf64Sym& sym = symbols[i];
        if (sym.st_shndx == SHN_UNDEF || sym.st_name >= strtab->sh_size) continue;
        const char* name = strings + sym.st_name;
        if (strcmp(name, symbolName) == 0) {
            *outValue = sym.st_value;
            found = true;
            break;
        }
    }

    montauk::free(fileData);
    return found;
}

static LibHandle* findFreeLib() {
    for (int i = 0; i < MaxLoadedLibs; i++) {
        if (g_loadedLibs[i].handle <= 0) return &g_loadedLibs[i];
    }
    return nullptr;
}

} // namespace

static LibHandle* findLoadedLib(const char* path) {
    for (int i = 0; i < MaxLoadedLibs; i++) {
        if (g_loadedLibs[i].handle > 0 &&
            strcmp(g_loadedLibs[i].path, path) == 0) {
            return &g_loadedLibs[i];
        }
    }
    return nullptr;
}

LibHandle* dlopen(const char* path) {
    if (!path) return nullptr;

    // Check if already loaded
    LibHandle* existing = findLoadedLib(path);
    if (existing) {
        // Just return existing handle
        return existing;
    }

    // Load the library via syscall
    int handle = montauk::load_lib(path);
    if (handle <= 0) {
        return nullptr;
    }

    LibHandle* lib = findFreeLib();
    if (!lib) {
        montauk::unload_lib(handle);
        return nullptr;
    }

    lib->handle = handle;
    lib->base = montauk::get_libbase(handle);
    strcpy(lib->path, path);

    return lib;
}

void* dlsym(LibHandle* lib, const char* symbolName) {
    if (!lib || !symbolName) return nullptr;

    // Get base address from kernel if we don't have it
    if (lib->base == 0) {
        lib->base = montauk::get_libbase(lib->handle);
    }

    uint64_t value = 0;
    if (!lookupSymbolValue(lib->path, symbolName, &value)) return nullptr;
    return (void*)(lib->base + value);
}

int dlclose(LibHandle* lib) {
    if (!lib) return -1;

    int result = montauk::unload_lib(lib->handle);
    if (result == 0) {
        lib->handle = 0;
        lib->base = 0;
        lib->path[0] = '\0';
    }

    return result;
}

}  // namespace libloader