#include "dr_api.h"
#include "dr_tools.h"
#include "drwrap.h"
#include "dr_ir_opnd.h"
#include "dr_ir_instr.h"

#include <ctype.h>
#include <string.h>
#include <windows.h>
#include <commdlg.h>
#include <stdint.h>
#include <wchar.h>
#include <shobjidl.h>

// Note: Some Qt mangled names might vary slightly by compiler version.
// These are typical for MSVC.
const char *MANGLED_QString_const = "??0QString@@QEAA@PEBD@Z";
const char *MANGLED_QUrl_const = "??0QUrl@@QEAA@AEBVQString@@W4ParsingMode@0@@Z";
const char *MANGLED_QFileDialog_getOpenFileName = "?getOpenFileName@QFileDialog@@SA?AVQString@@PEAVQWidget@@AEBV2@11PEAV2@V?$QFlags@W4Option@QFileDialog@@@@@Z";
const char *MANGLED_QFileDialog_getOpenFileNames = "?getOpenFileNames@QFileDialog@@SA?AV?$QList@VQString@@@@PEAVQWidget@@AEBVQString@@11PEAV4@V?$QFlags@W4Option@QFileDialog@@@@@Z";
const char *MANGLED_QFileDialog_getOpenFileUrl = "?getOpenFileUrl@QFileDialog@@SA?AVQUrl@@PEAVQWidget@@AEBVQString@@AEBV2@1PEAV4@V?$QFlags@W4Option@QFileDialog@@@@AEBV?$QList@VQString@@@@@Z";
const char *MANGLED_QFileDialog_getOpenFileUrls = "?getOpenFileUrls@QFileDialog@@SA?AV?$QList@VQUrl@@@@PEAVQWidget@@AEBVQString@@AEBVQUrl@@1PEAV4@V?$QFlags@W4Option@QFileDialog@@@@AEBV?$QList@VQString@@@@@Z";
const char *MANGLED_QFileDialog_getOpenFileContent = "?getOpenFileContent@QFileDialog@@SAXAEBVQString@@AEBV?$function@$$A6AXAEBVQString@@AEBVQByteArray@@@Z@std@@PEAVQWidget@@@Z";
const char *MANGLED_QFileDialog_getSaveFileName = "?getSaveFileName@QFileDialog@@SA?AVQString@@PEAVQWidget@@AEBV2@11PEAV2@V?$QFlags@W4Option@QFileDialog@@@@@Z";
const char *MANGLED_QFileDialog_getSaveFileUrl = "?getSaveFileUrl@QFileDialog@@SA?AVQUrl@@PEAVQWidget@@AEBVQString@@AEBV2@1PEAV4@V?$QFlags@W4Option@QFileDialog@@@@AEBV?$QList@VQString@@@@@Z";

char *OPEN_FILE = "C:\\Users\\vboxuser\\Downloads\\spoofed_number.txt";
char *SAVE_FILE = "C:\\Users\\vboxuser\\Downloads\\spoofed_save.txt";

// Function pointers for functions we need to call ourselves
static void (*fn_QString_const)(void *qstring, const char *str) = NULL;
static void (*fn_QUrl_const)(void *qurl, void *qstring, int parsing_mode) = NULL;
static LPVOID (*fn_CoTaskMemAlloc)(size_t cb) = NULL;
static HRESULT (*fn_SHCreateItemFromParsingName)(PWSTR, void *, void *, void **);

// --- drwrap Callbacks for IFileOpenDialog (COM) ---

struct per_thread_info
{
    // Information needed for COM hooking
    IShellItem *sitem_to_wrap;
};

static void pre_CoCreateInstance(void *wrapctx, void **user_data);
static void post_CoCreateInstance(void *wrapctx, void *user_data);
static void post_IFileOpenDialog_GetResult(void *wrapctx, void *user_data);
static void pre_IShellItem_GetDisplayName(void *wrapctx, void **user_data);
static void pre_IFileOpenDialog_GetResult(void *wrapctx, void **user_data);
static void pre_IFileOpenDialog_Show(void *wrapctx, void **user_data);

static HRESULT custom_IFileDialog_Show(void *, void *);
static HRESULT custom_IFileOpenDialog_GetResult(void *, IShellItem **ppsitem);
static HRESULT custom_IFileSaveDialog_GetResult(void *, IShellItem **ppsitem);

struct IFileOpenDialogVtbl customFODVtbl = {};
struct IFileSaveDialogVtbl customFSDVtbl = {};

// --- drwrap Callbacks for WinAPI Functions ---

static void pre_GetOpenFileNameA(void *wrapctx, void **user_data)
{
    OPENFILENAMEA *ofn = (OPENFILENAMEA *)drwrap_get_arg(wrapctx, 0);
    if (ofn && ofn->lpstrFile)
    {
        strncpy(ofn->lpstrFile, OPEN_FILE, ofn->nMaxFile - 1);
        ofn->lpstrFile[ofn->nMaxFile - 1] = '\0';
    }

    drwrap_skip_call(wrapctx, (void *)1, 0);
}

static void pre_GetOpenFileNameW(void *wrapctx, void **user_data)
{
    OPENFILENAMEW *ofn = (OPENFILENAMEW *)drwrap_get_arg(wrapctx, 0);
    if (ofn && ofn->lpstrFile)
    {
        mbstowcs(ofn->lpstrFile, OPEN_FILE, ofn->nMaxFile - 1);
        ofn->lpstrFile[ofn->nMaxFile - 1] = L'\0';
    }

    drwrap_skip_call(wrapctx, (void *)1, 0);
}

static void pre_GetSaveFileNameA(void *wrapctx, void **user_data)
{
    OPENFILENAMEA *ofn = (OPENFILENAMEA *)drwrap_get_arg(wrapctx, 0);
    if (ofn && ofn->lpstrFile)
    {
        strncpy(ofn->lpstrFile, SAVE_FILE, ofn->nMaxFile - 1);
        ofn->lpstrFile[ofn->nMaxFile - 1] = '\0';
    }

    drwrap_skip_call(wrapctx, (void *)1, 0);
}

static void pre_GetSaveFileNameW(void *wrapctx, void **user_data)
{
    OPENFILENAMEW *ofn = (OPENFILENAMEW *)drwrap_get_arg(wrapctx, 0);
    if (ofn && ofn->lpstrFile)
    {
        mbstowcs(ofn->lpstrFile, SAVE_FILE, ofn->nMaxFile - 1);
        ofn->lpstrFile[ofn->nMaxFile - 1] = L'\0';
    }

    drwrap_skip_call(wrapctx, (void *)1, 0);
}

// --- drwrap Callbacks for Qt Functions ---

static void pre_QFileDialog_getOpenFileName(void *wrapctx, void **user_data)
{
    if (fn_QString_const == NULL)
        return;
    // The return value is passed as a pointer in the first argument (sret)
    void *qstring_ret_ptr = (void *)drwrap_get_arg(wrapctx, 0);
    fn_QString_const(qstring_ret_ptr, OPEN_FILE);
    drwrap_skip_call(wrapctx, NULL, 0);
}

static void pre_QFileDialog_getOpenFileUrl(void *wrapctx, void **user_data)
{
    if (fn_QString_const == NULL || fn_QUrl_const == NULL)
        return;

    const char *prefix = "file:///";
    char urlbuf[MAX_PATH + 10];
    snprintf(urlbuf, sizeof(urlbuf), "%s%s", prefix, OPEN_FILE);

    // Replace backslashes for URL format
    for (char *p = urlbuf; *p; ++p)
    {
        if (*p == '\\')
            *p = '/';
    }

    // Create a temporary QString for the URL
    char local_qstring[24]; // Allocate space for a QString object on the stack
    fn_QString_const(local_qstring, urlbuf);

    // Get the return QUrl object pointer (sret)
    void *qurl_ret_ptr = (void *)drwrap_get_arg(wrapctx, 0);
    fn_QUrl_const(qurl_ret_ptr, local_qstring, 0); // 0 is TolerantMode

    drwrap_skip_call(wrapctx, NULL, 0);
}

static void pre_QFileDialog_getSaveFileName(void *wrapctx, void **user_data)
{
    if (fn_QString_const == NULL)
        return;

    // The return value is passed as a pointer in the first argument (sret)
    void *qstring_ret_ptr = (void *)drwrap_get_arg(wrapctx, 0);
    fn_QString_const(qstring_ret_ptr, SAVE_FILE);
    drwrap_skip_call(wrapctx, NULL, 0);
}

static void pre_QFileDialog_getSaveFileUrl(void *wrapctx, void **user_data)
{
    if (fn_QString_const == NULL || fn_QUrl_const == NULL)
        return;

    const char *prefix = "file:///";
    char urlbuf[MAX_PATH + 10];
    snprintf(urlbuf, sizeof(urlbuf), "%s%s", prefix, SAVE_FILE);
    
    // Replace backslashes for URL format
    for (char *p = urlbuf; *p; ++p)
    {
        if (*p == '\\')
            *p = '/';
    }

    // Create a temporary QString for the URL
    char local_qstring[24]; // Allocate space for a QString object on the stack
    fn_QString_const(local_qstring, urlbuf);

    // Get the return QUrl object pointer (sret)
    void *qurl_ret_ptr = (void *)drwrap_get_arg(wrapctx, 0);
    fn_QUrl_const(qurl_ret_ptr, local_qstring, 0); // 0 is TolerantMode

    drwrap_skip_call(wrapctx, NULL, 0);
}

// --- COM Wrapping Logic ---

const int TYPE_IFileOpenDialog = 1;
const int TYPE_IFileSaveDialog = 2;

struct CoCreateInstance_data {
    int type;
    void *dialog;
};

static void pre_CoCreateInstance(void *wrapctx, void **user_data)
{
    *user_data = NULL;

    // Arguments: rclsid, pUnkOuter, dwClsContext, riid, ppv
    const IID *riid = (const IID *)drwrap_get_arg(wrapctx, 3);
    const CLSID *rclsid = (const CLSID *)drwrap_get_arg(wrapctx, 0);
    if (!rclsid || !riid) return;

    int type = 0;
    void *dialog = NULL;

    if (IsEqualIID(rclsid, &CLSID_FileOpenDialog) && IsEqualIID(riid, &IID_IFileOpenDialog))
    {
        dr_printf("drwrap: Intercepting CoCreateInstance for FileOpenDialog\n");
        // Pass the pointer to the output interface pointer (ppv) to the post-hook
        dialog = drwrap_get_arg(wrapctx, 4);
        type = TYPE_IFileOpenDialog;
    }

    if (IsEqualIID(rclsid, &CLSID_FileSaveDialog) && IsEqualIID(riid, &IID_IFileSaveDialog))
    {
        dr_printf("drwrap: Intercepting CoCreateInstance for FileSaveDialog\n");
        // Pass the pointer to the output interface pointer (ppv) to the post-hook
        dialog = drwrap_get_arg(wrapctx, 4);
        type = TYPE_IFileSaveDialog;
    }

    if (!type || !dialog) return;

    struct CoCreateInstance_data *m_user_data = malloc(sizeof(*m_user_data));
    if (!m_user_data) {
        perror("malloc");
        dr_exit_process(1);
    }

    m_user_data->type = type;
    m_user_data->dialog = dialog;

    *user_data = (void *)m_user_data;
}

static void post_CoCreateInstance(void *wrapctx, void *user_data)
{
    // If user_data is NULL, this wasn't the CoCreateInstance we cared about
    if (!user_data)
        return;

    struct CoCreateInstance_data m_user_data;
    m_user_data = *(struct CoCreateInstance_data *)user_data;

    free(user_data);


    HRESULT hr = (HRESULT)drwrap_get_retval(wrapctx);
    if (SUCCEEDED(hr))
    {
        if (m_user_data.type == TYPE_IFileOpenDialog) {
            IFileOpenDialog **ppv = (IFileOpenDialog **)m_user_data.dialog;
            if (!ppv || !*ppv) return;
            
            IFileOpenDialog *fod = *ppv;
            
            WCHAR wfile[MAX_PATH];
            mbstowcs(wfile, OPEN_FILE, MAX_PATH - 1);
            wfile[MAX_PATH - 1] = L'\0';

            fod->lpVtbl->SetFileName(fod, wfile);

            customFODVtbl = *fod->lpVtbl;
            customFODVtbl.Show = custom_IFileDialog_Show;
            customFODVtbl.GetResult = custom_IFileOpenDialog_GetResult;

            fod->lpVtbl = &customFODVtbl;
        }

        if (m_user_data.type == TYPE_IFileSaveDialog) {
            IFileSaveDialog **ppv = (IFileSaveDialog **)m_user_data.dialog;
            if (!ppv || !*ppv) return;
            
            IFileSaveDialog *fsd = *ppv;
            
            WCHAR wfile[MAX_PATH];
            mbstowcs(wfile, SAVE_FILE, MAX_PATH - 1);
            wfile[MAX_PATH - 1] = L'\0';

            customFSDVtbl = *fsd->lpVtbl;
            customFSDVtbl.Show = custom_IFileDialog_Show;
            customFSDVtbl.GetResult = custom_IFileSaveDialog_GetResult;

            fsd->lpVtbl = &customFSDVtbl;
        }
    }
}

static HRESULT custom_IFileDialog_Show(void *self, void *b)
{
    return S_OK;
}

static HRESULT custom_IFileSaveDialog_GetResult(void *self, IShellItem **ppsitem)
{
    WCHAR file[MAX_PATH];
    mbstowcs(file, SAVE_FILE, MAX_PATH);
    file[MAX_PATH - 1] = L'\0';

    HANDLE handle = CreateFileA(SAVE_FILE, GENERIC_ALL,
                                FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
                                NULL, CREATE_NEW, FILE_ATTRIBUTE_NORMAL, NULL);
    CloseHandle(handle);

    return fn_SHCreateItemFromParsingName(file, NULL, &IID_IShellItem, (void **)ppsitem);
}

static HRESULT custom_IFileOpenDialog_GetResult(void *self, IShellItem **ppsitem)
{
    WCHAR file[MAX_PATH];
    mbstowcs(file, OPEN_FILE, MAX_PATH);
    file[MAX_PATH - 1] = L'\0';

    return fn_SHCreateItemFromParsingName(file, NULL, &IID_IShellItem, (void **)ppsitem);
}

// --- Main Events ---

static void event_module_load(void *drcontext, const module_data_t *mod, bool loaded)
{
    // Wrap functions as their modules are loaded.
    // This is much more efficient than iterating every basic block.

#define WRAP_FUNCTION(name, pre_hook, post_hook)                                           \
    {                                                                                      \
        app_pc addr = (app_pc)dr_get_proc_address(mod->handle, name);                      \
        if (addr != NULL)                                                                  \
        {                                                                                  \
            drwrap_wrap(addr, pre_hook, post_hook);                                        \
            dr_printf("Wrapped %s!%s at %p\n", dr_module_preferred_name(mod), name, addr); \
        }                                                                                  \
    }

#define GET_FUNCTION_POINTER(var, name)                                                              \
    {                                                                                                \
        if (var == NULL)                                                                             \
        {                                                                                            \
            var = (void *)dr_get_proc_address(mod->handle, name);                                    \
            if (var != NULL)                                                                         \
                dr_printf("Found function pointer %s in %s\n", name, dr_module_preferred_name(mod)); \
        }                                                                                            \
    }

    // WinAPI Hooks
    WRAP_FUNCTION("GetOpenFileNameA", pre_GetOpenFileNameA, NULL);
    WRAP_FUNCTION("GetOpenFileNameW", pre_GetOpenFileNameW, NULL);
    WRAP_FUNCTION("GetSaveFileNameA", pre_GetSaveFileNameA, NULL);
    WRAP_FUNCTION("GetSaveFileNameW", pre_GetSaveFileNameW, NULL);
    WRAP_FUNCTION("CoCreateInstance", pre_CoCreateInstance, post_CoCreateInstance);

    // Qt Hooks
    WRAP_FUNCTION(MANGLED_QFileDialog_getOpenFileName, pre_QFileDialog_getOpenFileName, NULL);
    WRAP_FUNCTION(MANGLED_QFileDialog_getOpenFileUrl, pre_QFileDialog_getOpenFileUrl, NULL);
    WRAP_FUNCTION(MANGLED_QFileDialog_getSaveFileName, pre_QFileDialog_getSaveFileName, NULL);
    WRAP_FUNCTION(MANGLED_QFileDialog_getSaveFileUrl, pre_QFileDialog_getSaveFileUrl, NULL);

    // Get pointers to functions we need to call ourselves
    GET_FUNCTION_POINTER(fn_QString_const, MANGLED_QString_const);
    GET_FUNCTION_POINTER(fn_QUrl_const, MANGLED_QUrl_const);
    GET_FUNCTION_POINTER(fn_CoTaskMemAlloc, "CoTaskMemAlloc");
    GET_FUNCTION_POINTER(fn_SHCreateItemFromParsingName, "SHCreateItemFromParsingName");
}

static void event_exit(void)
{
    drwrap_exit();
}

DR_EXPORT void dr_client_main(client_id_t id, int argc, const char *argv[])
{
    dr_set_client_name("File Dialog Interceptor", "http://example.com/");

    if (!drwrap_init())
    {
        DR_ASSERT_MSG(false, "drwrap_init failed");
    }

    dr_register_module_load_event(event_module_load);
    dr_register_exit_event(event_exit);
}