#include "il2cpp_resolver.hpp"
#include <string>
#include <cstring>
#include <thread>
#include <chrono>
#include <fstream>
#include <filesystem>

#ifdef _WIN32
#include <windows.h>
#include <psapi.h>
#else
#include <dlfcn.h>
#include <link.h>
#include <unistd.h>
#endif

namespace devour {

#ifdef _WIN32
static HMODULE g_gameAsm = nullptr;

uintptr_t Il2CppResolver::gameAssemblyBase() {
    return reinterpret_cast<uintptr_t>(g_gameAsm);
}
size_t Il2CppResolver::gameAssemblySize() {
    if (!g_gameAsm) return 0;
    MODULEINFO mi{};
    if (GetModuleInformation(GetCurrentProcess(), g_gameAsm, &mi, sizeof(mi)))
        return mi.SizeOfImage;
    return 0;
}

void* Il2CppResolver::getProc(const char* name) {
    if (!g_gameAsm) return nullptr;
    return reinterpret_cast<void*>(GetProcAddress(g_gameAsm, name));
}

bool Il2CppResolver::init() {
    for (int i = 0; i < 600; ++i) { // 60s timeout
        g_gameAsm = GetModuleHandleA("GameAssembly.dll");
        if (g_gameAsm) return true;
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    return false;
}

static bool waitForDomain(int timeoutMs = 30000){
    auto domain_get = reinterpret_cast<void*(*)()>(Il2CppResolver::getProc("il2cpp_domain_get"));
    if(!domain_get) return false;
    for(int i=0;i<timeoutMs/100;i++){
        void* d = domain_get();
        if(d) return true;
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    return false;
}

// Simple pattern scan
std::vector<void*> Il2CppResolver::scanPattern(const char* pattern, const char* mask, uintptr_t start, size_t size) {
    std::vector<void*> out;
    size_t patLen = strlen(mask);
    for (size_t i = 0; i + patLen < size; ++i) {
        bool ok = true;
        for (size_t j = 0; j < patLen; ++j) {
            if (mask[j] == 'x' && reinterpret_cast<char*>(start + i)[j] != pattern[j]) { ok = false; break; }
        }
        if (ok) out.push_back(reinterpret_cast<void*>(start + i));
    }
    return out;
}

ResolvedMethod Il2CppResolver::resolve(const char* klassName, const char* methodName) {
    // Domain wait is now handled by installPatches once, not per-resolve to avoid 6*15s stalls
    // If domain not ready, fallback to string xref quickly
    auto domain_get_test = reinterpret_cast<void*(*)()>(getProc("il2cpp_domain_get"));
    if (domain_get_test && !domain_get_test()) {
        return resolveByStringXref(methodName);
    }
    auto il2cpp_domain_get = reinterpret_cast<void*(*)()>(getProc("il2cpp_domain_get"));
    auto il2cpp_thread_attach = reinterpret_cast<void*(*)(void*)>(getProc("il2cpp_thread_attach"));
    auto klass_from_name = reinterpret_cast<void*(*)(void*, const char*, const char*)>(getProc("il2cpp_class_from_name"));
    auto class_get_method = reinterpret_cast<void*(*)(void*, const char*)>(getProc("il2cpp_class_get_method_from_name"));
    auto domain_get_assemblies = reinterpret_cast<const void**(*)(void*, size_t*)>(getProc("il2cpp_domain_get_assemblies"));
    auto assembly_get_image = reinterpret_cast<void*(*)(void*)>(getProc("il2cpp_assembly_get_image"));

    if (!il2cpp_domain_get || !klass_from_name || !class_get_method || !domain_get_assemblies || !assembly_get_image) {
        return resolveByStringXref(methodName);
    }
    void* domain = il2cpp_domain_get();
    if (domain && il2cpp_thread_attach) il2cpp_thread_attach(domain);

    size_t nAssemblies = 0;
    const void** assemblies = nullptr;
    for(int w=0; w<100; w++){
        assemblies = domain_get_assemblies(domain, &nAssemblies);
        if(assemblies && nAssemblies>0) break;
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    if (!assemblies || nAssemblies==0) return resolveByStringXref(methodName);

    const char* candidates[] = {klassName, "NolanRankController", "LobbyController", "InventoryController", "GameController", "SurvivalLobbyController", "NolanRankController", nullptr};
    for (size_t ai = 0; ai < nAssemblies; ++ai) {
        void* image = assembly_get_image((void*)assemblies[ai]);
        if (!image) continue;
        for (int ci = 0; candidates[ci]; ++ci) {
            void* klass = klass_from_name(image, "", candidates[ci]);
            if (!klass) {
                // Try with no namespace but also try with "Devour" namespace?
                klass = klass_from_name(image, "Devour", candidates[ci]);
            }
            if (!klass) continue;
            void* method = class_get_method(klass, methodName);
            if (method) {
                void* ptr = *reinterpret_cast<void**>(reinterpret_cast<uint8_t*>(method) + 0x18);
                if (ptr) return {methodName, ptr};
            }
        }
    }
    // Brute force: enumerate all classes in all images for methodName (UpdateLocked* etc live in unknown class)
    auto image_get_class_count = reinterpret_cast<int(*)(void*)>(getProc("il2cpp_image_get_class_count"));
    auto image_get_class = reinterpret_cast<void*(*)(void*, int)>(getProc("il2cpp_image_get_class"));
    if (image_get_class_count && image_get_class) {
        for (size_t ai = 0; ai < nAssemblies; ++ai) {
            void* image = assembly_get_image((void*)assemblies[ai]);
            if (!image) continue;
            int count = image_get_class_count(image);
            for (int ci = 0; ci < count; ++ci) {
                void* klass = image_get_class(image, ci);
                if (!klass) continue;
                void* method = class_get_method(klass, methodName);
                if (method) {
                    void* ptr = *reinterpret_cast<void**>(reinterpret_cast<uint8_t*>(method) + 0x18);
                    if (ptr) return {methodName, ptr};
                }
            }
        }
    }
    return resolveByStringXref(methodName);
}

ResolvedMethod Il2CppResolver::resolveByStringXref(const char* str, int occurrence) {
    uintptr_t base = gameAssemblyBase();
    size_t sz = gameAssemblySize();
    if (!base || !sz) return {str, nullptr};

    // Fast find via memchr + memcmp — don't require null terminator (IsRobeUnlocked appears inside longer log string)
    size_t slen = strlen(str);
    uintptr_t strAddr = 0;
    const char* start = reinterpret_cast<const char*>(base);
    const char* end = start + sz - slen - 1;
    const char* p = start;
    while (p < end) {
        p = (const char*)memchr(p, str[0], end - p);
        if (!p) break;
        if (memcmp(p, str, slen)==0) {
            if (occurrence--==0) { strAddr = reinterpret_cast<uintptr_t>(p); break; }
        }
        p++;
    }
    if (!strAddr) return {str, nullptr};

    // Scan for 48 8D 0D ? ? ? ? (lea rcx, [rip+disp]) referencing strAddr
    // Simplified: look for 4-byte disp that resolves to strAddr
    for (size_t i = 0; i + 7 < sz; ++i) {
        uint8_t* p = reinterpret_cast<uint8_t*>(base + i);
        if (p[0] == 0x48 && p[1] == 0x8D && p[2] == 0x0D) {
            int32_t disp = *reinterpret_cast<int32_t*>(p + 3);
            uintptr_t target = base + i + 7 + disp;
            if (target == strAddr) {
                // Found xref — function start is previous 0xCC or 0x48 89 5C ? ... search back
                // Heuristic: walk back up to 0x200 for function prologue 40 53 / 48 83 EC
                for (int back = 0; back < 0x300; ++back) {
                    uint8_t* q = p - back;
                    if (q < reinterpret_cast<uint8_t*>(base)) break;
                    if (q[0] == 0x40 && q[1] == 0x53) return {str, q}; // push rbx
                    if (q[0] == 0x48 && q[1] == 0x83 && q[2] == 0xEC) return {str, q};
                }
                return {str, p}; // fallback
            }
        }
    }
    return {str, nullptr};
}

ResolvedMethod Il2CppResolver::isRobeUnlocked() { return resolve("NolanRankController", "IsRobeUnlocked"); }
ResolvedMethod Il2CppResolver::isCharacterUnlocked() { return resolve("NolanRankController", "IsCharacterUnlocked"); }
ResolvedMethod Il2CppResolver::hasOutfit() { return resolve("NolanRankController", "HasOutfit"); }
ResolvedMethod Il2CppResolver::hasPerk() { return resolve("NolanRankController", "HasPerk"); }
ResolvedMethod Il2CppResolver::hasPet() { return resolve("NolanRankController", "HasPet"); }
ResolvedMethod Il2CppResolver::updateLockedPets() { return resolve("SurvivalLobbyController", "UpdateLockedPets"); }
ResolvedMethod Il2CppResolver::updateLockedPerks() { return resolve("SurvivalLobbyController", "UpdateLockedPerks"); }
ResolvedMethod Il2CppResolver::updateLockedFlashlights() { return resolve("SurvivalLobbyController", "UpdateLockedFlashlights"); }
ResolvedMethod Il2CppResolver::updateLockedEmotes() { return resolve("SurvivalLobbyController", "UpdateLockedEmotes"); }
ResolvedMethod Il2CppResolver::updateLockedOutfits() { return resolve("SurvivalLobbyController", "UpdateLockedOutfits"); }
ResolvedMethod Il2CppResolver::updateLockedRobes() { return resolve("SurvivalLobbyController", "UpdateLockedRobes"); }
ResolvedMethod Il2CppResolver::waitForPurchasedPet() { return resolve("SurvivalLobbyController", "WaitForPurchasedPetValidation"); }

void Il2CppResolver::dumpDiagnostics(){
#ifdef _WIN32
    // Also log via OutputDebugString for Wine debug
    OutputDebugStringA("[Diag] dumpDiagnostics enter\n");
    try{
        char exe[MAX_PATH]; GetModuleFileNameA(nullptr, exe, MAX_PATH);
        char* slash = strrchr(exe, '\\'); if(!slash) slash = strrchr(exe, '/');
        std::string logPath = slash ? std::string(exe, slash - exe) + "\\DevourUnlocker.log" : "DevourUnlocker.log";
        std::ofstream f(logPath, std::ios::app);
        f << "[Diag] dumpDiagnostics ENTER\n"; f.flush();
        auto il2cpp_domain_get = reinterpret_cast<void*(*)()>(getProc("il2cpp_domain_get"));
        auto domain_get_assemblies = reinterpret_cast<const void**(*)(void*, size_t*)>(getProc("il2cpp_domain_get_assemblies"));
        auto assembly_get_image = reinterpret_cast<void*(*)(void*)>(getProc("il2cpp_assembly_get_image"));
        auto image_get_class_count = reinterpret_cast<int(*)(void*)>(getProc("il2cpp_image_get_class_count"));
        auto image_get_class = reinterpret_cast<void*(*)(void*, int)>(getProc("il2cpp_image_get_class"));
        auto class_get_name = reinterpret_cast<const char*(*)(void*)>(getProc("il2cpp_class_get_name"));
        auto class_get_namespace = reinterpret_cast<const char*(*)(void*)>(getProc("il2cpp_class_get_namespace"));
        auto class_get_method_count = reinterpret_cast<int(*)(void*)>(getProc("il2cpp_class_num_methods"));
        auto class_get_methods = reinterpret_cast<void**(*)(void*, void**)>(getProc("il2cpp_class_get_methods"));
        auto method_get_name = reinterpret_cast<const char*(*)(void*)>(getProc("il2cpp_method_get_name"));
        if(!il2cpp_domain_get || !domain_get_assemblies || !assembly_get_image || !image_get_class_count || !image_get_class) {
            f << "[Diag] missing IL2CPP exports for dump\n"; return;
        }
        void* domain = il2cpp_domain_get();
        if(!domain){ f << "[Diag] domain null\n"; return; }
        size_t nAssemblies=0; const void** assemblies = domain_get_assemblies(domain, &nAssemblies);
        f << "[Diag] assemblies=" << nAssemblies << "\n";
            for(size_t ai=0; ai<nAssemblies && ai<32; ++ai){
            void* image = assembly_get_image((void*)assemblies[ai]);
            if(!image) continue;
            auto image_get_name = reinterpret_cast<const char*(*)(void*)>(getProc("il2cpp_image_get_name"));
            const char* iname = image_get_name ? image_get_name(image) : "?";
            int cnt = image_get_class_count(image);
            f << "[Diag] image " << iname << " classes=" << cnt << " ptr=" << image << "\n";
            // Log first 10 class names of this image as sample
            for(int ci=0; ci<cnt && ci<10; ++ci){
                void* klass = image_get_class(image, ci);
                if(!klass) continue;
                const char* kname = class_get_name ? class_get_name(klass) : "?";
                const char* ns = class_get_namespace ? class_get_namespace(klass) : "";
                f << "    sample class " << ns << "." << kname << "\n";
            }
            // Now search for UpdateLocked etc in all classes of this image — also dump Nolan/Survival specifically
            for(int ci=0; ci<cnt && ci<2000; ++ci){
                void* klass = image_get_class(image, ci);
                if(!klass) continue;
                const char* kname = class_get_name ? class_get_name(klass) : "?";
                std::string kn=kname?kname:"";
                bool isTargetKlass = kn=="NolanRankController" || kn=="SurvivalLobbyController" || kn=="LobbyController";
                auto class_get_method = reinterpret_cast<void*(*)(void*, const char*)>(getProc("il2cpp_class_get_method_from_name"));
                if(isTargetKlass && class_get_method){
                    f << "  DIAG target klass " << kname << " in " << iname << "\n";
                    const char* targets[]={"UpdateLockedPets","UpdateLockedPerks","IsRobeUnlocked","IsCharacterUnlocked","HasPerk","HasOutfit","UpdateLockedRobes",nullptr};
                    for(int ti=0; targets[ti]; ++ti){
                        void* m = class_get_method(klass, targets[ti]);
                        f << "    try " << targets[ti] << " -> " << m;
                        if(m){
                            void* ptr = *reinterpret_cast<void**>(reinterpret_cast<uint8_t*>(m) + 0x18);
                            f << " ptr=" << ptr;
                            auto method_get_name2 = reinterpret_cast<const char*(*)(void*)>(getProc("il2cpp_method_get_name"));
                            if(method_get_name2) f << " name=" << method_get_name2(m);
                        }
                        f << "\n";
                    }
                    // Dump all methods of this klass
                    auto class_get_methods2 = reinterpret_cast<const void**(*)(void*, void**)>(getProc("il2cpp_class_get_methods"));
                    auto method_get_name2 = reinterpret_cast<const char*(*)(void*)>(getProc("il2cpp_method_get_name"));
                    if(class_get_methods2 && method_get_name2){
                        void* iter2=nullptr;
                        int mc2=0;
                        f << "    all methods of " << kname << ":\n";
                        while(true){
                            const void** methods = (const void**)class_get_methods2(klass, &iter2);
                            if(!methods || !*methods) break;
                            const char* mn = method_get_name2((void*)*methods);
                            f << "      " << (mn?mn:"<null>") << "\n";
                            if(++mc2>100) break;
                        }
                    }
                }
                if(class_get_method){
                    const char* targets[]={"UpdateLockedPets","UpdateLockedPerks","IsRobeUnlocked","IsCharacterUnlocked","HasPerk","HasOutfit",nullptr};
                    for(int ti=0; targets[ti]; ++ti){
                        void* m = class_get_method(klass, targets[ti]);
                        if(m) {
                            void* ptr = *reinterpret_cast<void**>(reinterpret_cast<uint8_t*>(m) + 0x18);
                            f << "  FOUND " << kname << "." << targets[ti] << " @ " << m << " -> " << ptr << " (image " << iname << ")\n";
                        }
                    }
                }
            }
        }
        f << "[Diag] dump done\n";
    }catch(...){}
#endif
}

#else // Linux stubs

uintptr_t Il2CppResolver::gameAssemblyBase() { return 0; }
size_t Il2CppResolver::gameAssemblySize() { return 0; }
void* Il2CppResolver::getProc(const char*) { return nullptr; }
bool Il2CppResolver::init() { return false; }
std::vector<void*> Il2CppResolver::scanPattern(const char*, const char*, uintptr_t, size_t) { return {}; }
ResolvedMethod Il2CppResolver::resolve(const char*, const char* m) { return {m, nullptr}; }
ResolvedMethod Il2CppResolver::resolveByStringXref(const char* s, int) { return {s, nullptr}; }
ResolvedMethod Il2CppResolver::isRobeUnlocked() { return {"IsRobeUnlocked", nullptr}; }
ResolvedMethod Il2CppResolver::isCharacterUnlocked() { return {"IsCharacterUnlocked", nullptr}; }
ResolvedMethod Il2CppResolver::hasOutfit() { return {"HasOutfit", nullptr}; }
ResolvedMethod Il2CppResolver::hasPerk() { return {"HasPerk", nullptr}; }
ResolvedMethod Il2CppResolver::hasPet() { return {"HasPet", nullptr}; }
ResolvedMethod Il2CppResolver::updateLockedPets() { return {"UpdateLockedPets", nullptr}; }
ResolvedMethod Il2CppResolver::updateLockedPerks() { return {"UpdateLockedPerks", nullptr}; }
ResolvedMethod Il2CppResolver::updateLockedFlashlights() { return {"UpdateLockedFlashlights", nullptr}; }
ResolvedMethod Il2CppResolver::updateLockedEmotes() { return {"UpdateLockedEmotes", nullptr}; }
ResolvedMethod Il2CppResolver::updateLockedOutfits() { return {"UpdateLockedOutfits", nullptr}; }
ResolvedMethod Il2CppResolver::updateLockedRobes() { return {"UpdateLockedRobes", nullptr}; }
ResolvedMethod Il2CppResolver::waitForPurchasedPet() { return {"WaitForPurchasedPetValidation", nullptr}; }

#endif

} // namespace devour