/*
    * LibSyscall.hpp
    * SYS_LOAD_LIB, SYS_UNLOAD_LIB, SYS_DLSYM syscalls
    * Copyright (c) 2026 Daniel Hammer
*/

#include <Memory/UserRange.hpp>
#include <cstdint>
#include <Sched/Scheduler.hpp>
#include <Sched/ElfLoader.hpp>
#include <Memory/Paging.hpp>
#include <Memory/HHDM.hpp>
#include <Memory/PageFrameAllocator.hpp>
#include <Terminal/Terminal.hpp>
#include <Libraries/Memory.hpp>
#include <Libraries/String.hpp>
#include <CppLib/Spinlock.hpp>
#include <Ipc/Ipc.hpp>

namespace montauk::abi {

    // Maximum libraries per process
    static constexpr int MaxLibsPerProcess = 8;
    static constexpr int MaxProcesses = Sched::MaxProcesses;

    // Library entry tracking
    struct LibEntry {
        char path[128];
        uint64_t loadBias;
        uint32_t refcount;
        bool inUse;
    };

    inline uint64_t GetLibSlotBase(int libSlot) {
        return Sched::LIB_BASE + (uint64_t(libSlot) * Sched::LIB_MAX_SIZE);
    }

    // Per-process library table
    inline LibEntry g_libTable[MaxProcesses][MaxLibsPerProcess] = {};
    inline kcp::Mutex g_libLocks[MaxProcesses];

    // Initialize library table for a process slot
    inline void InitLibTable(int slot) {
        if (slot < 0 || slot >= MaxProcesses) return;
        g_libLocks[slot].Acquire();
        for (int i = 0; i < MaxLibsPerProcess; i++) {
            g_libTable[slot][i].inUse = false;
            g_libTable[slot][i].refcount = 0;
            g_libTable[slot][i].loadBias = 0;
            g_libTable[slot][i].path[0] = '\0';
        }
        g_libLocks[slot].Release();
    }

    // Load a shared library into the current process's address space.
    // Returns a library handle (slot index + 1) on success, or 0 on failure.
    // The library is loaded at a fixed address based on the slot.
    inline uint64_t Sys_LoadLib(const char* path) {
        int slot = GetCurrentSlot();
        if (slot < 0) return 0;

        auto* proc = Sched::GetCurrentProcessPtr();
        if (proc == nullptr) return 0;

        // Validate path
        if (path == nullptr) return 0;
        size_t pathLen = Lib::strlen(path);
        if (pathLen == 0 || pathLen >= sizeof(LibEntry::path)) return 0;

        g_libLocks[slot].Acquire();

        // Find a free slot or reuse an existing slot with the same path
        int libSlot = -1;
        for (int i = 0; i < MaxLibsPerProcess; i++) {
            if (g_libTable[slot][i].inUse) {
                if (Lib::strncmp(g_libTable[slot][i].path, path, sizeof(LibEntry::path)) == 0) {
                    // Same library already loaded - just increment refcount
                    g_libTable[slot][i].refcount++;
                    g_libLocks[slot].Release();
                    return (uint64_t)(i + 1);  // Handle = slot + 1 (0 = invalid)
                }
            } else if (libSlot < 0) {
                libSlot = i;
            }
        }

        if (libSlot < 0) {
            Kt::KernelLogStream(Kt::ERROR, "Lib") << "No free library slots";
            g_libLocks[slot].Release();
            return 0;
        }

        // Load the library into the process's address space
        uint64_t loadBias = Sched::ElfLoadLib(path, proc->pml4Phys, libSlot);
        if (loadBias == 0) {
            Kt::KernelLogStream(Kt::ERROR, "Lib") << "Failed to load library: " << path;
            g_libLocks[slot].Release();
            return 0;
        }

        // Store library info
        g_libTable[slot][libSlot].inUse = true;
        g_libTable[slot][libSlot].refcount = 1;
        g_libTable[slot][libSlot].loadBias = loadBias;
        Lib::strncpy(g_libTable[slot][libSlot].path, path, sizeof(LibEntry::path));

        g_libLocks[slot].Release();

        Kt::KernelLogStream(Kt::OK, "Lib") << "Loaded library: " << path << " with load bias " << kcp::hex << loadBias << kcp::dec;

        return (uint64_t)(libSlot + 1);  // Handle = slot + 1 (0 = invalid)
    }

    // Unload a shared library.
    // Decrements refcount; actually unmaps when refcount reaches 0.
    inline int Sys_UnloadLib(uint64_t handle) {
        int slot = GetCurrentSlot();
        if (slot < 0) return GenericFailure;

        int libSlot = (int)handle - 1;
        if (libSlot < 0 || libSlot >= MaxLibsPerProcess) return InvalidArgument;

        g_libLocks[slot].Acquire();
        if (!g_libTable[slot][libSlot].inUse) {
            g_libLocks[slot].Release();
            return InvalidArgument;
        }

        g_libTable[slot][libSlot].refcount--;
        if (g_libTable[slot][libSlot].refcount == 0) {
            // Actually unload - free the pages
            uint64_t libBase = GetLibSlotBase(libSlot);
            uint64_t libEnd = libBase + Sched::LIB_MAX_SIZE;

            auto* proc = Sched::GetCurrentProcessPtr();
            if (proc != nullptr) {
                Memory::UnmapAndFreeUserRange(proc->pml4Phys, libBase,
                                           (libEnd - libBase) / 0x1000ULL);
            }

            g_libTable[slot][libSlot].inUse = false;
            g_libTable[slot][libSlot].loadBias = 0;
            g_libTable[slot][libSlot].path[0] = '\0';
        }

        g_libLocks[slot].Release();
        return Success;
    }

    // Resolve a symbol in a loaded library.
    // handle = library handle from LoadLib
    // symbolOffset = ELF symbol value from the library's symbol table
    // Returns the virtual address of the symbol, or 0 if not found.
    inline uint64_t Sys_DLSym(uint64_t handle, uint64_t symbolOffset) {
        int slot = GetCurrentSlot();
        if (slot < 0) return 0;

        int libSlot = (int)handle - 1;
        if (libSlot < 0 || libSlot >= MaxLibsPerProcess) return 0;

        g_libLocks[slot].Acquire();
        if (!g_libTable[slot][libSlot].inUse) {
            g_libLocks[slot].Release();
            return 0;
        }

        uint64_t loadBias = g_libTable[slot][libSlot].loadBias;
        g_libLocks[slot].Release();
        return loadBias + symbolOffset;
    }

    // Get library relocation base / load bias (for userspace symbol resolution)
    inline uint64_t Sys_GetLibBase(uint64_t handle) {
        int slot = GetCurrentSlot();
        if (slot < 0) return 0;

        int libSlot = (int)handle - 1;
        if (libSlot < 0 || libSlot >= MaxLibsPerProcess) return 0;

        g_libLocks[slot].Acquire();
        if (!g_libTable[slot][libSlot].inUse) {
            g_libLocks[slot].Release();
            return 0;
        }

        uint64_t loadBias = g_libTable[slot][libSlot].loadBias;
        g_libLocks[slot].Release();
        return loadBias;
    }

    // Cleanup library table for a process slot
    inline void CleanupLibTable(int slot) {
        if (slot < 0 || slot >= MaxProcesses) return;

        auto* proc = Sched::GetProcessSlot(slot);
        if (proc == nullptr) return;

        g_libLocks[slot].Acquire();

        // Unmap all libraries for this process
        for (int i = 0; i < MaxLibsPerProcess; i++) {
            if (g_libTable[slot][i].inUse) {
                uint64_t libBase = GetLibSlotBase(i);
                uint64_t libEnd = libBase + Sched::LIB_MAX_SIZE;

                Memory::UnmapAndFreeUserRange(proc->pml4Phys, libBase,
                                           (libEnd - libBase) / 0x1000ULL);

                g_libTable[slot][i].inUse = false;
                g_libTable[slot][i].refcount = 0;
                g_libTable[slot][i].loadBias = 0;
            }
        }
        g_libLocks[slot].Release();
    }

}