#include "Injector.h"
#include "SkeetInjector.h"
#include <windows.h>
#include <cstdio>
#include <thread>
#include <chrono>

// ─── 单例 ──────────────────────────────────────────────────────

injector* injector::Get()
{
    static injector instance;
    return &instance;
}

injector::~injector()
{
    if (m_thread.joinable())
        m_thread.join();
}

// ─── 公开接口 ─────────────────────────────────────────────────

bool injector::Launch()
{
    if (m_running)
    {
        printf("[Injector] > warn: injection already in progress\n");
        return false;
    }

    m_running = true;
    m_success = false;

    m_thread = std::thread(&injector::InjectionThread, this);

    printf("[Injector] > injection thread launched\n");
    return true;
}

bool injector::IsRunning() const
{
    return m_running;
}

bool injector::HasSucceeded() const
{
    return m_success;
}

void injector::Wait()
{
    if (m_thread.joinable())
        m_thread.join();
}

// ─── 内部实现 ─────────────────────────────────────────────────

std::string injector::WriteInjectorToTemp()
{
    char tempPath[MAX_PATH];
    DWORD len = GetTempPathA(sizeof(tempPath), tempPath);
    if (len == 0 || len > sizeof(tempPath))
    {
        printf("[Injector] > err: GetTempPathA failed (0x%lx)\n", GetLastError());
        return {};
    }

    std::string exePath = std::string(tempPath) + "SkeetInjector.exe";

    HANDLE hFile = CreateFileA(
        exePath.c_str(),
        GENERIC_WRITE,
        0,
        nullptr,
        CREATE_ALWAYS,
        FILE_ATTRIBUTE_NORMAL,
        nullptr
    );

    if (hFile == INVALID_HANDLE_VALUE)
    {
        printf("[Injector] > err: cannot create %s (0x%lx)\n", exePath.c_str(), GetLastError());
        return {};
    }

    DWORD written = 0;
    BOOL ok = WriteFile(hFile, SkeetInjector_exe, SkeetInjector_exe_len, &written, nullptr);
    CloseHandle(hFile);

    if (!ok || written != SkeetInjector_exe_len)
    {
        printf("[Injector] > err: write failed for %s (wrote %lu/%u)\n",
               exePath.c_str(), written, SkeetInjector_exe_len);
        DeleteFileA(exePath.c_str());
        return {};
    }

    printf("[Injector] > written: %s (%u bytes)\n", exePath.c_str(), SkeetInjector_exe_len);
    return exePath;
}

bool injector::RunInjector(const std::string& exePath)
{
    // 使用 ShellExecuteExA + runas 以管理员权限启动（UAC 提权）
    SHELLEXECUTEINFOA sei = {};
    sei.cbSize = sizeof(sei);
    sei.fMask = SEE_MASK_NOCLOSEPROCESS | SEE_MASK_FLAG_NO_UI;
    sei.hwnd = nullptr;
    sei.lpVerb = "runas";                    // 触发 UAC 管理员提权
    sei.lpFile = exePath.c_str();
    sei.lpParameters = nullptr;
    sei.lpDirectory = nullptr;
    sei.nShow = SW_HIDE;

    printf("[Injector] > executing (runas): %s\n", exePath.c_str());

    if (!ShellExecuteExA(&sei))
    {
        DWORD err = GetLastError();
        if (err == ERROR_CANCELLED)
            printf("[Injector] > err: user cancelled UAC prompt\n");
        else
            printf("[Injector] > err: ShellExecuteExA failed (0x%lx)\n", err);
        return false;
    }

    // 等待注入器完成
    if (sei.hProcess)
    {
        printf("[Injector] > process started (pid: %lu), waiting...\n", GetProcessId(sei.hProcess));
        WaitForSingleObject(sei.hProcess, INFINITE);

        DWORD exitCode = 0;
        GetExitCodeProcess(sei.hProcess, &exitCode);
        CloseHandle(sei.hProcess);

        printf("[Injector] > process exited with code: %lu\n", exitCode);
    }

    return true;
}

void injector::CleanupTempFile(const std::string& exePath)
{
    if (exePath.empty())
        return;

    // 短暂延迟，确保进程完全释放文件句柄
    std::this_thread::sleep_for(std::chrono::milliseconds(200));

    if (DeleteFileA(exePath.c_str()))
    {
        printf("[Injector] > cleaned up: %s\n", exePath.c_str());
    }
    else
    {
        // 可能仍在使用中，标记为重启后删除
        MoveFileExA(exePath.c_str(), nullptr, MOVEFILE_DELAY_UNTIL_REBOOT);
        printf("[Injector] > cleanup deferred (reboot): %s\n", exePath.c_str());
    }
}

// ─── 注入线程 ─────────────────────────────────────────────────

void injector::InjectionThread()
{
    printf("[Injector] > injection thread started\n");

    // 1. 将嵌入的 SkeetInjector.exe 写入临时目录
    std::string exePath = WriteInjectorToTemp();
    if (exePath.empty())
    {
        printf("[Injector] > err: failed to write injector to temp\n");
        m_running = false;
        return;
    }

    // 2. 执行注入器
    bool ran = RunInjector(exePath);
    if (!ran)
    {
        printf("[Injector] > err: injector execution failed\n");
        CleanupTempFile(exePath);
        m_running = false;
        return;
    }

    // 3. 清理临时文件
    CleanupTempFile(exePath);

    // 4. 标记完成
    m_success = true;
    m_running = false;

    printf("[Injector] > injection thread completed successfully\n");
}
