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

#include "hook/table.h"
#include "hook/iohook.h"
#include "hook/procaddr.h"

#include "board/guid.h"
#include "hooklib/setupapi.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();
static HRESULT winusb_handle_irp(struct irp *irp);

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;

    memcpy_s(&winusb_classes[winusb_nclasses++], sizeof(struct winusb_dev), dev, sizeof(*dev));
    hr = S_OK;

    setupapi_add_phantom_dev(&usbio_guid, dev->path);
    iohook_push_handler(winusb_handle_irp);

    dprintf("WinUSB: Init\n");

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));
    
    proc_addr_table_push(
            target,
            "WINUSB.dll",
            winusb_syms,
            _countof(winusb_syms));
}

static HRESULT winusb_handle_irp(struct irp *irp)
{
    bool found = false;
    if (irp->op != IRP_OP_OPEN && irp->op != IRP_OP_CLOSE) {
        return iohook_invoke_next(irp);
    }

    if (irp->open_filename == NULL || irp->open_filename[0] == L'\0') {
        return iohook_invoke_next(irp);
    }
    
    for (int i = 0; i < winusb_nclasses; i++) {
        if (!wcscmp(irp->open_filename, winusb_classes[i].path)) {
            dprintf("WinUSB: Open/close %ls\n", winusb_classes[i].path);
            irp->fd = winusb_classes[i].fd;
            found = true;
            break;
        }
    }

    if (!found) {
        return iohook_invoke_next(irp);
    }
    
    return S_OK;
}

// FIXME
int get_winusb_fake_dev_idx(WINUSB_INTERFACE_HANDLE InterfaceHandle)
{
    for (int i = 0; i < winusb_nclasses; i++) {
        if (winusb_classes[i].fd == InterfaceHandle) {
            return i;
        }
    }
    return -1;
}

BOOL my_WinUsb_QueryPipe(WINUSB_INTERFACE_HANDLE InterfaceHandle, UCHAR AlternateInterfaceNumber, UCHAR PipeIndex, PWINUSB_PIPE_INFORMATION PipeInformation)
{
    dprintf("WinUSB: my_WinUsb_QueryPipe\n");
    UCHAR this_pipeid = 0;

    int idx = 0;
    /*int idx = get_winusb_fake_dev_idx(InterfaceHandle);
    if (idx == -1) {
        return next_WinUsb_QueryPipe(InterfaceHandle, AlternateInterfaceNumber, PipeIndex, PipeInformation);
    }*/

    switch (PipeIndex) {
        case 0: this_pipeid = 0x02; break;
        case 1: this_pipeid = 0x81; break;
    }
    
    if (!this_pipeid) {
        return false;
    }
    
    PipeInformation->PipeType = UsbdPipeTypeBulk;
    PipeInformation->PipeId = this_pipeid;
    PipeInformation->Interval = 0x00; // never checked?
    PipeInformation->MaximumPacketSize = 0x40; // never checked?
    
    return true;
}

BOOL my_WinUsb_QueryInterfaceSettings(WINUSB_INTERFACE_HANDLE InterfaceHandle, UCHAR AlternateInterfaceNumber, PUSB_INTERFACE_DESCRIPTOR UsbAltInterfaceDescriptor)
{
    dprintf("WinUSB: my_WinUsb_QueryInterfaceSettings\n");
    int idx = 0;
    /*int idx = get_winusb_fake_dev_idx(InterfaceHandle);
    if (idx == -1) {
        return next_WinUsb_QueryInterfaceSettings(InterfaceHandle, AlternateInterfaceNumber, UsbAltInterfaceDescriptor);
    }*/

    UsbAltInterfaceDescriptor->bLength = 9;
    UsbAltInterfaceDescriptor->bDescriptorType = USB_INTERFACE_DESCRIPTOR_TYPE;
    UsbAltInterfaceDescriptor->bInterfaceNumber = 0;
    UsbAltInterfaceDescriptor->bAlternateSetting = 0;
    UsbAltInterfaceDescriptor->bNumEndpoints = 2; // nothing else is checked
    UsbAltInterfaceDescriptor->bInterfaceClass = 255;
    UsbAltInterfaceDescriptor->bInterfaceSubClass = 0;
    UsbAltInterfaceDescriptor->bInterfaceProtocol = 0;
    UsbAltInterfaceDescriptor->iInterface = 0;
    
    return true;
}

BOOL my_WinUsb_GetDescriptor(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   DescriptorType,
    UCHAR                   Index,
    USHORT                  LanguageID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred
)
{
    int idx = 0;
    /*int idx = get_winusb_fake_dev_idx(InterfaceHandle);
    if (idx == -1) {
        return next_WinUsb_GetDescriptor(
            InterfaceHandle, 
            DescriptorType, 
            Index, 
            LanguageID,
            Buffer,
            BufferLength,
            LengthTransferred
        );
    }*/

    dprintf("WinUSB: my_WinUsb_GetDescriptor %d\n", idx);

    if (BufferLength < winusb_classes[idx].descriptor.bLength) {
        SetLastError(ERROR_INSUFFICIENT_BUFFER);
        return false;
    }

    // Surly this won't explode
    memcpy_s(Buffer, BufferLength, &winusb_classes[idx].descriptor, winusb_classes[(uintptr_t)InterfaceHandle].descriptor.bLength);
    
    *LengthTransferred = winusb_classes[idx].descriptor.bLength;

    return true;
}

BOOL my_WinUsb_WritePipe(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   PipeID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred,
    LPOVERLAPPED            Overlapped
)
{
    int idx = 0;
    /*int idx = get_winusb_fake_dev_idx(InterfaceHandle);
    if (idx == -1) {
        return next_WinUsb_WritePipe(
            InterfaceHandle, 
            PipeID, 
            Buffer, 
            BufferLength,
            LengthTransferred,
            Overlapped
        );
    }*/
#if 0
    dprintf("WritePipe %d\t", (int)BufferLength);
    for (int i = 0; i < BufferLength; i++) {
        dprintf("%02X ", (uint8_t)Buffer[i]);
    }
    dprintf("\n");
#endif
    winusb_classes[idx].write_pipe(NULL, Buffer, BufferLength, (uint32_t *)LengthTransferred);
    return true;
}

BOOL my_WinUsb_QueryDeviceInformation(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    ULONG                   InformationType,
    PULONG                  BufferLength,
    PVOID                   Buffer
)
{
    //dprintf("WinUSB: my_WinUsb_QueryDeviceInformation\n");
    memset(Buffer, 0x03, 1);
    return true;
}

BOOL my_WinUsb_ReadPipe(
    WINUSB_INTERFACE_HANDLE InterfaceHandle,
    UCHAR                   PipeID,
    PUCHAR                  Buffer,
    ULONG                   BufferLength,
    PULONG                  LengthTransferred,
    LPOVERLAPPED            Overlapped
)
{
    int idx = 0;
    /*int idx = get_winusb_fake_dev_idx(InterfaceHandle);
    if (idx == -1) {
        return next_WinUsb_WritePipe(
            InterfaceHandle, 
            PipeID, 
            Buffer, 
            BufferLength,
            LengthTransferred,
            Overlapped
        );
    }*/
#if 0
    dprintf("ReadPipe %d\n", (int)BufferLength);
#endif
    winusb_classes[idx].read_pipe(NULL, Buffer, BufferLength, (uint32_t *)LengthTransferred);
    return true; 
}

BOOL my_WinUsb_Initialize(HANDLE DeviceHandle, PWINUSB_INTERFACE_HANDLE InterfaceHandle)
{
    dprintf("WinUSB: my_WinUsb_Initialize\n");
    for (int i = 0; i < winusb_nclasses; i++) {
        if (winusb_classes[i].fd == DeviceHandle) {
            dprintf("WinUSB: Init %ls\n", winusb_classes[i].path);
            InterfaceHandle = winusb_classes[i].fd;
            return true;
        }
    }

    return next_WinUsb_Initialize(DeviceHandle, InterfaceHandle);
}

BOOL my_WinUsb_Free(WINUSB_INTERFACE_HANDLE InterfaceHandle)
{
    if (get_winusb_fake_dev_idx(InterfaceHandle) == -1) {
        return next_WinUsb_Free(InterfaceHandle);
    }
    
    InterfaceHandle = NULL;
    return true;
    
}