#include <windows.h>
#include <assert.h>
#include <stdbool.h>
#include <stdlib.h>
#include <stdint.h>

#include "hook/table.h"

#include "hooklib/winusb.h"

#include "util/dprintf.h"

static bool winusb_initted;
static CRITICAL_SECTION winusb_lock;
static struct winusb_dev *winusb_classes;
static size_t winusb_nclasses;

static void winusb_hook_init();

BOOL my_WinUsb_QueryPipe(
    WINUSB_INTERFACE_HANDLE  InterfaceHandle,
    UCHAR                    AlternateInterfaceNumber,
    UCHAR                    PipeIndex,
    PWINUSB_PIPE_INFORMATION PipeInformation
);

BOOL my_WinUsb_QueryInterfaceSettings(
    WINUSB_INTERFACE_HANDLE   InterfaceHandle,
    UCHAR                     AlternateInterfaceNumber,
    PUSB_INTERFACE_DESCRIPTOR UsbAltInterfaceDescriptor
);

BOOL my_WinUsb_GetDescriptor(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   DescriptorType,
    UCHAR                   Index,
    USHORT                  LanguageID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred
);

BOOL my_WinUsb_WritePipe(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   PipeID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred,
    LPOVERLAPPED            Overlapped
);

BOOL my_WinUsb_QueryDeviceInformation(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    ULONG                   InformationType,
    PULONG                  BufferLength,
    PVOID                   Buffer
);

BOOL my_WinUsb_ReadPipe(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   PipeID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred,
    LPOVERLAPPED            Overlapped
);

BOOL my_WinUsb_Initialize(
    HANDLE                   DeviceHandle,
    PWINUSB_INTERFACE_HANDLE InterfaceHandle
);

BOOL my_WinUsb_Free(
    WINUSB_INTERFACE_HANDLE InterfaceHandle
);

BOOL (*next_WinUsb_QueryPipe)(
    WINUSB_INTERFACE_HANDLE  InterfaceHandle,
    UCHAR                    AlternateInterfaceNumber,
    UCHAR                    PipeIndex,
    PWINUSB_PIPE_INFORMATION PipeInformation
);

BOOL (*next_WinUsb_QueryInterfaceSettings)(
    WINUSB_INTERFACE_HANDLE   InterfaceHandle,
    UCHAR                     AlternateInterfaceNumber,
    PUSB_INTERFACE_DESCRIPTOR UsbAltInterfaceDescriptor
);

BOOL (*next_WinUsb_GetDescriptor)(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   DescriptorType,
    UCHAR                   Index,
    USHORT                  LanguageID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred
);

BOOL (*next_WinUsb_WritePipe)(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   PipeID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred,
    LPOVERLAPPED            Overlapped
);

BOOL (*next_WinUsb_QueryDeviceInformation)(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    ULONG                   InformationType,
    PULONG                  BufferLength,
    PVOID                   Buffer
);

BOOL (*next_WinUsb_ReadPipe)(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   PipeID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred,
    LPOVERLAPPED            Overlapped
);

BOOL (*next_WinUsb_Initialize)(
    HANDLE                   DeviceHandle,
    PWINUSB_INTERFACE_HANDLE InterfaceHandle
);

BOOL (*next_WinUsb_Free)(
    WINUSB_INTERFACE_HANDLE InterfaceHandle
);

static const struct hook_symbol winusb_syms[] = {
    {
        .name   = "WinUsb_QueryPipe",
        .patch  = my_WinUsb_QueryPipe,
        .link   = (void *) &next_WinUsb_QueryPipe,
    }, {
        .name   = "WinUsb_QueryInterfaceSettings",
        .patch  = my_WinUsb_QueryInterfaceSettings,
        .link   = (void *) &next_WinUsb_QueryInterfaceSettings,
    }, {
        .name   = "WinUsb_GetDescriptor",
        .patch  = my_WinUsb_GetDescriptor,
        .link   = (void *) &next_WinUsb_GetDescriptor,
    }, {
        .name   = "WinUsb_WritePipe",
        .patch  = my_WinUsb_WritePipe,
        .link   = (void *) &next_WinUsb_WritePipe,
    }, {
        .name   = "WinUsb_QueryDeviceInformation",
        .patch  = my_WinUsb_QueryDeviceInformation,
        .link   = (void *) &next_WinUsb_QueryDeviceInformation,
    }, {
        .name   = "WinUsb_ReadPipe",
        .patch  = my_WinUsb_ReadPipe,
        .link   = (void *) &next_WinUsb_ReadPipe,
    }, {
        .name   = "WinUsb_Initialize",
        .patch  = my_WinUsb_Initialize,
        .link   = (void *) &next_WinUsb_Initialize,
    }, {
        .name   = "WinUsb_Free",
        .patch  = my_WinUsb_Free,
        .link   = (void *) &next_WinUsb_Free,
    }
};

HRESULT winusb_add_phantom_dev(const struct winusb_dev *dev)
{
    struct winusb_dev *class_;
    struct winusb_dev *new_array;
    HRESULT hr;

    assert(dev != NULL);

    winusb_hook_init();

    EnterCriticalSection(&winusb_lock);

    new_array = realloc(
            winusb_classes,
            (winusb_nclasses + 1) * sizeof(struct winusb_dev));

    if (new_array == NULL) {
        hr = E_OUTOFMEMORY;

        goto end;
    }

    winusb_classes = new_array;

    // TODO: Copy the dev class to the new slot in the array
    class_ = &winusb_classes[winusb_nclasses++];
    hr = S_OK;

end:
    LeaveCriticalSection(&winusb_lock);

    return hr;
}

static void winusb_hook_init()
{
    if (winusb_initted) {
        return;
    }

    winusb_hook_insert_hooks(NULL);

    InitializeCriticalSection(&winusb_lock);
    winusb_initted = true;
}

void winusb_hook_insert_hooks(HMODULE target)
{
    hook_table_apply(
            target,
            "winusb.dll",
            winusb_syms,
            _countof(winusb_syms));
}