// 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 <ddk/binding.h>
#include <ddk/debug.h>
#include <ddk/protocol/usb.h>
#include <ddk/protocol/usb/composite.h>
#include <usb/usb-request.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include "usb-composite.h"
#include "usb-interface.h"

namespace usb_composite {

static inline const usb_descriptor_header_t* NextDescriptor(const usb_descriptor_header_t* header) {
    return reinterpret_cast<const usb_descriptor_header_t*>(
                            reinterpret_cast<const uint8_t*>(header) + header->bLength);
}

zx_status_t UsbInterface::Create(zx_device_t* parent, UsbComposite* composite,
                                 const ddk::UsbProtocolClient& usb,
                                 const usb_interface_descriptor_t* interface_desc,
                                 size_t desc_length,
                                 fbl::RefPtr<UsbInterface>* out_interface) {
    fbl::AllocChecker ac;
    auto interface = fbl::MakeRefCountedChecked<UsbInterface>(&ac, composite->zxdev(), composite,
                                                              usb);
    if (!ac.check()) {
        return ZX_ERR_NO_MEMORY;
    }

    auto* device_desc = composite->device_descriptor();
    uint8_t usb_class, usb_subclass, usb_protocol;

    if (interface_desc->bInterfaceClass == 0) {
        usb_class = device_desc->bDeviceClass;
        usb_subclass = device_desc->bDeviceSubClass;
        usb_protocol = device_desc->bDeviceProtocol;
    } else {
        // class/subclass/protocol defined per-interface
        usb_class = interface_desc->bInterfaceClass;
        usb_subclass = interface_desc->bInterfaceSubClass;
        usb_protocol = interface_desc->bInterfaceProtocol;
   }

    auto status = interface->Init(interface_desc, desc_length, interface_desc->bInterfaceNumber,
                                  usb_class, usb_subclass, usb_protocol);
    if (status != ZX_OK) {
        return status;
    }

    status = interface->ConfigureEndpoints(interface_desc->bInterfaceNumber, 0);
    if (status != ZX_OK) {
        return status;
    }

    *out_interface = interface;
    return ZX_OK;
}

zx_status_t UsbInterface::Create(zx_device_t* parent, UsbComposite* composite,
                                 const ddk::UsbProtocolClient& usb,
                                 const usb_interface_assoc_descriptor_t* assoc_desc,
                                 size_t desc_length,
                                 fbl::RefPtr<UsbInterface>* out_interface) {
    fbl::AllocChecker ac;
    auto interface = fbl::MakeRefCountedChecked<UsbInterface>(&ac, composite->zxdev(), composite,
                                                              usb);
    if (!ac.check()) {
        return ZX_ERR_NO_MEMORY;
    }

    auto* device_desc = composite->device_descriptor();
    uint8_t usb_class, usb_subclass, usb_protocol;

    if (assoc_desc->bFunctionClass == 0) {
        usb_class = device_desc->bDeviceClass;
        usb_subclass = device_desc->bDeviceSubClass;
        usb_protocol = device_desc->bDeviceProtocol;
    } else {
        // class/subclass/protocol defined per-interface
        usb_class = assoc_desc->bFunctionClass;
        usb_subclass = assoc_desc->bFunctionSubClass;
        usb_protocol = assoc_desc->bFunctionProtocol;
   }

   // Interfaces in an IAD interface collection must be contiguous.
    auto last_interface_id = assoc_desc->bFirstInterface + assoc_desc->bInterfaceCount - 1;
    auto status = interface->Init(assoc_desc, desc_length, static_cast<uint8_t>(last_interface_id),
                                  usb_class, usb_subclass, usb_protocol);
    if (status != ZX_OK) {
        return status;
    }

    auto* header = reinterpret_cast<const usb_descriptor_header_t*>(assoc_desc);
    auto* end = reinterpret_cast<const usb_descriptor_header_t*>(
                                    reinterpret_cast<const uint8_t*>(assoc_desc) + desc_length);

    while (header < end) {
        if (header->bDescriptorType == USB_DT_INTERFACE) {
            auto* intf_desc = reinterpret_cast<const usb_interface_descriptor_t*>(header);
            if (intf_desc->bAlternateSetting == 0) {
                zx_status_t status = interface->ConfigureEndpoints(intf_desc->bInterfaceNumber, 0);
                if (status != ZX_OK) {
                    return status;
                }
            }
        }
        header = NextDescriptor(header);
    }

    *out_interface = interface;
    return ZX_OK;
}

zx_status_t UsbInterface::Init(const void* descriptors, size_t desc_length,
                               uint8_t last_interface_id, uint8_t usb_class,
                               uint8_t usb_subclass, uint8_t usb_protocol) {
    fbl::AllocChecker ac;
    auto desc_bytes = new (&ac) uint8_t[desc_length];
    if (!ac.check()) {
        return ZX_ERR_NO_MEMORY;
    }
    memcpy(desc_bytes, descriptors, desc_length);
    descriptors_.reset(desc_bytes, desc_length);
    last_interface_id_ = last_interface_id;
    usb_class_ = usb_class;
    usb_subclass_ = usb_subclass;
    usb_protocol_ = usb_protocol;

    return ZX_OK;
}

zx_status_t UsbInterface::DdkGetProtocol(uint32_t proto_id, void* protocol) {
    switch (proto_id) {
    case ZX_PROTOCOL_USB: {
        auto* proto = static_cast<usb_protocol_t*>(protocol);
        proto->ctx = this;
        proto->ops = &usb_protocol_ops_;
        return ZX_OK;
    }
    case ZX_PROTOCOL_USB_COMPOSITE: {
        auto* proto = static_cast<usb_composite_protocol_t*>(protocol);
        proto->ctx = this;
        proto->ops = &usb_composite_protocol_ops_;
        return ZX_OK;
    }
    default:
        return ZX_ERR_NOT_SUPPORTED;
    }
}

void UsbInterface::DdkUnbind() {
    DdkRemove();
}

void UsbInterface::DdkRelease() {
    // Release the reference now that devmgr no longer has a pointer to this object.
    __UNUSED bool dummy = Release();
}

// for determining index into active_endpoints[]
// bEndpointAddress has 4 lower order bits, plus high bit to signify direction
// shift high bit to bit 4 so index is in range 0 - 31.
static inline uint8_t GetEndpointIndex(const usb_endpoint_descriptor_t* ep) {
    return static_cast<uint8_t>(((ep)->bEndpointAddress & 0x0F) | ((ep)->bEndpointAddress >> 3));
}

zx_status_t UsbInterface::ConfigureEndpoints(uint8_t interface_id, uint8_t alt_setting) {
    usb_endpoint_descriptor_t* new_endpoints[USB_MAX_EPS] = {};
    bool interface_endpoints[USB_MAX_EPS] = {};
    zx_status_t status = ZX_OK;

    // iterate through our descriptors to find which endpoints should be active
    auto* header = reinterpret_cast<const usb_descriptor_header_t*>(descriptors_.get());
    auto* end = reinterpret_cast<const usb_descriptor_header_t*>(descriptors_.get()
                                                                 + descriptors_.size());
    int cur_interface = -1;

    bool enable_endpoints = false;
    while (header < end) {
        if (header->bDescriptorType == USB_DT_INTERFACE) {
            auto* intf_desc = reinterpret_cast<const usb_interface_descriptor_t*>(header);
            cur_interface = intf_desc->bInterfaceNumber;
            enable_endpoints = (intf_desc->bAlternateSetting == alt_setting);
        } else if (header->bDescriptorType == USB_DT_ENDPOINT && cur_interface == interface_id) {
            usb_endpoint_descriptor_t* ep = (usb_endpoint_descriptor_t*)header;
            auto ep_index = GetEndpointIndex(ep);
            interface_endpoints[ep_index] = true;
            if (enable_endpoints) {
                new_endpoints[ep_index] = ep;
            }
        }
        header = NextDescriptor(header);
    }

    // update to new set of endpoints
    // FIXME - how do we recover if we fail half way through processing the endpoints?
    for (size_t i = 0; i < countof(new_endpoints); i++) {
        if (interface_endpoints[i]) {
            auto* old_ep = active_endpoints_[i];
            auto* new_ep = new_endpoints[i];
            if (old_ep != new_ep) {
                if (old_ep) {
                    zx_status_t ret = usb_.EnableEndpoint(old_ep, nullptr, false);
                    if (ret != ZX_OK) status = ret;
                }
                if (new_ep) {
                    usb_ss_ep_comp_descriptor_t* ss_comp_desc = nullptr;
                    usb_descriptor_header_t* next =
                                (usb_descriptor_header_t *)((uint8_t *)new_ep + new_ep->bLength);
                    if (next + sizeof(*ss_comp_desc) <= end
                        && next->bDescriptorType == USB_DT_SS_EP_COMPANION) {
                        ss_comp_desc = (usb_ss_ep_comp_descriptor_t *)next;
                    }
                    zx_status_t ret = usb_.EnableEndpoint(new_ep, ss_comp_desc, true);
                    if (ret != ZX_OK) {
                        status = ret;
                    }
                }
                active_endpoints_[i] = new_ep;
            }
        }
    }
    return status;
}

zx_status_t UsbInterface::UsbControlOut(uint8_t request_type, uint8_t request, uint16_t value,
                                        uint16_t index, zx_time_t timeout, const void* write_buffer,
                                        size_t write_size) {
    return usb_.ControlOut(request_type, request, value, index, timeout, write_buffer, write_size);
}

zx_status_t UsbInterface::UsbControlIn(uint8_t request_type, uint8_t request, uint16_t value,
                                       uint16_t index, zx_time_t timeout, void* out_read_buffer,
                                       size_t read_size, size_t* out_read_actual) {
    return usb_.ControlIn(request_type, request, value, index, timeout, out_read_buffer, read_size,
                          out_read_actual);
}

void UsbInterface::UsbRequestQueue(usb_request_t* usb_request,
                                   const usb_request_complete_t* complete_cb) {
    usb_.RequestQueue(usb_request, complete_cb);
}

usb_speed_t UsbInterface::UsbGetSpeed() {
    return usb_.GetSpeed();
}

zx_status_t UsbInterface::UsbSetInterface(uint8_t interface_number, uint8_t alt_setting) {
    return composite_->SetInterface(interface_number, alt_setting);
}

uint8_t UsbInterface::UsbGetConfiguration() {
    return usb_.GetConfiguration();
}

zx_status_t UsbInterface::UsbSetConfiguration(uint8_t configuration) {
    return usb_.SetConfiguration(configuration);
}

zx_status_t UsbInterface::UsbEnableEndpoint(const usb_endpoint_descriptor_t* ep_desc,
                                            const usb_ss_ep_comp_descriptor_t* ss_com_desc,
                                            bool enable) {
    return ZX_ERR_NOT_SUPPORTED;
}

zx_status_t UsbInterface::UsbResetEndpoint(uint8_t ep_address) {
    return usb_.ResetEndpoint(ep_address);
}

zx_status_t UsbInterface::UsbResetDevice() {
    return usb_.ResetDevice();
}

size_t UsbInterface::UsbGetMaxTransferSize(uint8_t ep_address) {
    return usb_.GetMaxTransferSize(ep_address);
}

uint32_t UsbInterface::UsbGetDeviceId() {
    return usb_.GetDeviceId();
}

void UsbInterface::UsbGetDeviceDescriptor(usb_device_descriptor_t* out_desc) {
    return usb_.GetDeviceDescriptor(out_desc);
}

zx_status_t UsbInterface::UsbGetConfigurationDescriptorLength(uint8_t configuration,
                                                              size_t* out_length) {
    return usb_.GetConfigurationDescriptorLength(configuration, out_length);
}

zx_status_t UsbInterface::UsbGetConfigurationDescriptor(uint8_t configuration,
                                                        void* out_desc_buffer,
                                                        size_t desc_size, size_t*
                                                        out_desc_actual) {
    return usb_.GetConfigurationDescriptor(configuration, out_desc_buffer,
                                           desc_size, out_desc_actual);
}

size_t UsbInterface::UsbGetDescriptorsLength() {
    return descriptors_.size();
}

void UsbInterface::UsbGetDescriptors(void* out_descs_buffer, size_t descs_size,
                                     size_t* out_descs_actual) {
    size_t length = descriptors_.size();
    if (length > descs_size) {
        length = descs_size;
    }
    memcpy(out_descs_buffer, descriptors_.get(), length);
    *out_descs_actual = length;
}

size_t UsbInterface::UsbCompositeGetAdditionalDescriptorLength() {
    auto* config = composite_->GetConfigurationDescriptor();
    auto* header = NextDescriptor(reinterpret_cast<const usb_descriptor_header_t*>(config));
    auto* end = reinterpret_cast<const usb_descriptor_header_t*>(
                    reinterpret_cast<const uint8_t*>(config) + le16toh(config->wTotalLength));

    const usb_interface_descriptor_t* interface = nullptr;
    while (header < end) {
        if (header->bDescriptorType == USB_DT_INTERFACE) {
            usb_interface_descriptor_t* test_intf = (usb_interface_descriptor_t*)header;
            // We are only interested in descriptors past the last stored descriptor
            // for the current interface.
            if (test_intf->bAlternateSetting == 0 &&
                test_intf->bInterfaceNumber > last_interface_id_) {
                interface = test_intf;
                break;
            }
        }
        header = NextDescriptor(header);
    }
    if (!interface) {
        return 0;
    }
    return reinterpret_cast<uintptr_t>(end) - reinterpret_cast<uintptr_t>(interface);
}

zx_status_t UsbInterface::UsbCompositeGetAdditionalDescriptorList(uint8_t* out_desc_list,
                                                                  size_t desc_count,
                                                                  size_t* out_desc_actual) {
    return composite_->GetAdditionalDescriptorList(last_interface_id_, out_desc_list, desc_count,
                                                   out_desc_actual);
}

zx_status_t UsbInterface::UsbGetStringDescriptor(uint8_t desc_id, uint16_t lang_id,
                                                 uint16_t* out_lang_id,
                                                 void* out_string_buffer, size_t string_size,
                                                 size_t* out_string_actual) {
    return usb_.GetStringDescriptor(desc_id, lang_id, out_lang_id, out_string_buffer, string_size,
                                    out_string_actual);
}

zx_status_t UsbInterface::UsbCancelAll(uint8_t ep_address) {
    return usb_.CancelAll(ep_address);
}

uint64_t UsbInterface::UsbGetCurrentFrame() {
    return usb_.GetCurrentFrame();
}

size_t UsbInterface::UsbGetRequestSize() {
    return usb_.GetRequestSize();
}

zx_status_t UsbInterface::UsbCompositeClaimInterface(const usb_interface_descriptor_t* desc,
                                                     uint32_t length) {

    auto status = composite_->ClaimInterface(desc->bInterfaceNumber);
    if (status != ZX_OK) {
        return status;
    }
    // Copy claimed interface descriptors to end of descriptor array.
    fbl::AllocChecker ac;
    size_t old_length = descriptors_.size();
    size_t new_length = old_length + length;
    auto* new_descriptors = new (&ac) uint8_t[new_length];
    if (!ac.check()) {
        return ZX_ERR_NO_MEMORY;
    }
    memcpy(new_descriptors, descriptors_.get(), old_length);
    memcpy(new_descriptors + old_length, desc, length);
    descriptors_.reset(new_descriptors, new_length);

    if (desc->bInterfaceNumber > last_interface_id_) {
        last_interface_id_ = desc->bInterfaceNumber;
    }
    return ZX_OK;
}

bool UsbInterface::ContainsInterface(uint8_t interface_id) {
    auto* header = reinterpret_cast<const usb_descriptor_header_t*>(descriptors_.get());
    auto* end = reinterpret_cast<const usb_descriptor_header_t*>(descriptors_.get()
                                                                 + descriptors_.size());

    while (header < end) {
        if (header->bDescriptorType == USB_DT_INTERFACE) {
            auto* intf_desc = reinterpret_cast<const usb_interface_descriptor_t*>(header);
            if (intf_desc->bInterfaceNumber == interface_id) {
                return true;
            }
        }
        header = NextDescriptor(header);
    }
    return false;
}

zx_status_t UsbInterface::SetAltSetting(uint8_t interface_id, uint8_t alt_setting) {
    zx_status_t status = ConfigureEndpoints(interface_id, alt_setting);
    if (status != ZX_OK) {
        return status;
    }

    return UsbControlOut(USB_DIR_OUT | USB_TYPE_STANDARD | USB_RECIP_INTERFACE,
                         USB_REQ_SET_INTERFACE, alt_setting, interface_id, ZX_TIME_INFINITE,
                         nullptr, 0);
}

} // namespace usb_composite