#include "pch.h"
#include "Il2CppApi.h"
#include "Il2CppHelper.h"
#include "Il2CppMetadata.h"
#include "ProtoMetadata.h"
#include "ScriptJson.h"
#include "ProgressReporter.h"

#include <functional>
#include <chrono>
#include <iomanip>
#include <fstream>
#define DEBUG_VER
//#define RELEASE_VER

static DWORD WINAPI UnloadDll(LPVOID d) {

    std::cout << "Dll Unloaded\n";

    FreeConsole();
    FreeLibraryAndExitThread((HMODULE)d, 0);
}

BOOL WINAPI DllMain(const HMODULE instance, DWORD reason)
{
    switch (reason)
    {
    case DLL_PROCESS_ATTACH:
        DisableThreadLibraryCalls(instance);
        CreateThread(nullptr, 0, MainThread, instance, 0, nullptr);
        break;
    case DLL_PROCESS_DETACH:

        break;
    }


    return true;

}

DWORD WINAPI MainThread(LPVOID lpReserved) {
    // Setup console for UTF-8 output
    AllocConsole();
    SetConsoleOutputCP(CP_UTF8);
    SetConsoleCP(CP_UTF8);
    FILE* console;
    freopen_s(&console, "CONOUT$", "w", stdout);
    freopen_s(&console, "CONIN$", "r", stdin);

    // Initial messages
    std::cout << R"(
███████╗ ██████╗██╗     ██╗██████╗ ███████╗███████╗██████╗ ██╗   ██╗███╗   ███╗██████╗ ███████╗██████╗ 
██╔════╝██╔════╝██║     ██║██╔══██╗██╔════╝██╔════╝██╔══██╗██║   ██║████╗ ████║██╔══██╗██╔════╝██╔══██╗
█████╗  ██║     ██║     ██║██████╔╝███████╗█████╗  ██║  ██║██║   ██║██╔████╔██║██████╔╝█████╗  ██████╔╝
██╔══╝  ██║     ██║     ██║██╔═══╝ ╚════██║██╔══╝  ██║  ██║██║   ██║██║╚██╔╝██║██╔═══╝ ██╔══╝  ██╔══██╗
███████╗╚██████╗███████╗██║██║     ███████║███████╗██████╔╝╚██████╔╝██║ ╚═╝ ██║██║     ███████╗██║  ██║
╚══════╝ ╚═════╝╚══════╝╚═╝╚═╝     ╚══════╝╚══════╝╚═════╝  ╚═════╝ ╚═╝     ╚═╝╚═╝     ╚══════╝╚═╝  ╚═╝
                                                                                                       )" << "\n";
    std::cout << "Waiting for il2cpp initialization... (Press Enter when ready)\n";
    while (!(GetAsyncKeyState(VK_RETURN) & 0x8000)) {
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }

    
    std::cout << "\n=== Starting Dump Process ===\n";
    Il2CppApi::GCDisable();

    std::ofstream dumpFile("dump.cs");
    if (!dumpFile.is_open()) {
        std::cerr << "FATAL ERROR: Failed to create output file!\n";
        return 1;
    }

    size_t count = 0;
    Il2CppAssembly** assemblies = Il2CppApi::DomainGetAssemblies(&count);

    ProgressReporter dumpReporter(count, 50);
    for (size_t index = 0; index < count; ++index) {
        Il2CppAssembly* assembly = assemblies[index];
        if (!assembly) continue;

        auto image = Il2CppApi::AssemblyGetImage(assembly);
        auto assemblyName = Il2CppApi::AssemblyGetName(assembly);
        
        auto classCount = Il2CppApi::ImageGetClassCount(image);
        for(int i = 0; i < classCount; i++) {
            Il2CppClass* klass = Il2CppApi::ImageGetClass(image, i);
            if(!klass) continue;

            auto namespaze = Il2CppApi::ClassGetNamespace(klass);
            auto className = Il2CppApi::ClassGetName(klass);

            try {
                dumpFile << "// Module: " << assemblyName << "\n";
                dumpFile << "// Namespace: " << namespaze << "\n";
                if(!std::string(namespaze).empty()) dumpFile << "// FullName: " << namespaze << "." << className << "\n";
                else dumpFile << "// FullName: " << className << "\n";

                std::string classDump = Il2CppMetadata::DumpClass(klass);
                if(!classDump.empty()) dumpFile << classDump << "\n\n";
            }
            catch (const std::exception& e) {
                std::cout << "An exception happened! " << e.what() << "\n";
            }
        }

        dumpReporter.update(index);
    }

    dumpReporter.finish();
    dumpFile.close();

    std::cout << "\n=== Generating Proto File ===\n";
    Il2CppApi::ThreadAttach(Il2CppApi::DomainGet());

    std::ofstream protoFile("main.proto");
    if (!protoFile.is_open()) {
        std::cerr << "FATAL ERROR: Failed to create proto file!\n";
        return 1;
    }

    const Il2CppAssembly* protoAssembly = Il2CppApi::DomainAssemblyOpen("Assembly-CSharp.dll");
    const Il2CppImage* protoImage = Il2CppApi::AssemblyGetImage(protoAssembly);

    auto protoCount = Il2CppApi::ImageGetClassCount(protoImage);

    ProgressReporter protoReporter(protoCount, 50);
    for(uint32_t i = 0; i < protoCount; i++) {
        Il2CppClass* klass = Il2CppApi::ImageGetClass(protoImage, i);
        if(!klass) continue;

        auto parent = Il2CppApi::ClassGetParent(klass);
        if(!parent) continue;

        auto parent_name = Il2CppApi::ClassGetName(parent);
        if(strcmp(parent_name, PROTO_CLASS) != 0) continue;

        try {
            std::string protoDump = ProtoMetadata::DumpProto(klass);
            if(!protoDump.empty()) protoFile << protoDump << "\n\n";
        }
        catch (const std::exception& e) {
            std::cout << "An exception happened! " << e.what() << "\n";
        }

        protoReporter.update(i);
    }

    protoReporter.finish();
    protoFile.close();

    // Generate script.json (optional, keep if needed)
    std::cout << "\n=== Generating Script JSON ===\n";
    std::ofstream jsonFile("script.json");
    if (jsonFile.is_open()) {
        jsonFile << ScriptJsonGenerator::GenerateJson();
        jsonFile.close();
        std::cout << "JSON metadata saved to script.json\n";
    }
    else {
        std::cerr << "Failed to create script.json\n";
    }
    
    std::cout << "\nDumping completed\n";
    UnloadDll(lpReserved);
    return 0;
}