#include "ModuleManagement.h"

#include "HooksManagement.h"
#include "ModuleAllocator.h"
#include "ModuleContainer.h"
#include "ModuleDataPersistence.h"
#include "ModuleLinkInformationFactory.h"
#include "ModuleMetaInformation.h"
#include "ModuleMetaInformationFactory.h"
#include "RelocationUtils.h"
#include "globals.h"
#include "utils/LoadedFile.h"
#include "utils/StringTools.h"

#include "utils/logger.h"

#include <algorithm>
#include <sys/dirent.h>
#include <vector>

namespace WUMSLoader::Modules {

namespace {

struct ModuleInfoWithLoadedFileWrapper {
    ModuleMetaInformation metaInfo;
    Utils::LoadedFile file;
};

std::vector<std::string> getModuleFilePaths(const std::string &basePath) {
    PROFILE_FUNCTION();
    std::vector<std::string> result;
    struct dirent *dp;
    DIR *dfd;

    if (basePath.empty()) {
        DEBUG_FUNCTION_LINE_ERR("Failed to scan module dir: Path was empty");
        return result;
    }

    if ((dfd = opendir(basePath.c_str())) == nullptr) {
        DEBUG_FUNCTION_LINE_ERR("Couldn't open dir %s", basePath.c_str());
        return result;
    }

    while ((dp = readdir(dfd)) != nullptr) {
        if (dp->d_type == DT_DIR) {
            continue;
        }

        std::string_view fileName(dp->d_name);
        if (fileName.starts_with('.') || fileName.starts_with('_') || !fileName.ends_with(".wms")) {
            DEBUG_FUNCTION_LINE_WARN("Skip file %s/%s", basePath.c_str(), dp->d_name);
            continue;
        }

        result.push_back(basePath + "/" + dp->d_name);
    }
    closedir(dfd);
    return result;
}

bool CheckModulesByDependencies(const std::vector<ModuleInfoWithLoadedFileWrapper> &modules) {
    PROFILE_FUNCTION();
    std::set<std::string, std::less<>> loaderModuleNames;

    for (const auto &curModule : modules) {
        const auto &exportName = curModule.metaInfo.getExportName();

        DEBUG_FUNCTION_LINE_VERBOSE("Check if we can load %s", exportName.c_str());

        for (const auto &curRPL : curModule.metaInfo.getDependencyList()) {
            if (!curRPL.starts_with("homebrew")) {
                continue;
            }
            if (curRPL == "homebrew_wupsbackend") {
                OSFatal("Error: module depends on homebrew_wupsbackend, this is not supported");
            }
            if (!loaderModuleNames.contains(curRPL)) {
                DEBUG_FUNCTION_LINE_VERBOSE("%s requires %s which is not loaded yet", exportName.c_str(), curRPL.c_str());
                return false;
            } else {
                DEBUG_FUNCTION_LINE_VERBOSE("Used %s, but it's already loaded", curRPL.c_str());
            }
        }
        loaderModuleNames.insert(exportName);
    }
    return true;
}

std::vector<ModuleInfoWithLoadedFileWrapper> OrderModulesMetaByDependencies(std::vector<ModuleInfoWithLoadedFileWrapper> &&modulesMetaInfo) {
    PROFILE_FUNCTION();
    std::vector<std::reference_wrapper<ModuleInfoWithLoadedFileWrapper>> pendingModules;
    pendingModules.reserve(modulesMetaInfo.size());
    for (auto &mod : modulesMetaInfo) {
        pendingModules.push_back(std::ref(mod));
    }

    std::vector<std::reference_wrapper<ModuleInfoWithLoadedFileWrapper>> orderedRefs;
    orderedRefs.reserve(modulesMetaInfo.size());

    std::set<std::string, std::less<>> loadedModulesExportNames;

    while (!pendingModules.empty()) {
        const size_t previousSize = pendingModules.size();

        std::erase_if(pendingModules, [&](auto &ref) {
            auto &metaInfo = ref.get();

            const bool hasUnresolvedDeps = std::ranges::any_of(
                    metaInfo.metaInfo.getDependencyList(),
                    [&](const auto &curImportRPL) {
                        if (!curImportRPL.starts_with("homebrew")) {
                            return false;
                        }

                        if (curImportRPL == "homebrew_wupsbackend") {
                            OSFatal("Error: module depends on homebrew_wupsbackend, this is not supported");
                        }

                        if (!loadedModulesExportNames.contains(curImportRPL)) {
                            DEBUG_FUNCTION_LINE_VERBOSE("We can't load the module, because %s is not loaded yet", curImportRPL.c_str());
                            return true;
                        }
                        return false;
                    });

            if (!hasUnresolvedDeps) {
                orderedRefs.push_back(ref);
                loadedModulesExportNames.insert(metaInfo.metaInfo.getExportName());
                return true;
            }

            return false;
        });

        if (pendingModules.size() == previousSize) {
            OSFatal("Failed to resolve dependencies.");
        }
    }

    std::vector<ModuleInfoWithLoadedFileWrapper> finalOrder;
    finalOrder.reserve(orderedRefs.size());

    for (auto &ref : orderedRefs) {
        finalOrder.push_back(std::move(ref.get()));
    }

    return finalOrder;
}

void CallInitHooksForModule(const ModuleContainer &curModule) {
    HooksManagement::CallHook(curModule, WUMS_HOOK_INIT_WUT_THREAD);
    HooksManagement::CallHook(curModule, WUMS_HOOK_INIT_WUT_MALLOC);
    HooksManagement::CallHook(curModule, WUMS_HOOK_INIT_WUT_NEWLIB);
    HooksManagement::CallHook(curModule, WUMS_HOOK_INIT_WUT_STDCPP);
    HooksManagement::CallHook(curModule, WUMS_HOOK_INIT_WUT_DEVOPTAB);
    HooksManagement::CallHook(curModule, WUMS_HOOK_INIT_WUT_SOCKETS);
    HooksManagement::CallHook(curModule, WUMS_HOOK_INIT_WRAPPER, !curModule.getMetaInformation().isSkipInitFini());
    HooksManagement::CallHook(curModule, WUMS_HOOK_INIT);
}

std::vector<ModuleInfoWithLoadedFileWrapper> loadModuleMetaInformationForPath(std::vector<Utils::LoadedFile> &&loadedFiles) {
    PROFILE_FUNCTION();

    std::vector<ModuleInfoWithLoadedFileWrapper> preloadedModules;

    for (auto &loadedFile : loadedFiles) {
        ModuleParseErrors error = MODULE_PARSE_ERROR_NONE;
        if (auto metaInfoOpt = ModuleMetaInformationFactory::loadModuleMetaInfo(std::span(loadedFile.data(), loadedFile.size()), error)) {
            preloadedModules.push_back({std::move(*metaInfoOpt), std::move(loadedFile)});
        } else {
            auto errMsg = string_format("Failed to load module meta information");
            if (error == MODULE_PARSE_ERROR_INCOMPATIBLE_VERSION) {
                errMsg += ". Incompatible version.";
            }
            DEBUG_FUNCTION_LINE_ERR("%s", errMsg.c_str());
            OSFatal(errMsg.c_str());
        }
    }
    return preloadedModules;
}

std::vector<Utils::LoadedFile> loadModulesIntoMemory(const std::string &basePath) {
    PROFILE_FUNCTION();
    std::vector<Utils::LoadedFile> result;
    for (const auto &modulePath : getModuleFilePaths(basePath)) {
        auto fileOpt = Utils::LoadFileToMem(modulePath);
        if (!fileOpt) {
            DEBUG_FUNCTION_LINE_ERR("Failed to load file");
            return {};
        }
        result.emplace_back(std::move(*fileOpt));
    }
    return result;
}

std::pair<std::vector<ModuleInfoWithLoadedFileWrapper>, std::vector<ModuleInfoWithLoadedFileWrapper>> getModulesInCorrectOrderAndSplit(std::vector<ModuleInfoWithLoadedFileWrapper> &&metaInfosWithLoadedFiles) {
    PROFILE_FUNCTION();
    // Order modules list by dependencies.
    metaInfosWithLoadedFiles = OrderModulesMetaByDependencies(std::move(metaInfosWithLoadedFiles));

    // Make sure the plugin backend module is at the very end as module might depend on all other modules
    const auto backendIt = std::ranges::find_if(metaInfosWithLoadedFiles, [](const auto &cur) {
        return std::string_view(cur.metaInfo.getExportName()) == "homebrew_wupsbackend";
    });

    if (backendIt != metaInfosWithLoadedFiles.end()) {
        std::ranges::rotate(backendIt, backendIt + 1, metaInfosWithLoadedFiles.end());
    }

    // Make sure the "basemodule" is the first of the "regular" modules.
    auto regularModulesMetaInfoRange = std::ranges::stable_partition(metaInfosWithLoadedFiles, [](const auto &cur) {
        return cur.metaInfo.isInitBeforeRelocationDoneHook();
    });
    //  Ensure homebrew_basemodule is the first of the regular modules (if it exists)
    const auto baseModuleIt = std::ranges::find_if(regularModulesMetaInfoRange, [](const auto &cur) {
        return std::string_view(cur.metaInfo.getExportName()) == "homebrew_basemodule";
    });

    if (baseModuleIt != regularModulesMetaInfoRange.end()) {
        // Rotate the base module to the start of the regular modules subrange
        std::ranges::rotate(regularModulesMetaInfoRange.begin(), baseModuleIt, baseModuleIt + 1);
    }

    // Final sanity check that the order is still meeting dependency requirements
    if (!CheckModulesByDependencies(metaInfosWithLoadedFiles)) {
        OSFatal("Failed to order Aroma Modules by depenencies");
    }

    // split into "early and "regular" modules
    std::vector<ModuleInfoWithLoadedFileWrapper> earlyModules;
    std::vector<ModuleInfoWithLoadedFileWrapper> regularModules;

    // Use the iterator returned from stable_partition to know exactly where the split is
    const auto boundaryIt = regularModulesMetaInfoRange.begin();

    // Move elements to avoid expensive copying
    std::move(metaInfosWithLoadedFiles.begin(), boundaryIt, std::back_inserter(earlyModules));
    std::move(boundaryIt, metaInfosWithLoadedFiles.end(), std::back_inserter(regularModules));

    return {std::move(earlyModules), std::move(regularModules)};
}


} // namespace

std::vector<ModuleContainer> ModuleManagement::loadAndLinkModules(const std::string &basePath, std::map<std::string, OSDynLoad_Module, std::less<>> &usedRPls, MEMHeapHandle defaultHeapHandle, PersistedModuleData &persistedData) {
    PROFILE_FUNCTION();
    auto loadedModules   = loadModulesIntoMemory(basePath);
    auto moduleMetaInfos = loadModuleMetaInformationForPath(std::move(loadedModules));

    // Sort modules to satisfy dependencies; split into two groups
    auto [earlyModulesMetaInfoWithLoadedFiles, regularModuleMetaInfoWithLoadedFiles] = getModulesInCorrectOrderAndSplit(std::move(moduleMetaInfos));

    std::vector<ModuleContainer> moduleContainers;
    moduleContainers.reserve(earlyModulesMetaInfoWithLoadedFiles.size() + regularModuleMetaInfoWithLoadedFiles.size());

#ifdef VERBOSE_DEBUG
    DEBUG_FUNCTION_LINE_VERBOSE("Final order of early modules");
    for (auto &[metaInfo, file] : earlyModulesMetaInfoWithLoadedFiles) {
        DEBUG_FUNCTION_LINE_VERBOSE("%s", metaInfo.getExportName().c_str());
    }
    DEBUG_FUNCTION_LINE_VERBOSE("Final order of regular modules");
    for (auto &[metaInfo, file] : regularModuleMetaInfoWithLoadedFiles) {
        DEBUG_FUNCTION_LINE_VERBOSE("%s", metaInfo.getExportName().c_str());
    }
#endif

    // ========================================================================
    // PHASE 1: Load and link "early" modules.
    // These will get loaded into this smaller heap living between 0x00800000 and 0x01000000
    // ========================================================================
    ModuleAllocator defaultAllocator = ModuleAllocator::FromExpHeap(defaultHeapHandle);
    for (auto &[metaInfo, file] : earlyModulesMetaInfoWithLoadedFiles) {
        auto error = ModuleLinkErrors::MODULE_LINK_ERROR_NONE;
        if (auto linkInfo = ModuleLinkInformationFactory::loadAndLink(std::span(file.data(), file.size()), defaultAllocator, error)) {
            moduleContainers.emplace_back(std::move(metaInfo), std::move(*linkInfo));
        } else {
            auto errMsg = string_format("Failed to load/link module %s. %s", metaInfo.getExportName().c_str(), ModuleLinkErrorsToString(error).data());
            DEBUG_FUNCTION_LINE_ERR("%s", errMsg.c_str());
            OSFatal(errMsg.c_str());
        }
    }

    DEBUG_FUNCTION_LINE_VERBOSE("Resolve relocations without loading additional rpls");
    RelocationUtils::ResolveRelocations(moduleContainers, RelocationUtils::ExternalRPLLoadingStrategy::IGNORE_EXTERNAL_RPLS, usedRPls);

    // Build struct passed into INIT Hook
    persistedData.build(moduleContainers);

    for (auto &curModule : moduleContainers) {
        CallInitHooksForModule(curModule);
    }

    // This sets the custom rpl allocator so we properly load external RPL now via ExternalRPLLoadingStrategy::LOAD_EXTERNAL_RPLS
    HooksManagement::CallHook(moduleContainers, WUMS_HOOK_GET_CUSTOM_RPL_ALLOCATOR);

    // ========================================================================
    // PHASE 2: LOAD REGULAR MODULES. First get a module allocator which uses the mapped memory
    // then load/link "regular" modules
    // ========================================================================
    auto mappedAllocatorOpt = ModuleAllocator::FromMappedMemory(moduleContainers);
    if (!mappedAllocatorOpt) {
        DEBUG_FUNCTION_LINE_WARN("Failed to get memory module allocator, fallback to default one");
        mappedAllocatorOpt = defaultAllocator;
    }

    for (auto &[metaInfo, file] : regularModuleMetaInfoWithLoadedFiles) {
        ModuleLinkErrors error = ModuleLinkErrors::MODULE_LINK_ERROR_NONE;
        if (auto linkInfo = ModuleLinkInformationFactory::loadAndLink(std::span(file.data(), file.size()), *mappedAllocatorOpt, error)) {
            moduleContainers.emplace_back(std::move(metaInfo), std::move(*linkInfo));
        } else {
            auto errMsg = string_format("Failed to load/link module %s. %s", metaInfo.getExportName().c_str(), ModuleLinkErrorsToString(error).data());
            DEBUG_FUNCTION_LINE_ERR("%s", errMsg.c_str());
            OSFatal(errMsg.c_str());
        }
    }

    // Resolve relocations for regular modules as well, this time we can actually load external RPLs
    RelocationUtils::ResolveRelocations(moduleContainers, RelocationUtils::ExternalRPLLoadingStrategy::LOAD_EXTERNAL_RPLS, usedRPls);

    // rebuild struct passed into INIT Hook. The pointer passed to early module is still valid! (just hope they didn't make a copy)
    persistedData.build(moduleContainers);

    DEBUG_FUNCTION_LINE_VERBOSE("Call Relocations done hook to replace memory alloc/free functions");
    HooksManagement::CallHook(moduleContainers, WUMS_HOOK_RELOCATIONS_DONE);

    DEBUG_FUNCTION_LINE("Call init hooks for regular modules");
    for (auto &moduleContainer : moduleContainers) {
        if (moduleContainer.getMetaInformation().isInitBeforeRelocationDoneHook()) {
            // These have already been initialized
            continue;
        }
        CallInitHooksForModule(moduleContainer);
    }

    DEBUG_FUNCTION_LINE("%d modules loaded successfully", moduleContainers.size());

    return moduleContainers;
}

} // namespace WUMSLoader::Modules