#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <stdio.h>
#include <stdint.h>

#include "../include/qemu_a6_usb_protocol.h"

#define QA6_DEFAULT_HOST "127.0.0.1"
#define QA6_DEFAULT_PORT "26050"
#define QA6_DRIVER_PATH L"\\\\.\\QemuA6Ude0"

static volatile LONG Qa6StopRequested;

static BOOL WINAPI
Qa6ConsoleHandler(DWORD event)
{
    if (event == CTRL_C_EVENT || event == CTRL_BREAK_EVENT ||
        event == CTRL_CLOSE_EVENT) {
        InterlockedExchange(&Qa6StopRequested, 1);
        return TRUE;
    }
    return FALSE;
}

static int
Qa6SendAll(SOCKET socketHandle, const unsigned char *buffer, int length)
{
    int offset = 0;

    while (offset < length) {
        int sent = send(socketHandle, (const char *)buffer + offset,
                        length - offset, 0);
        if (sent == SOCKET_ERROR || sent == 0) {
            return -1;
        }
        offset += sent;
    }
    return 0;
}

static int
Qa6ReceiveAll(SOCKET socketHandle, unsigned char *buffer, int length)
{
    int offset = 0;

    while (offset < length) {
        int received = recv(socketHandle, (char *)buffer + offset,
                            length - offset, 0);
        if (received == SOCKET_ERROR || received == 0) {
            return -1;
        }
        offset += received;
    }
    return 0;
}

static SOCKET
Qa6ConnectQemu(const char *host, const char *port)
{
    struct addrinfo hints;
    struct addrinfo *addresses = NULL;
    struct addrinfo *address;
    SOCKET socketHandle = INVALID_SOCKET;
    int result;

    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    result = getaddrinfo(host, port, &hints, &addresses);
    if (result != 0) {
        fprintf(stderr, "Resolution de %s:%s impossible: %d\n", host, port,
                result);
        return INVALID_SOCKET;
    }

    for (address = addresses; address != NULL; address = address->ai_next) {
        socketHandle = socket(address->ai_family, address->ai_socktype,
                              address->ai_protocol);
        if (socketHandle == INVALID_SOCKET) {
            continue;
        }
        if (connect(socketHandle, address->ai_addr,
                    (int)address->ai_addrlen) == 0) {
            break;
        }
        closesocket(socketHandle);
        socketHandle = INVALID_SOCKET;
    }

    freeaddrinfo(addresses);
    return socketHandle;
}

static HANDLE
Qa6OpenDriver(void)
{
    return CreateFileW(QA6_DRIVER_PATH, GENERIC_READ | GENERIC_WRITE,
                       FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
                       OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
}

static int
Qa6ValidateFrame(const unsigned char *frame, DWORD length,
                 unsigned short expectedType)
{
    const Qa6UsbFrameHeader *header;

    if (length < QA6_USB_HEADER_SIZE) {
        return -1;
    }
    header = (const Qa6UsbFrameHeader *)frame;
    if (header->magic != QA6_USB_MAGIC ||
        header->version != QA6_USB_PROTOCOL_VERSION ||
        header->type != expectedType ||
        header->payload_length > QA6_USB_MAX_PAYLOAD ||
        length != QA6_USB_HEADER_SIZE + header->payload_length) {
        return -1;
    }
    return 0;
}

static int
Qa6RelayLoop(HANDLE driver, SOCKET qemu)
{
    unsigned char request[QA6_USB_MAX_FRAME_SIZE];
    unsigned char response[QA6_USB_MAX_FRAME_SIZE];

    while (InterlockedCompareExchange(&Qa6StopRequested, 0, 0) == 0) {
        Qa6UsbFrameHeader *requestHeader;
        Qa6UsbFrameHeader *responseHeader;
        DWORD requestLength = 0;
        DWORD responseLength;
        DWORD written = 0;

        if (!ReadFile(driver, request, sizeof(request), &requestLength,
                      NULL)) {
            DWORD error = GetLastError();
            fprintf(stderr, "Unable to read from driver: error %lu\n",
                    error);
            return -1;
        }
        if (Qa6ValidateFrame(request, requestLength,
                             QA6_USB_MESSAGE_REQUEST) != 0) {
            fprintf(stderr, "Invalid frame received from driver (%lu bytes)\n",
                    requestLength);
            return -1;
        }

        requestHeader = (Qa6UsbFrameHeader *)request;
        printf("USB EP0 #%lu: bm=%02x req=%02x value=%02x%02x "
               "index=%02x%02x len=%u\n",
               requestHeader->request_id,
               requestHeader->setup[0], requestHeader->setup[1],
               requestHeader->setup[3], requestHeader->setup[2],
               requestHeader->setup[5], requestHeader->setup[4],
               (unsigned)(requestHeader->setup[6] |
                          (requestHeader->setup[7] << 8)));
        fflush(stdout);

        if (Qa6SendAll(qemu, request, (int)requestLength) != 0) {
            fprintf(stderr, "QEMU connection closed while sending.\n");
            return -1;
        }
        if (Qa6ReceiveAll(qemu, response, QA6_USB_HEADER_SIZE) != 0) {
            fprintf(stderr,
                    "QEMU connection closed while receiving a response.\n");
            return -1;
        }

        responseHeader = (Qa6UsbFrameHeader *)response;
        if (responseHeader->payload_length > QA6_USB_MAX_PAYLOAD) {
            fprintf(stderr, "Invalid QEMU response length.\n");
            return -1;
        }
        responseLength = QA6_USB_HEADER_SIZE +
                         responseHeader->payload_length;
        if (responseHeader->payload_length != 0 &&
            Qa6ReceiveAll(qemu, response + QA6_USB_HEADER_SIZE,
                          (int)responseHeader->payload_length) != 0) {
            fprintf(stderr, "QEMU connection closed while receiving data.\n");
            return -1;
        }
        if (Qa6ValidateFrame(response, responseLength,
                             QA6_USB_MESSAGE_RESPONSE) != 0 ||
            responseHeader->request_id != requestHeader->request_id) {
            fprintf(stderr, "Invalid QEMU response for request #%lu.\n",
                    requestHeader->request_id);
            return -1;
        }

        if (!WriteFile(driver, response, responseLength, &written, NULL) ||
            written != responseLength) {
            DWORD error = GetLastError();
            fprintf(stderr, "Unable to write to driver: error %lu\n",
                    error);
            return -1;
        }
    }

    return 0;
}

int
main(int argc, char **argv)
{
    const char *host = argc > 1 ? argv[1] : QA6_DEFAULT_HOST;
    const char *port = argc > 2 ? argv[2] : QA6_DEFAULT_PORT;
    WSADATA winsockData;
    HANDLE driver = INVALID_HANDLE_VALUE;
    SOCKET qemu = INVALID_SOCKET;
    int exitCode = 1;

    if (sizeof(Qa6UsbFrameHeader) != 32) {
        fprintf(stderr, "Internal error: USB protocol size=%zu\n",
                sizeof(Qa6UsbFrameHeader));
        return 2;
    }

    SetConsoleCtrlHandler(Qa6ConsoleHandler, TRUE);
    if (WSAStartup(MAKEWORD(2, 2), &winsockData) != 0) {
        fprintf(stderr, "Unable to initialize Winsock.\n");
        return 2;
    }

    printf("Native Windows QEMU A6 USB bridge\n");
    printf("Driver : \\\\.\\QemuA6Ude0\n");
    printf("QEMU   : %s:%s\n", host, port);

    while (InterlockedCompareExchange(&Qa6StopRequested, 0, 0) == 0) {
        driver = Qa6OpenDriver();
        if (driver != INVALID_HANDLE_VALUE) {
            break;
        }
        printf("Waiting for the QemuA6Ude driver...\n");
        Sleep(1000);
    }
    if (driver == INVALID_HANDLE_VALUE) {
        goto cleanup;
    }

    while (InterlockedCompareExchange(&Qa6StopRequested, 0, 0) == 0) {
        qemu = Qa6ConnectQemu(host, port);
        if (qemu != INVALID_SOCKET) {
            break;
        }
        printf("Waiting for QEMU at %s:%s...\n", host, port);
        Sleep(1000);
    }
    if (qemu == INVALID_SOCKET) {
        goto cleanup;
    }

    printf("Connected. idevicerestore can use the virtual DFU iPhone.\n");
    exitCode = Qa6RelayLoop(driver, qemu) == 0 ? 0 : 1;

cleanup:
    if (qemu != INVALID_SOCKET) {
        shutdown(qemu, SD_BOTH);
        closesocket(qemu);
    }
    if (driver != INVALID_HANDLE_VALUE) {
        CloseHandle(driver);
    }
    WSACleanup();
    return exitCode;
}