// Native input-value constructor for the local development bridge.
// WitcherScript exposes SInputAction fields as const. We construct the native
// value and then pass it through the game's existing script input handlers.
#include <winsock2.h>
#include <windows.h>
#include <xinput.h>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <atomic>
#include <string>
#include <WitcherNativeBridgeApi.h>

using GetStateFn = DWORD (WINAPI*)(DWORD, XINPUT_STATE*);
using BindFn = int (WSAAPI*)(SOCKET, const sockaddr*, int);
static GetStateFn originalGetState = nullptr;
static BindFn originalBind = nullptr;
static SRWLOCK padLock = SRWLOCK_INIT;
static XINPUT_STATE pad{};
static std::atomic<bool> agentControl{false};
static ULONGLONG expires = 0;
static LONG polls = 0;
struct ImportHook { void** slot; void* original; void* replacement; };
static ImportHook hooks[2]{};
static int hookCount = 0;

static void RestoreImports() {
    for (int i = hookCount - 1; i >= 0; --i) {
        auto& hook = hooks[i]; DWORD protection;
        if (*hook.slot != hook.replacement) continue;
        if (VirtualProtect(hook.slot, sizeof(void*), PAGE_READWRITE, &protection)) {
            InterlockedCompareExchangePointer(reinterpret_cast<void* volatile*>(hook.slot), hook.original, hook.replacement);
            DWORD unused; VirtualProtect(hook.slot, sizeof(void*), protection, &unused);
        }
    }
}

static DWORD WINAPI GetState(DWORD index, XINPUT_STATE* out) {
    // Human mode never waits on our lock or adds polling/timers to real input.
    if (index != 0 || !out || !agentControl.load(std::memory_order_acquire))
        return originalGetState(index, out);
    AcquireSRWLockExclusive(&padLock);
    if (!agentControl.load(std::memory_order_relaxed)) {
        ReleaseSRWLockExclusive(&padLock);
        return originalGetState(index, out);
    }
    InterlockedIncrement(&polls);
    if (expires && GetTickCount64() >= expires) {
        pad.Gamepad = {};
        ++pad.dwPacketNumber;
        expires = 0;
    }
    *out = pad;
    ReleaseSRWLockExclusive(&padLock);
    return ERROR_SUCCESS;
}

static void SetAgentControl(void*, void* frame, void*) {
    const int value = WNB_ReadIntParameter(frame);
    AcquireSRWLockExclusive(&padLock);
    pad.Gamepad = {}; ++pad.dwPacketNumber; expires = 0;
    agentControl.store(value == 1, std::memory_order_release);
    ReleaseSRWLockExclusive(&padLock);
    WNB_ReturnVoid(frame);
}

static void AgentControl(void*, void* frame, void* result) {
    WNB_ReturnBool(frame, result, agentControl.load(std::memory_order_acquire));
}

// The built-in debugger binds to all interfaces. Restrict its known TCP port
// in this process; other game/Steam socket binds pass through unchanged.
static int WSAAPI Bind(SOCKET socket, const sockaddr* address, int length) {
    if (address && length >= static_cast<int>(sizeof(sockaddr_in)) && address->sa_family == AF_INET) {
        sockaddr_in local = *reinterpret_cast<const sockaddr_in*>(address);
        if (local.sin_port == htons(37001)) {
            local.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
            return originalBind(socket, reinterpret_cast<const sockaddr*>(&local), sizeof(local));
        }
    }
    return originalBind(socket, address, length);
}

static bool HookImport(const char* dll, const char* name, WORD ordinal, void* replacement, void** original) {
    auto base = reinterpret_cast<unsigned char*>(GetModuleHandleW(nullptr));
    auto dos = reinterpret_cast<IMAGE_DOS_HEADER*>(base);
    if (dos->e_magic != IMAGE_DOS_SIGNATURE) return false;
    auto pe = reinterpret_cast<IMAGE_NT_HEADERS64*>(base + dos->e_lfanew);
    if (pe->Signature != IMAGE_NT_SIGNATURE || pe->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) return false;
    auto directory = pe->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (!directory.VirtualAddress || directory.VirtualAddress >= pe->OptionalHeader.SizeOfImage) return false;
    auto imports = reinterpret_cast<IMAGE_IMPORT_DESCRIPTOR*>(base + directory.VirtualAddress);
    for (; imports->Name; ++imports) {
        if (_stricmp(reinterpret_cast<char*>(base + imports->Name), dll)) continue;
        if (!imports->OriginalFirstThunk) return false;
        auto lookup = reinterpret_cast<IMAGE_THUNK_DATA64*>(base + imports->OriginalFirstThunk);
        auto target = reinterpret_cast<IMAGE_THUNK_DATA64*>(base + imports->FirstThunk);
        for (; lookup->u1.AddressOfData; ++lookup, ++target) {
            bool match = IMAGE_SNAP_BY_ORDINAL64(lookup->u1.Ordinal)
                ? ordinal && IMAGE_ORDINAL64(lookup->u1.Ordinal) == ordinal
                : name && !std::strcmp(reinterpret_cast<IMAGE_IMPORT_BY_NAME*>(base + lookup->u1.AddressOfData)->Name, name);
            if (!match) continue;
            DWORD protection;
            if (!VirtualProtect(&target->u1.Function, sizeof(void*), PAGE_READWRITE, &protection)) return false;
            *original = reinterpret_cast<void*>(target->u1.Function);
            hooks[hookCount++] = {reinterpret_cast<void**>(&target->u1.Function), *original, replacement};
            InterlockedExchangePointer(reinterpret_cast<void* volatile*>(&target->u1.Function), replacement);
            DWORD unused;
            VirtualProtect(&target->u1.Function, sizeof(void*), protection, &unused);
            return true;
        }
    }
    return false;
}

static SHORT Axis(float value) {
    if (!std::isfinite(value)) return 0;
    return static_cast<SHORT>(std::clamp(value, -1.f, 1.f) * 32767.f);
}

static void PadState(void*, void* frame, void* result) {
    XINPUT_GAMEPAD next{};
    next.sThumbLX = Axis(WNB_ReadFloatParameter(frame));
    next.sThumbLY = Axis(WNB_ReadFloatParameter(frame));
    next.sThumbRX = Axis(WNB_ReadFloatParameter(frame));
    next.sThumbRY = Axis(WNB_ReadFloatParameter(frame));
    const int buttons = WNB_ReadIntParameter(frame);
    const int left = WNB_ReadIntParameter(frame), right = WNB_ReadIntParameter(frame);
    const int lease = WNB_ReadIntParameter(frame);
    if (!originalGetState || buttons < 0 || buttons > 65535 || (buttons & 0x0c00) ||
        left < 0 || left > 255 || right < 0 || right > 255 || lease < 1 || lease > 10000) {
        WNB_ReturnBool(frame, result, false); return;
    }
    next.wButtons = static_cast<WORD>(buttons);
    next.bLeftTrigger = static_cast<BYTE>(left); next.bRightTrigger = static_cast<BYTE>(right);
    AcquireSRWLockExclusive(&padLock);
    if (!agentControl.load(std::memory_order_relaxed)) {
        ReleaseSRWLockExclusive(&padLock);
        WNB_ReturnBool(frame, result, false); return;
    }
    pad.Gamepad = next; ++pad.dwPacketNumber;
    expires = GetTickCount64() + lease;
    ReleaseSRWLockExclusive(&padLock);
    WNB_ReturnBool(frame, result, true);
}

static void PadRelease(void*, void* frame, void*) {
    AcquireSRWLockExclusive(&padLock);
    pad.Gamepad = {}; ++pad.dwPacketNumber; expires = 0;
    ReleaseSRWLockExclusive(&padLock);
    WNB_ReturnVoid(frame);
}

static void PadPolls(void*, void* frame, void* result) {
    WNB_ReturnInt(frame, result, InterlockedCompareExchange(&polls, 0, 0));
}

// A script guard checks this on each natural invocation, even if the owner is
// gone or simulation ticks are frozen. No engine offsets or new import hooks.
static std::atomic<ULONGLONG> requestLeaseExpires{0};
static void RequestLease(void*, void* frame, void* result) {
    const int milliseconds = WNB_ReadIntParameter(frame);
    if (milliseconds < 0 || milliseconds > 60000) {
        WNB_ReturnBool(frame, result, false); return;
    }
    requestLeaseExpires = milliseconds ? GetTickCount64() + milliseconds : 0;
    WNB_ReturnBool(frame, result, true);
}
static void RequestLeaseAlive(void*, void* frame, void* result) {
    const auto deadline = requestLeaseExpires.load();
    WNB_ReturnBool(frame, result, deadline && GetTickCount64() < deadline);
}

// 5.01 strings use UTF-8. Encode JSON without script string helpers, which
// normalize some control characters (StrChar(13) became a space in live tests).
static void JsonQuote(void*, void* frame, void* result) {
    WNB_String input{};
    WNB_ReadStringParameter(frame, &input);
    std::wstring value;
    if (input.data && input.size > 1 && input.size < 1048576) {
        auto bytes = reinterpret_cast<const char*>(input.data);
        int size = static_cast<int>(input.size - 1);
        int n = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, bytes, size, nullptr, 0);
        if (n > 0) {
            value.resize(n);
            MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, bytes, size, value.data(), n);
        }
    }
    std::wstring out = L"\"";
    constexpr wchar_t hex[] = L"0123456789abcdef";
    for (wchar_t c : value) {
        if (c == L'"' || c == L'\\') { out += L'\\'; out += c; }
        else if (c < 32 || c >= 127) {
            out += L"\\u";
            out += hex[(c >> 12) & 15]; out += hex[(c >> 8) & 15];
            out += hex[(c >> 4) & 15]; out += hex[c & 15];
        } else out += c;
    }
    out += L'"';
    WNB_ReturnString(frame, result, out.c_str(), static_cast<uint32_t>(out.size()));
}

struct InputAction {
    WNB_Name name;
    float value;
    float previous;
    bool patterned;
    char padding[3];
    WNB_Name pattern;
};
static_assert(sizeof(InputAction) == 20);

static void MakeInput(void*, void* frame, void* result) {
    InputAction action{};
    action.name = WNB_ReadNameParameter(frame);
    action.value = WNB_ReadFloatParameter(frame);
    action.previous = WNB_ReadFloatParameter(frame);
    if (result) std::memcpy(result, &action, sizeof(action));
    WNB_ReturnVoid(frame);
}

BOOL APIENTRY DllMain(HMODULE module, DWORD reason, LPVOID) {
    if (reason == DLL_PROCESS_ATTACH) {
        DisableThreadLibraryCalls(module);
        if (WNB_GetApiVersion() != 2) return FALSE;
        if (!HookImport("XINPUT9_1_0.dll", "XInputGetState", 0, reinterpret_cast<void*>(&GetState), reinterpret_cast<void**>(&originalGetState))) return FALSE;
        if (!HookImport("WS2_32.dll", "bind", 2, reinterpret_cast<void*>(&Bind), reinterpret_cast<void**>(&originalBind))) {
            RestoreImports(); return FALSE;
        }
        // Remain loaded once callbacks have been queued: unloading would leave
        // the registration queue pointing into unmapped plugin memory.
        bool registered = WNB_RegisterNative("PD_MakeInput", &MakeInput);
        registered = WNB_RegisterNative("PD_SetAgentControl", &SetAgentControl) && registered;
        registered = WNB_RegisterNative("PD_AgentControl", &AgentControl) && registered;
        registered = WNB_RegisterNative("PD_PadState", &PadState) && registered;
        registered = WNB_RegisterNative("PD_PadRelease", &PadRelease) && registered;
        registered = WNB_RegisterNative("PD_PadPolls", &PadPolls) && registered;
        registered = WNB_RegisterNative("PD_JsonQuote", &JsonQuote) && registered;
        registered = WNB_RegisterNative("PD_RequestLease", &RequestLease) && registered;
        registered = WNB_RegisterNative("PD_RequestLeaseAlive", &RequestLeaseAlive) && registered;
        if (!registered) OutputDebugStringA("PathInput: native registration failed; inspect WitcherNativeBridge.log");
    } else if (reason == DLL_PROCESS_DETACH) {
        RestoreImports();
    }
    return TRUE;
}