// Inject.cpp - Modern D3D9 implementation of Inject.DLL for Decal
//
// This DLL gets injected into the Asheron's Call game client process.
// It hooks the D3D9 rendering pipeline to enable Decal's overlay system,
// manages the plugin lifecycle, and provides vtable hooking utilities.

#define WIN32_LEAN_AND_MEAN
#define INJECT_EXPORTS
#include <windows.h>
#include <d3d9.h>
#include <tchar.h>
#include <string>

#include "Inject.h"

#pragma comment(lib, "d3d9.lib")

// ---------------------------------------------------------------------------
// Shared data section - accessible across processes using this DLL
// ---------------------------------------------------------------------------
#pragma data_seg(".InjectDll")
static HHOOK g_hHook = NULL;
#pragma data_seg()
#pragma comment(linker, "/SECTION:.InjectDll,RWS")

// ---------------------------------------------------------------------------
// Globals
// ---------------------------------------------------------------------------
static HINSTANCE g_hInstance = NULL;
static HWND g_hGameWindow = NULL;
static IDirect3DDevice9* g_pDevice = NULL;
static WNDPROC g_pfnOrigWndProc = NULL;
static bool g_bInitialized = false;
static bool g_bContainerMode = false;

// COM interface pointers for managed Decal core
typedef HRESULT(__stdcall* PFN_CO_CREATE)(REFCLSID, LPUNKNOWN, DWORD, REFIID, LPVOID*);
static IUnknown* g_pDecalCore = NULL;

// Original D3D9 method pointers (saved before hooking)
extern "C" INJECT_API void* BeginSceneO = NULL;
extern "C" INJECT_API void* EndSceneO = NULL;

// Saved vtable entries for unhooking
struct VTablePatch
{
    void* pInterface;
    int nIndex;
    void* pOriginal;
};
static VTablePatch g_patches[64];
static int g_nPatchCount = 0;

// Registry path constants
static const TCHAR* REG_DECAL_KEY = _T("SOFTWARE\\Decal");
static const TCHAR* REG_AGENT_KEY = _T("SOFTWARE\\Decal\\Agent");

// ---------------------------------------------------------------------------
// Path mapping
// ---------------------------------------------------------------------------
static std::basic_string<TCHAR> GetAgentPath()
{
    HKEY hKey;
    TCHAR szPath[MAX_PATH] = {};
    DWORD dwSize = MAX_PATH * sizeof(TCHAR);

    if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, REG_AGENT_KEY, 0, KEY_READ, &hKey) == ERROR_SUCCESS)
    {
        RegQueryValueEx(hKey, _T("AgentPath"), NULL, NULL, (LPBYTE)szPath, &dwSize);
        RegCloseKey(hKey);
    }

    return szPath;
}

extern "C" INJECT_API LPTSTR __stdcall InjectMapPath(int pathType, LPCTSTR szFilename, LPTSTR szBuffer)
{
    std::basic_string<TCHAR> base;

    if (pathType == 0) // eInjectPathDatFile - AC installation directory
    {
        HKEY hKey;
        TCHAR szPath[MAX_PATH] = {};
        DWORD dwSize = MAX_PATH * sizeof(TCHAR);

        if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Microsoft\\Microsoft Games\\Asheron's Call\\1.00"),
                         0, KEY_READ, &hKey) == ERROR_SUCCESS)
        {
            RegQueryValueEx(hKey, _T("GamePath"), NULL, NULL, (LPBYTE)szPath, &dwSize);
            RegCloseKey(hKey);
        }
        base = szPath;
    }
    else // eInjectPathAgent - Decal agent directory
    {
        base = GetAgentPath();
    }

    // Ensure trailing backslash
    if (!base.empty() && base.back() != _T('\\'))
        base += _T('\\');

    _tcscpy(szBuffer, (base + szFilename).c_str());
    return szBuffer;
}

// ---------------------------------------------------------------------------
// VTable hooking utilities
// ---------------------------------------------------------------------------
extern "C" INJECT_API void* __stdcall GetVTableEntry(void* pInterface, int nIndex)
{
    if (!pInterface)
        return NULL;

    // COM interface vtable: first DWORD at *pInterface is pointer to vtable array
    void** vtable = *(void***)pInterface;
    return vtable[nIndex];
}

extern "C" INJECT_API BOOL __stdcall HookVTable(void* pInterface, int nIndex, void* pNewFunction)
{
    if (!pInterface || !pNewFunction)
        return FALSE;

    void** vtable = *(void***)pInterface;
    void* pOriginal = vtable[nIndex];

    // Change page protection to allow write
    DWORD dwOldProtect;
    if (!VirtualProtect(&vtable[nIndex], sizeof(void*), PAGE_READWRITE, &dwOldProtect))
        return FALSE;

    // Patch the vtable entry
    vtable[nIndex] = pNewFunction;

    // Restore protection
    VirtualProtect(&vtable[nIndex], sizeof(void*), dwOldProtect, &dwOldProtect);

    // Save for unhooking
    if (g_nPatchCount < _countof(g_patches))
    {
        g_patches[g_nPatchCount].pInterface = pInterface;
        g_patches[g_nPatchCount].nIndex = nIndex;
        g_patches[g_nPatchCount].pOriginal = pOriginal;
        g_nPatchCount++;
    }

    return TRUE;
}

extern "C" INJECT_API BOOL __stdcall UnhookVTable(void* pInterface, int nIndex)
{
    // Find the saved original and restore it
    for (int i = 0; i < g_nPatchCount; i++)
    {
        if (g_patches[i].pInterface == pInterface && g_patches[i].nIndex == nIndex)
        {
            void** vtable = *(void***)pInterface;

            DWORD dwOldProtect;
            if (!VirtualProtect(&vtable[nIndex], sizeof(void*), PAGE_READWRITE, &dwOldProtect))
                return FALSE;

            vtable[nIndex] = g_patches[i].pOriginal;
            VirtualProtect(&vtable[nIndex], sizeof(void*), dwOldProtect, &dwOldProtect);

            // Remove from patch list
            g_patches[i] = g_patches[g_nPatchCount - 1];
            g_nPatchCount--;

            return TRUE;
        }
    }

    return FALSE;
}

// ---------------------------------------------------------------------------
// Memory patching
// ---------------------------------------------------------------------------
extern "C" INJECT_API BOOL __stdcall PatchMemory(void* pAddress, const void* pData, DWORD dwSize)
{
    DWORD dwOldProtect;
    if (!VirtualProtect(pAddress, dwSize, PAGE_EXECUTE_READWRITE, &dwOldProtect))
        return FALSE;

    memcpy(pAddress, pData, dwSize);

    VirtualProtect(pAddress, dwSize, dwOldProtect, &dwOldProtect);
    FlushInstructionCache(GetCurrentProcess(), pAddress, dwSize);

    return TRUE;
}

// ---------------------------------------------------------------------------
// D3D9 hook trampolines
// ---------------------------------------------------------------------------
typedef HRESULT(__stdcall* PFN_BeginScene)(IDirect3DDevice9*);
typedef HRESULT(__stdcall* PFN_EndScene)(IDirect3DDevice9*);
typedef HRESULT(__stdcall* PFN_Reset)(IDirect3DDevice9*, D3DPRESENT_PARAMETERS*);

static PFN_Reset g_pfnOrigReset = NULL;

// Forward declarations for managed-layer callbacks
typedef void(__stdcall* PFN_DecalCallback)();
static PFN_DecalCallback g_pfnRender3D = NULL;
static PFN_DecalCallback g_pfnRender2D = NULL;
static PFN_DecalCallback g_pfnPreReset = NULL;
static PFN_DecalCallback g_pfnPostReset = NULL;

// Hooked BeginScene - called every frame before rendering
static HRESULT __stdcall Hooked_BeginScene(IDirect3DDevice9* pDevice)
{
    PFN_BeginScene pfnOrig = (PFN_BeginScene)BeginSceneO;
    HRESULT hr = pfnOrig(pDevice);

    // Store device reference on first call
    if (!g_pDevice)
    {
        g_pDevice = pDevice;
    }

    return hr;
}

// Hooked EndScene - called every frame after rendering; we draw our overlay here
static HRESULT __stdcall Hooked_EndScene(IDirect3DDevice9* pDevice)
{
    // Render 3D overlays (D3DService markers) before EndScene
    if (g_pfnRender3D)
        g_pfnRender3D();

    // Render 2D overlays (HUDs, UI) before EndScene
    if (g_pfnRender2D)
        g_pfnRender2D();

    PFN_EndScene pfnOrig = (PFN_EndScene)EndSceneO;
    return pfnOrig(pDevice);
}

// Hooked Reset - device lost/restore cycle
static HRESULT __stdcall Hooked_Reset(IDirect3DDevice9* pDevice, D3DPRESENT_PARAMETERS* pParams)
{
    // Notify managed layer to release D3D resources
    if (g_pfnPreReset)
        g_pfnPreReset();

    HRESULT hr = g_pfnOrigReset(pDevice, pParams);

    // Notify managed layer to recreate D3D resources
    if (g_pfnPostReset)
        g_pfnPostReset();

    return hr;
}

// ---------------------------------------------------------------------------
// D3D9 device discovery and hooking
// ---------------------------------------------------------------------------
static bool HookD3D9Device()
{
    // Create a temporary D3D9 device to get its vtable
    IDirect3D9* pD3D = Direct3DCreate9(D3D_SDK_VERSION);
    if (!pD3D)
        return false;

    D3DPRESENT_PARAMETERS pp = {};
    pp.Windowed = TRUE;
    pp.SwapEffect = D3DSWAPEFFECT_DISCARD;
    pp.BackBufferFormat = D3DFMT_UNKNOWN;
    pp.hDeviceWindow = g_hGameWindow;

    IDirect3DDevice9* pTempDevice = NULL;
    HRESULT hr = pD3D->CreateDevice(
        D3DADAPTER_DEFAULT, D3DDEVTYPE_NULLREF, g_hGameWindow,
        D3DCREATE_SOFTWARE_VERTEXPROCESSING,
        &pp, &pTempDevice);

    if (FAILED(hr) || !pTempDevice)
    {
        pD3D->Release();
        return false;
    }

    // Get vtable pointers from the temporary device
    void** vtable = *(void***)pTempDevice;

    // IDirect3DDevice9 vtable indices:
    // [41] = BeginScene
    // [42] = EndScene
    // [16] = Reset
    BeginSceneO = vtable[41];
    EndSceneO = vtable[42];
    g_pfnOrigReset = (PFN_Reset)vtable[16];

    pTempDevice->Release();
    pD3D->Release();

    // Now hook the actual game's D3D9 device vtable
    // We patch the device class vtable (shared by all D3D9 devices of the same type)
    // The temp device gave us the vtable layout; actual hooking happens via HookVTable
    // when the real device is discovered.

    return true;
}

// ---------------------------------------------------------------------------
// Window subclass procedure
// ---------------------------------------------------------------------------
static LRESULT CALLBACK InjectWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
    if (uMsg == WM_DESTROY)
    {
        // Game is shutting down - clean up
        Container_StopPlugins();
        Container_Terminate();
    }

    return CallWindowProc(g_pfnOrigWndProc, hWnd, uMsg, wParam, lParam);
}

// ---------------------------------------------------------------------------
// Startup / lifecycle
// ---------------------------------------------------------------------------
extern "C" INJECT_API void __stdcall DecalStartup()
{
    if (g_bInitialized)
        return;

    // Find the game window
    g_hGameWindow = FindWindow(NULL, _T("Asheron's Call"));
    if (!g_hGameWindow)
        return;

    // Subclass the game window to intercept messages
    g_pfnOrigWndProc = (WNDPROC)SetWindowLongPtr(g_hGameWindow, GWLP_WNDPROC, (LONG_PTR)InjectWndProc);

    // Initialize COM
    CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);

    // Discover D3D9 vtable layout for hooking
    HookD3D9Device();

    // Create the managed Decal core (Decal.Core COM server)
    // CLSID_DecalCore = {4557D5A1-...}
    CLSID clsidDecal;
    CLSIDFromProgID(L"Decal.Core", &clsidDecal);

    HRESULT hr = CoCreateInstance(clsidDecal, NULL, CLSCTX_INPROC_SERVER,
                                  IID_IUnknown, (void**)&g_pDecalCore);

    if (SUCCEEDED(hr) && g_pDecalCore)
    {
        // The managed Decal.Core will enumerate and start services/plugins
        // via registry entries under HKLM\SOFTWARE\Decal\Services and \Plugins
    }

    g_bInitialized = true;
}

extern "C" INJECT_API void __stdcall Container_Initialize(HWND hWnd, void* pDevice)
{
    g_hGameWindow = hWnd;
    g_pDevice = (IDirect3DDevice9*)pDevice;
    g_bContainerMode = true;

    CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);

    CLSID clsidDecal;
    CLSIDFromProgID(L"Decal.Core", &clsidDecal);
    CoCreateInstance(clsidDecal, NULL, CLSCTX_INPROC_SERVER,
                     IID_IUnknown, (void**)&g_pDecalCore);
}

extern "C" INJECT_API void __stdcall Container_StartPlugins()
{
    // Delegate to managed layer
    // The managed Decal.Core handles plugin enumeration and loading
}

extern "C" INJECT_API void __stdcall Container_StopPlugins()
{
    // Delegate to managed layer
}

extern "C" INJECT_API void __stdcall Container_Terminate()
{
    if (g_pDecalCore)
    {
        g_pDecalCore->Release();
        g_pDecalCore = NULL;
    }

    // Restore window procedure
    if (g_hGameWindow && g_pfnOrigWndProc)
    {
        SetWindowLongPtr(g_hGameWindow, GWLP_WNDPROC, (LONG_PTR)g_pfnOrigWndProc);
        g_pfnOrigWndProc = NULL;
    }

    // Unhook all vtable patches
    while (g_nPatchCount > 0)
    {
        int i = g_nPatchCount - 1;
        void** vtable = *(void***)g_patches[i].pInterface;

        DWORD dwOldProtect;
        VirtualProtect(&vtable[g_patches[i].nIndex], sizeof(void*), PAGE_READWRITE, &dwOldProtect);
        vtable[g_patches[i].nIndex] = g_patches[i].pOriginal;
        VirtualProtect(&vtable[g_patches[i].nIndex], sizeof(void*), dwOldProtect, &dwOldProtect);

        g_nPatchCount--;
    }

    g_pDevice = NULL;
    g_bInitialized = false;
}

extern "C" INJECT_API void __stdcall Container_Draw()
{
    if (g_pfnRender2D)
        g_pfnRender2D();
}

// ---------------------------------------------------------------------------
// COM server stubs (ATL-free, minimal)
// ---------------------------------------------------------------------------
static LONG g_lLockCount = 0;

STDAPI DllCanUnloadNow()
{
    return (g_lLockCount == 0 && !g_bInitialized) ? S_OK : S_FALSE;
}

STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID* ppv)
{
    *ppv = NULL;
    return CLASS_E_CLASSNOTAVAILABLE;
}

STDAPI DllRegisterServer()
{
    return S_OK;
}

STDAPI DllUnregisterServer()
{
    return S_OK;
}

// ---------------------------------------------------------------------------
// DllMain
// ---------------------------------------------------------------------------
BOOL WINAPI DllMain(HINSTANCE hInstance, DWORD dwReason, LPVOID lpReserved)
{
    switch (dwReason)
    {
    case DLL_PROCESS_ATTACH:
        g_hInstance = hInstance;
        DisableThreadLibraryCalls(hInstance);

        // Check if we're in the AC client process
        {
            TCHAR szFilename[MAX_PATH];
            GetModuleFileName(NULL, szFilename, MAX_PATH);

            // Extract just the filename
            LPTSTR pName = _tcsrchr(szFilename, _T('\\'));
            if (pName)
                pName++;
            else
                pName = szFilename;

            // If running inside client.exe, auto-start Decal
            if (_tcsnicmp(pName, _T("client"), 6) == 0 ||
                _tcsnicmp(pName, _T("acclient"), 8) == 0)
            {
                // Increment our own ref count to prevent premature unloading
                TCHAR szModule[MAX_PATH];
                GetModuleFileName(hInstance, szModule, MAX_PATH);
                LoadLibrary(szModule);

                // Start Decal hooking
                DecalStartup();
            }
        }
        break;

    case DLL_PROCESS_DETACH:
        Container_Terminate();
        break;
    }

    return TRUE;
}
