// Copyright 2015-2020 Espressif Systems (Shanghai) PTE LTD
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "USBHID.h"

#if CONFIG_TINYUSB_HID_ENABLED

#include "USB.h"
#include "esp32-hal-tinyusb.h"
#include "esp_hid_common.h"

#define USB_HID_DEVICES_MAX 10

ESP_EVENT_DEFINE_BASE(ARDUINO_USB_HID_EVENTS);
esp_err_t arduino_usb_event_post(
    esp_event_base_t event_base, int32_t event_id, void *event_data, size_t event_data_size,
    TickType_t ticks_to_wait
);
esp_err_t arduino_usb_event_handler_register_with(
    esp_event_base_t event_base, int32_t event_id, esp_event_handler_t event_handler, void *event_handler_arg
);

typedef struct {
    USBHIDDevice *device;
    uint8_t reports_num;
    uint8_t *report_ids;
} tinyusb_hid_device_t;

static tinyusb_hid_device_t tinyusb_hid_devices[USB_HID_DEVICES_MAX];

static uint8_t tinyusb_hid_devices_num = 0;
static bool tinyusb_hid_devices_is_initialized = false;
static SemaphoreHandle_t tinyusb_hid_device_input_sem = NULL;
static SemaphoreHandle_t tinyusb_hid_device_input_mutex = NULL;

static bool tinyusb_hid_is_initialized = false;
static hid_interface_protocol_enum_t tinyusb_interface_protocol = HID_ITF_PROTOCOL_NONE;
static uint8_t tinyusb_loaded_hid_devices_num = 0;
static uint16_t tinyusb_hid_device_descriptor_len = 0;
static uint8_t *tinyusb_hid_device_descriptor = NULL;
#if ARDUHAL_LOG_LEVEL >= ARDUHAL_LOG_LEVEL_DEBUG
static const char *tinyusb_hid_device_report_types[4] = {"INVALID", "INPUT", "OUTPUT", "FEATURE"};
#endif

static bool tinyusb_enable_hid_device(uint16_t descriptor_len, USBHIDDevice *device) {
    if (tinyusb_hid_is_initialized) {
        log_e("TinyUSB HID has already started! Device not enabled");
        return false;
    }
    if (tinyusb_loaded_hid_devices_num >= USB_HID_DEVICES_MAX) {
        log_e("Maximum devices already enabled! Device not enabled");
        return false;
    }
    tinyusb_hid_device_descriptor_len += descriptor_len;
    tinyusb_hid_devices[tinyusb_loaded_hid_devices_num++].device = device;

    log_d("Device[%u] len: %u", tinyusb_loaded_hid_devices_num - 1, descriptor_len);
    return true;
}

USBHIDDevice *tinyusb_get_device_by_report_id(uint8_t report_id) {
    for (uint8_t i = 0; i < tinyusb_loaded_hid_devices_num; i++) {
        tinyusb_hid_device_t *device = &tinyusb_hid_devices[i];
        if (device->device && device->reports_num) {
            for (uint8_t r = 0; r < device->reports_num; r++) {
                if (report_id == device->report_ids[r]) { return device->device; }
            }
        }
    }
    return NULL;
}

static uint16_t tinyusb_on_get_feature(uint8_t report_id, uint8_t *buffer, uint16_t reqlen) {
    USBHIDDevice *device = tinyusb_get_device_by_report_id(report_id);
    if (device) { return device->_onGetFeature(report_id, buffer, reqlen); }
    return 0;
}

static bool tinyusb_on_set_feature(uint8_t report_id, const uint8_t *buffer, uint16_t reqlen) {
    USBHIDDevice *device = tinyusb_get_device_by_report_id(report_id);
    if (device) {
        device->_onSetFeature(report_id, buffer, reqlen);
        return true;
    }
    return false;
}

static bool tinyusb_on_set_output(uint8_t report_id, const uint8_t *buffer, uint16_t reqlen) {
    USBHIDDevice *device = tinyusb_get_device_by_report_id(report_id);
    if (device) {
        device->_onOutput(report_id, buffer, reqlen);
        return true;
    }
    return false;
}

static uint16_t tinyusb_on_add_descriptor(uint8_t device_index, uint8_t *dst) {
    uint16_t res = 0;
    uint8_t report_id = 0, reports_num = 0;
    tinyusb_hid_device_t *device = &tinyusb_hid_devices[device_index];
    if (device->device) {
        res = device->device->_onGetDescriptor(dst);
        if (res) {

            esp_hid_report_map_t *hid_report_map = esp_hid_parse_report_map(dst, res);
            if (hid_report_map) {
                if (device->report_ids) { free(device->report_ids); }
                device->reports_num = hid_report_map->reports_len;
                device->report_ids = (uint8_t *)malloc(device->reports_num);
                memset(device->report_ids, 0, device->reports_num);
                reports_num = device->reports_num;

                for (uint8_t i = 0; i < device->reports_num; i++) {
                    if (hid_report_map->reports[i].protocol_mode == ESP_HID_PROTOCOL_MODE_REPORT) {
                        report_id = hid_report_map->reports[i].report_id;
                        for (uint8_t r = 0; r < device->reports_num; r++) {
                            if (!report_id) {
                                // todo: handle better when device has no report ID set
                                break;
                            } else if (report_id == device->report_ids[r]) {
                                // already added
                                reports_num--;
                                break;
                            } else if (!device->report_ids[r]) {
                                // empty slot
                                device->report_ids[r] = report_id;
                                break;
                            }
                        }
                    } else {
                        reports_num--;
                    }
                }
                device->reports_num = reports_num;
                esp_hid_free_report_map(hid_report_map);
            }
        }
    }
    return res;
}

static bool tinyusb_load_enabled_hid_devices() {
    if (tinyusb_hid_device_descriptor != NULL) { return true; }
    tinyusb_hid_device_descriptor = (uint8_t *)malloc(tinyusb_hid_device_descriptor_len);
    if (tinyusb_hid_device_descriptor == NULL) {
        log_e("HID Descriptor Malloc Failed");
        return false;
    }
    uint8_t *dst = tinyusb_hid_device_descriptor;

    for (uint8_t i = 0; i < tinyusb_loaded_hid_devices_num; i++) {
        uint16_t len = tinyusb_on_add_descriptor(i, dst);
        if (!len) {
            break;
        } else {
            dst += len;
        }
    }

    esp_hid_report_map_t *hid_report_map =
        esp_hid_parse_report_map(tinyusb_hid_device_descriptor, tinyusb_hid_device_descriptor_len);
    if (hid_report_map) {
        log_d("Loaded HID Descriptor with the following reports:");
        for (uint8_t i = 0; i < hid_report_map->reports_len; i++) {
            if (hid_report_map->reports[i].protocol_mode == ESP_HID_PROTOCOL_MODE_REPORT) {
                log_d(
                    "  ID: %3u, Type: %7s, Size: %2u, Usage: %8s",
                    hid_report_map->reports[i].report_id,
                    esp_hid_report_type_str(hid_report_map->reports[i].report_type),
                    hid_report_map->reports[i].value_len,
                    esp_hid_usage_str(hid_report_map->reports[i].usage)
                );
            }
        }
        esp_hid_free_report_map(hid_report_map);
    } else {
        log_e("Failed to parse the hid report descriptor!");
        return false;
    }

    return true;
}

extern "C" uint16_t tusb_hid_load_descriptor(uint8_t *dst, uint8_t *itf) {
    if (tinyusb_hid_is_initialized) { return 0; }
    tinyusb_hid_is_initialized = true;

    uint8_t str_index = tinyusb_add_string_descriptor("TinyUSB HID");
    // For keyboard boot protocol, we've already called tinyusb_enable_interface2(reserve_endpoints=true)
    uint8_t ep_in =
        tinyusb_interface_protocol == HID_ITF_PROTOCOL_KEYBOARD ? 1 : tinyusb_get_free_in_endpoint();
    TU_VERIFY(ep_in != 0);
    uint8_t ep_out =
        tinyusb_interface_protocol == HID_ITF_PROTOCOL_KEYBOARD ? 1 : tinyusb_get_free_out_endpoint();
    TU_VERIFY(ep_out != 0);
    uint8_t descriptor[TUD_HID_INOUT_DESC_LEN] = {
        // HID Input & Output descriptor
        // Interface number, string index, protocol, report descriptor len, EP OUT & IN address, size &
        // polling interval
        TUD_HID_INOUT_DESCRIPTOR(
            *itf,
            str_index,
            tinyusb_interface_protocol,
            tinyusb_hid_device_descriptor_len,
            ep_out,
            (uint8_t)(0x80 | ep_in),
            CFG_TUD_ENDOINT_SIZE,
            1
        )
    };
    *itf += 1;
    memcpy(dst, descriptor, TUD_HID_INOUT_DESC_LEN);
    return TUD_HID_INOUT_DESC_LEN;
}

// Invoked when received GET HID REPORT DESCRIPTOR request
// Application return pointer to descriptor, whose contents must exist long enough for transfer to complete
uint8_t const *tud_hid_descriptor_report_cb(uint8_t instance) {
    log_v("instance: %u", instance);
    if (!tinyusb_load_enabled_hid_devices()) { return NULL; }
    return tinyusb_hid_device_descriptor;
}

// Invoked when received SET_PROTOCOL request
// protocol is either HID_PROTOCOL_BOOT (0) or HID_PROTOCOL_REPORT (1)
void tud_hid_set_protocol_cb(uint8_t instance, uint8_t protocol) {
    log_v("instance: %u, protocol:%u", instance, protocol);
    arduino_usb_hid_event_data_t p;
    p.instance = instance;
    p.set_protocol.protocol = protocol;
    arduino_usb_event_post(
        ARDUINO_USB_HID_EVENTS,
        ARDUINO_USB_HID_SET_PROTOCOL_EVENT,
        &p,
        sizeof(arduino_usb_hid_event_data_t),
        portMAX_DELAY
    );
}

// Invoked when received SET_IDLE request. return false will stall the request
// - Idle Rate = 0 : only send report if there is changes, i.e skip duplication
// - Idle Rate > 0 : skip duplication, but send at least 1 report every idle rate (in unit of 4 ms).
bool tud_hid_set_idle_cb(uint8_t instance, uint8_t idle_rate) {
    log_v("instance: %u, idle_rate:%u", instance, idle_rate);
    arduino_usb_hid_event_data_t p;
    p.instance = instance;
    p.set_idle.idle_rate = idle_rate;
    arduino_usb_event_post(
        ARDUINO_USB_HID_EVENTS,
        ARDUINO_USB_HID_SET_IDLE_EVENT,
        &p,
        sizeof(arduino_usb_hid_event_data_t),
        portMAX_DELAY
    );
    return true;
}

// Invoked when received GET_REPORT control request
// Application must fill buffer report's content and return its length.
// Return zero will cause the stack to STALL request
uint16_t tud_hid_get_report_cb(
    uint8_t instance, uint8_t report_id, hid_report_type_t report_type, uint8_t *buffer, uint16_t reqlen
) {
    uint16_t res = tinyusb_on_get_feature(report_id, buffer, reqlen);
    if (!res) {
        log_d(
            "instance: %u, report_id: %u, report_type: %s, reqlen: %u",
            instance,
            report_id,
            tinyusb_hid_device_report_types[report_type],
            reqlen
        );
    }
    return res;
}

// Invoked when received SET_REPORT control request or
// received data on OUT endpoint ( Report ID = 0, Type = 0 )
void tud_hid_set_report_cb(
    uint8_t instance, uint8_t report_id, hid_report_type_t report_type, uint8_t const *buffer,
    uint16_t bufsize
) {
    if (!report_id && (!report_type || report_type == HID_REPORT_TYPE_OUTPUT)) {
        if (!tinyusb_on_set_output(0, buffer, bufsize) &&
            !tinyusb_on_set_output(buffer[0], buffer + 1, bufsize - 1)) {
            log_d(
                "instance: %u, report_id: %u, report_type: %s, bufsize: %u",
                instance,
                buffer[0],
                tinyusb_hid_device_report_types[HID_REPORT_TYPE_OUTPUT],
                bufsize - 1
            );
        }
    } else {
        if (!tinyusb_on_set_feature(report_id, buffer, bufsize)) {
            log_d(
                "instance: %u, report_id: %u, report_type: %s, bufsize: %u",
                instance,
                report_id,
                tinyusb_hid_device_report_types[report_type],
                bufsize
            );
        }
    }
}

USBHID::USBHID(hid_interface_protocol_enum_t itf_protocol) {
    if (!tinyusb_hid_devices_is_initialized) {
        tinyusb_hid_devices_is_initialized = true;
        for (uint8_t i = 0; i < USB_HID_DEVICES_MAX; i++) {
            memset(&tinyusb_hid_devices[i], 0, sizeof(tinyusb_hid_device_t));
        }
        tinyusb_hid_devices_num = 0;
        tinyusb_interface_protocol = itf_protocol;
        tinyusb_enable_interface2(
            USB_INTERFACE_HID,
            TUD_HID_INOUT_DESC_LEN,
            tusb_hid_load_descriptor,
            itf_protocol == HID_ITF_PROTOCOL_KEYBOARD
        );
    }
}

void USBHID::begin() {
    if (tinyusb_hid_device_input_sem == NULL) { tinyusb_hid_device_input_sem = xSemaphoreCreateBinary(); }
    if (tinyusb_hid_device_input_mutex == NULL) { tinyusb_hid_device_input_mutex = xSemaphoreCreateMutex(); }
}

void USBHID::end() {
    if (tinyusb_hid_device_input_sem != NULL) {
        vSemaphoreDelete(tinyusb_hid_device_input_sem);
        tinyusb_hid_device_input_sem = NULL;
    }
    if (tinyusb_hid_device_input_mutex != NULL) {
        vSemaphoreDelete(tinyusb_hid_device_input_mutex);
        tinyusb_hid_device_input_mutex = NULL;
    }
}

bool USBHID::ready(void) { return tud_hid_n_ready(0); }

// TinyUSB is in the process of changing the type of the last argument to
// tud_hid_report_complete_cb(), so extract the type from the version of TinyUSB that we're
// compiled with.
template <class F> struct ArgType;

template <class R, class T1, class T2, class T3> struct ArgType<R (*)(T1, T2, T3)> {
    typedef T1 type1;
    typedef T2 type2;
    typedef T3 type3;
};

typedef ArgType<decltype(&tud_hid_report_complete_cb)>::type3 tud_hid_report_complete_cb_len_t;

void tud_hid_report_complete_cb(
    uint8_t instance, uint8_t const *report, tud_hid_report_complete_cb_len_t len
) {
    if (tinyusb_hid_device_input_sem) { xSemaphoreGive(tinyusb_hid_device_input_sem); }
}

bool USBHID::SendReport(uint8_t id, const void *data, size_t len, uint32_t timeout_ms) {
    if (!tinyusb_hid_device_input_sem || !tinyusb_hid_device_input_mutex) {
        log_e("TX Semaphore is NULL. You must call USBHID::begin() before you can send reports");
        return false;
    }

    if (xSemaphoreTake(tinyusb_hid_device_input_mutex, timeout_ms / portTICK_PERIOD_MS) != pdTRUE) {
        log_e("report %u mutex failed", id);
        return false;
    }

    // If we're configured to support boot protocol, and the host has requested boot protocol, prevent
    // sending of report ID, by passing report ID of 0 to tud_hid_n_report().
    uint8_t effective_id = ((tinyusb_interface_protocol != HID_ITF_PROTOCOL_NONE) &&
                            (tud_hid_n_get_protocol(0) == HID_PROTOCOL_BOOT))
                               ? 0
                               : id;

    bool res = ready();
    if (!res) {
        log_e("not ready");
    } else {
        // The semaphore may be given if the last SendReport() timed out waiting for the report to
        // be sent. Or, tud_hid_report_complete_cb() may be called an extra time, causing the
        // semaphore to be given. In these cases, take the semaphore to clear its state so that
        // we can wait for it to be given after calling tud_hid_n_report().
        xSemaphoreTake(tinyusb_hid_device_input_sem, 0);

        res = tud_hid_n_report(0, effective_id, data, len);
        if (!res) {
            log_e("report %u failed", id);
        } else {
            if (xSemaphoreTake(tinyusb_hid_device_input_sem, timeout_ms / portTICK_PERIOD_MS) != pdTRUE) {
                log_e("report %u wait failed", id);
                res = false;
            }
        }
    }

    xSemaphoreGive(tinyusb_hid_device_input_mutex);
    return res;
}

bool USBHID::addDevice(USBHIDDevice *device, uint16_t descriptor_len) {
    if (device && tinyusb_loaded_hid_devices_num < USB_HID_DEVICES_MAX) {
        if (!tinyusb_enable_hid_device(descriptor_len, device)) { return false; }
        return true;
    }
    return false;
}

void USBHID::onEvent(esp_event_handler_t callback) { onEvent(ARDUINO_USB_HID_ANY_EVENT, callback); }
void USBHID::onEvent(arduino_usb_hid_event_t event, esp_event_handler_t callback) {
    arduino_usb_event_handler_register_with(ARDUINO_USB_HID_EVENTS, event, callback, this);
}

#endif /* CONFIG_TINYUSB_HID_ENABLED */