// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>

#include <usbhost/usbhost.h>
#include <unittest/unittest.h>
#include <zircon/device/usb-function-test.h>
#include <zircon/device/usb-peripheral.h>

namespace {

static struct usb_device* dev = nullptr;
static struct usb_endpoint_descriptor* bulk_out_ep = nullptr;
static struct usb_endpoint_descriptor* bulk_in_ep = nullptr;
static struct usb_endpoint_descriptor* intr_ep = nullptr;

static constexpr size_t BUFFER_SIZE = 4096;

// Data to send to the device
static uint8_t send_buf[BUFFER_SIZE] = {};
// Buffer for receiving data from the device
static int receive_buf[BUFFER_SIZE] = {};

static constexpr int TIMEOUT = 1000; // 1 seecond


// Fill send_buf with random values.
static void randomize() {
    // Generate some random data.
    for (size_t i = 0; i < sizeof(send_buf); i++) {
        send_buf[i] = static_cast<uint8_t>(random());
    }
}

// This thread waits for completion of an interrupt request.
static void* intr_complete_thread(void* arg) {
    auto* req = static_cast<struct usb_request**>(arg);
    *req = usb_request_wait(dev, TIMEOUT);
    return nullptr;
}

// This thread waits for completion of a bulk OUT request and bulk IN request.
static void* bulk_complete_thread(void* arg) {
    auto* reqs = static_cast<struct usb_request**>(arg);
    *reqs++ = usb_request_wait(dev, TIMEOUT);
    *reqs = usb_request_wait(dev, TIMEOUT);
    return nullptr;
}

// Tests control and interrupt transfers with specified transfer size.
static bool control_interrupt_test(size_t transfer_size) {
    BEGIN_TEST;

    randomize();

    // Send data to device via OUT control request.
    int ret = usb_device_control_transfer(dev,
                            USB_DIR_OUT | USB_TYPE_VENDOR | USB_TYPE_VENDOR | USB_RECIP_INTERFACE,
                            USB_FUNCTION_TEST_SET_DATA, 0, 0, send_buf, transfer_size, TIMEOUT);
    EXPECT_EQ(ret, transfer_size);

    // Receive data back from device via IN control request.
    ret = usb_device_control_transfer(dev,
                            USB_DIR_IN | USB_TYPE_VENDOR | USB_TYPE_VENDOR | USB_RECIP_INTERFACE,
                            USB_FUNCTION_TEST_GET_DATA, 0, 0, receive_buf, transfer_size, TIMEOUT);
    EXPECT_EQ(ret, transfer_size);

    // Sent and received data should match.
    EXPECT_EQ(memcmp(send_buf, receive_buf, transfer_size), 0);

    // Create a thread to wait for interrupt request.
    pthread_t thread;
    struct usb_request* complete_req = nullptr;
    ret = pthread_create(&thread, nullptr, intr_complete_thread, &complete_req); 
    EXPECT_EQ(ret, 0);

    // Queue read for interrupt request
    auto* req = usb_request_new(dev, intr_ep);
    EXPECT_NE(req, nullptr);
    req->buffer = receive_buf;
    req->buffer_length = static_cast<int>(transfer_size);
    ret = usb_request_queue(req);
    EXPECT_EQ(ret, 0);

    // Ask the device to send us an interrupt request containing the data we sent earlier.
    ret = usb_device_control_transfer(dev,
                            USB_DIR_OUT | USB_TYPE_VENDOR | USB_TYPE_VENDOR | USB_RECIP_INTERFACE,
                            USB_FUNCTION_TEST_SEND_INTERUPT, 0, 0, nullptr, 0, TIMEOUT);
    EXPECT_EQ(ret, 0);

    pthread_join(thread, nullptr);

    EXPECT_EQ(complete_req, req);

    // Sent data should match payload of interrupt request.
    EXPECT_EQ(memcmp(send_buf, receive_buf, transfer_size), 0);

    usb_request_free(req);

    END_TEST;
}

// Test control and interrupt requests with 8 byte transfer size.
static bool control_interrupt_test_8() {
    return control_interrupt_test(8);
}

// Test control and interrupt requests with 64 byte transfer size.
static bool control_interrupt_test_64() {
    return control_interrupt_test(64);
}

// Test control and interrupt requests with 100 byte transfer size.
static bool control_interrupt_test_100() {
    return control_interrupt_test(100);
}

// Test control and interrupt requests with 256 byte transfer size.
static bool control_interrupt_test_256() {
    return control_interrupt_test(256);
}

// Test control and interrupt requests with 1000 byte transfer size.
static bool control_interrupt_test_1000() {
    return control_interrupt_test(1000);
}

// Tests bulk OUT and IN transfers.
// Send BUFFER_SIZE bytes to device, read back and compare.
static bool bulk_test() {
    BEGIN_TEST;

    auto* send_req = usb_request_new(dev, bulk_out_ep);
    EXPECT_NE(send_req, nullptr);
    send_req->buffer = send_buf;
    send_req->buffer_length = static_cast<int>(BUFFER_SIZE);

    auto* receive_req = usb_request_new(dev, bulk_in_ep);
    EXPECT_NE(send_req, nullptr);
    receive_req->buffer = receive_buf;
    receive_req->buffer_length = static_cast<int>(BUFFER_SIZE);

    for (int i = 0; i < 10; i++) {
        randomize();

        // Create a thread to wait for request completions.
        pthread_t thread;
        struct usb_request* complete_reqs[2] = {};
        int ret = pthread_create(&thread, nullptr, bulk_complete_thread, complete_reqs); 
        EXPECT_EQ(ret, 0);

        // Queue requests in both directions
        ret = usb_request_queue(receive_req);
        EXPECT_EQ(ret, 0);
        ret = usb_request_queue(send_req);
        EXPECT_EQ(ret, 0);

        pthread_join(thread, nullptr);

        EXPECT_NE(complete_reqs[0], nullptr);
        EXPECT_NE(complete_reqs[1], nullptr);

        // Sent and received data should match.
        EXPECT_EQ(memcmp(send_buf, receive_buf, BUFFER_SIZE), 0);
    }

    usb_request_free(send_req);
    usb_request_free(receive_req);

    END_TEST;
}

// usb_host_load() will call this for all connected USB devices.
static int usb_device_added(const char *dev_name, void *client_data) {
    usb_descriptor_iter iter;
    struct usb_descriptor_header* header;
    struct usb_interface_descriptor* intf = nullptr;
    int ret;
    
    dev = usb_device_open(dev_name);
    if (!dev) {
        fprintf(stderr, "usb_device_open failed for %s\n", dev_name);
        return 0;
    }

    uint16_t vid = usb_device_get_vendor_id(dev);
    uint16_t pid = usb_device_get_product_id(dev);

    if (vid != GOOGLE_USB_VID || pid != GOOGLE_USB_FUNCTION_TEST_PID) {
        // Device doesn't match, so keep looking.
        return 0;
    }

    usb_descriptor_iter_init(dev, &iter);

    while ((header = usb_descriptor_iter_next(&iter)) != nullptr) {
        if (header->bDescriptorType == USB_DT_INTERFACE) {
            intf = reinterpret_cast<struct usb_interface_descriptor*>(header);
        } else if (header->bDescriptorType == USB_DT_ENDPOINT) {
            auto* ep = reinterpret_cast<struct usb_endpoint_descriptor*>(header);
            if (usb_endpoint_type(ep) == USB_ENDPOINT_XFER_BULK) {
                if (usb_endpoint_dir_in(ep)) {
                    bulk_in_ep = ep;
                } else {
                    bulk_out_ep = ep;
                }
            } else if (usb_endpoint_type(ep) == USB_ENDPOINT_XFER_INT) {
                intr_ep = ep;
            }
        }
    }

    if (!intf || !bulk_out_ep || !bulk_in_ep || !intr_ep) {
        fprintf(stderr, "could not find all our endpoints\n");
        goto fail;
    }

    ret = usb_device_claim_interface(dev, intf->bInterfaceNumber);
    if (ret < 0) {
        fprintf(stderr, "usb_device_claim_interface failed\n");
        goto fail;
    }

    // Test done, exit from usb_host_load().
    return 1;

fail:
    usb_device_close(dev);
    dev = nullptr;
    // Test done, exit from usb_host_load().
    return 1;
}

static int usb_device_removed(const char *dev_name, void *client_data) {
    return 0;
}

static int usb_discovery_done(void *client_data) {
    return 0;
}

} // anonymous namespace

BEGIN_TEST_CASE(usb_function_tests)
RUN_TEST(control_interrupt_test_8);
RUN_TEST(control_interrupt_test_64);
RUN_TEST(control_interrupt_test_100);
RUN_TEST(control_interrupt_test_256);
RUN_TEST(control_interrupt_test_1000);
RUN_TEST(bulk_test);
END_TEST_CASE(usb_function_tests)

int main(int argc, char** argv) {
    struct usb_host_context* context = usb_host_init();
    if (!context) {
        fprintf(stderr, "usb_host_context failed\n");
        return -1;
    }

    auto ret = usb_host_load(context, usb_device_added, usb_device_removed, usb_discovery_done,
                             nullptr);
    if (ret < 0) {
        fprintf(stderr, "usb_host_load failed!\n");
        usb_host_cleanup(context);
        goto fail;
    }
    if (!dev) {
        fprintf(stderr, "No USB Function Test devices found!\n");
        goto fail;
    }

    ret = unittest_run_all_tests(argc, argv) ? 0 : -1;

fail:
    if (dev) {
        usb_device_close(dev);
    }
    usb_host_cleanup(context);
    return ret;
}