#include <stdlib.h>
#include <libopencm3/usb/usbd.h>
#include <libopencm3/usb/cdc.h>

#include <SEGGER_RTT.h>
#include <string.h>

#include "rndis.h"
#include "util.h"
#include "rndis_protocol.h"

/* Buffer to be used for control requests. */
static uint8_t usbd_control_buffer[1025];

static const uint8_t station_hwaddr[6] = { RNDIS_HWADDR };
static const uint8_t permanent_hwaddr[6] = { RNDIS_HWADDR };

static const struct usb_device_descriptor dev = {
        .bLength = USB_DT_DEVICE_SIZE,
        .bDescriptorType = USB_DT_DEVICE,
        .bcdUSB = 0x0200,
        .bDeviceClass = 0xE0,  // Wireless Controller
        .bDeviceSubClass = 0,
        .bDeviceProtocol = 0,
        .bMaxPacketSize0 = 64,
        .idVendor = 0x4E44,
        .idProduct = 0x4953,
        .bcdDevice = 0xffff,
        .iManufacturer = 1,
        .iProduct = 2,
        .iSerialNumber = 3,
        .bNumConfigurations = 1,
};

static const struct usb_endpoint_descriptor comm_endp[] = {{
    .bLength = USB_DT_ENDPOINT_SIZE,
    .bDescriptorType = USB_DT_ENDPOINT,
    .bEndpointAddress = 0x83,
    .bmAttributes = USB_ENDPOINT_ATTR_INTERRUPT,
    .wMaxPacketSize = 8,
    .bInterval = 0xf,
}};

static const struct {
    struct usb_cdc_header_descriptor header;
    struct usb_cdc_call_management_descriptor call_mgmt;
    struct usb_cdc_acm_descriptor acm;
    struct usb_cdc_union_descriptor cdc_union;
} __attribute__((packed)) cdcacm_functional_descriptors = {
        .header = {
                .bFunctionLength = sizeof(struct usb_cdc_header_descriptor),
                .bDescriptorType = CS_INTERFACE,
                .bDescriptorSubtype = USB_CDC_TYPE_HEADER,
                .bcdCDC = 0x0110,
        },
        .call_mgmt = {
                .bFunctionLength =
                sizeof(struct usb_cdc_call_management_descriptor),
                .bDescriptorType = CS_INTERFACE,
                .bDescriptorSubtype = USB_CDC_TYPE_CALL_MANAGEMENT,
                .bmCapabilities = 0,
                .bDataInterface = 1,
        },
        .acm = {
                .bFunctionLength = sizeof(struct usb_cdc_acm_descriptor),
                .bDescriptorType = CS_INTERFACE,
                .bDescriptorSubtype = USB_CDC_TYPE_ACM,
                .bmCapabilities = 0,
        },
        .cdc_union = {
                .bFunctionLength = sizeof(struct usb_cdc_union_descriptor),
                .bDescriptorType = CS_INTERFACE,
                .bDescriptorSubtype = USB_CDC_TYPE_UNION,
                .bControlInterface = 0,
                .bSubordinateInterface0 = 1,
        },
};

// RNDIS Communications Control
static const struct usb_interface_descriptor comm_iface[] = {{
    .bLength = USB_DT_INTERFACE_SIZE,
    .bDescriptorType = USB_DT_INTERFACE,
    .bInterfaceNumber = 0,
    .bAlternateSetting = 0,
    .bNumEndpoints = 1,
    // Linux kernel calls this "RNDIS for tethering"
    .bInterfaceClass = 0xE0,  // Wireless Controller Class
    .bInterfaceSubClass = 0x01,
    .bInterfaceProtocol = 0x03,

     // https://docs.microsoft.com/en-us/windows-hardware/drivers/network/communication-class-interface
//    .bInterfaceClass = 0x02,
//    .bInterfaceSubClass = 0x02,
//    .bInterfaceProtocol = 0xff,

    .iInterface = 0,

    .endpoint = comm_endp,

    .extra = &cdcacm_functional_descriptors,
    .extralen = sizeof(cdcacm_functional_descriptors),
}};

static const struct usb_endpoint_descriptor data_endp[] = {{
   .bLength = USB_DT_ENDPOINT_SIZE,
   .bDescriptorType = USB_DT_ENDPOINT,
   .bEndpointAddress = 0x01,
   .bmAttributes = USB_ENDPOINT_ATTR_BULK,
   .wMaxPacketSize = 64,
   .bInterval = 0,
}, {
   .bLength = USB_DT_ENDPOINT_SIZE,
   .bDescriptorType = USB_DT_ENDPOINT,
   .bEndpointAddress = 0x82,
   .bmAttributes = USB_ENDPOINT_ATTR_BULK,
   .wMaxPacketSize = 64,
   .bInterval = 0,
}};

static const struct usb_interface_descriptor data_iface[] = {{
    .bLength = USB_DT_INTERFACE_SIZE,
    .bDescriptorType = USB_DT_INTERFACE,
    .bInterfaceNumber = 1,
    .bAlternateSetting = 0,
    .bNumEndpoints = 2,
    .bInterfaceClass = USB_CLASS_DATA,
    .bInterfaceSubClass = 0,
    .bInterfaceProtocol = 0,
    .iInterface = 0,

    .endpoint = data_endp,
}};

static const struct usb_interface ifaces[] = {{
    .num_altsetting = 1,
    .altsetting = comm_iface,
}, {
    .num_altsetting = 1,
    .altsetting = data_iface,
}};

static const struct usb_config_descriptor config = {
        .bLength = USB_DT_CONFIGURATION_SIZE,
        .bDescriptorType = USB_DT_CONFIGURATION,
        .wTotalLength = 0,
        .bNumInterfaces = 2,
        .bConfigurationValue = 1,
        .iConfiguration = 0,
        .bmAttributes = 0x80,
        .bMaxPower = 0x32,

        .interface = ifaces,
};

static const char *usb_strings[] = {
        "Arti",  // iManufacturer
        "STM32-RNDIS",  // iProduct
        "lrndis demo",  // iSerialNumber
};

static void rndis_query_cmplt32(uint8_t **buf, uint32_t status, uint32_t data)
{
    rndis_query_cmplt_t *c = (rndis_query_cmplt_t *)*buf;
    c->MessageType = REMOTE_NDIS_QUERY_CMPLT;
    c->MessageLength = sizeof(rndis_query_cmplt_t) + 4;
    c->Status = status;
    c->InformationBufferLength = 4;
    c->InformationBufferOffset = 16;
    *(uint32_t *)(c + 1) = data;
}

static void rndis_query_cmplt(uint8_t **buf, uint32_t status, const void *data, uint32_t size){
    rndis_query_cmplt_t *c = (rndis_query_cmplt_t *)*buf;
    c->MessageType = REMOTE_NDIS_QUERY_CMPLT;
    c->MessageLength = sizeof(rndis_query_cmplt_t) + size;
    c->Status = status;
    c->InformationBufferLength = size;
    c->InformationBufferOffset = 16;
    memcpy(c + 1, data, size);
}

static void rndis_handle_query(uint8_t **buf, uint16_t *len)
{
    union {
        void                    *buf;
        rndis_generic_msg_t     *header;
        rndis_query_msg_t       *get;
        rndis_query_cmplt_t     *get_c;
    } u;

    u.buf = *buf;

    SEGGER_RTT_printf(0, "NDIS Query oid: 0x%08x, len: 0x%02x\n", u.get->Oid, u.get->InformationBufferLength);

    switch (u.get->Oid)
    {
        case OID_GEN_PHYSICAL_MEDIUM:
            rndis_query_cmplt32(buf, RNDIS_STATUS_SUCCESS, NDIS_MEDIUM_802_3);
            return;
        case OID_802_3_PERMANENT_ADDRESS:
            rndis_query_cmplt(buf, RNDIS_STATUS_SUCCESS, &permanent_hwaddr, sizeof(permanent_hwaddr));
            return;
        default:
            rndis_query_cmplt(buf, RNDIS_STATUS_FAILURE, NULL, 0);
            return;
    }
}

static enum usbd_request_return_codes rndis_control_request(usbd_device *usbd_dev, struct usb_setup_data *req, uint8_t **buf,
                                                             uint16_t *len, void (**complete)(usbd_device *usbd_dev, struct usb_setup_data *req))
{
    (void)req;
    (void)complete;

    union {
        void                        *buf;
        rndis_generic_msg_t         *header;
        rndis_initialize_msg_t      *init;
        rndis_initialize_cmplt_t    *init_c;
        rndis_query_msg_t           *get;
        rndis_query_cmplt_t         *get_c;
        rndis_set_msg_t             *set;
        rndis_set_cmplt_t           *set_c;
        rndis_indicate_status_t     *indic;
    } u;

    //SEGGER_RTT_printf(0, "bmRequestType: 0x%02x, bRequest: 0x%02x\n", req->bmRequestType, req->bRequest);

    u.buf = *buf;
    switch (u.header->MessageType)
    {
        case REMOTE_NDIS_INITIALIZE_MSG:
            SEGGER_RTT_printf(0, "NDIS Init: HOST Ver %d.%d, Max MTU: %dbytes \n", u.init->MajorVersion, u.init->MinorVersion, u.init->MaxTransferSize);
            u.init_c->MessageType = REMOTE_NDIS_INITIALIZE_CMPLT;
            return USBD_REQ_HANDLED;
        case REMOTE_NDIS_INITIALIZE_CMPLT:
            SEGGER_RTT_printf(0, "NDIS Init complete, MTU: %d\n", RNDIS_RX_BUFFER_SIZE);
            u.init_c->MessageType = REMOTE_NDIS_INITIALIZE_CMPLT;
            u.init_c->MessageLength = sizeof(rndis_initialize_cmplt_t);
            // u.int_c->RequestId stays the same
            u.init_c->Status = RNDIS_STATUS_SUCCESS;
            u.init_c->MajorVersion = RNDIS_MAJOR_VERSION;
            u.init_c->MinorVersion = RNDIS_MINOR_VERSION;
            u.init_c->DeviceFlags = RNDIS_DF_CONNECTIONLESS;
            u.init_c->Medium = RNDIS_MEDIUM_802_3;
            u.init_c->MaxPacketsPerTransfer = 1;
            u.init_c->MaxTransferSize = RNDIS_RX_BUFFER_SIZE;
            u.init_c->PacketAlignmentFactor = 0;
            u.init_c->AfListOffset = 0;
            u.init_c->AfListSize = 0;
            //*len = sizeof(rndis_initialize_cmplt_t);
            usbd_ep_write_packet(usbd_dev, 0x83, "\x01\x00\x00\x00\x00\x00\x00\x00", 8);
            return USBD_REQ_HANDLED;
        case REMOTE_NDIS_QUERY_MSG:
            rndis_handle_query(buf, len);
            usbd_ep_write_packet(usbd_dev, 0x83, "\x01\x00\x00\x00\x00\x00\x00\x00", 8);
            //*len = u.get_c->MessageLength;
            return USBD_REQ_HANDLED;
        case REMOTE_NDIS_QUERY_CMPLT:
            SEGGER_RTT_printf(0, "NDIS Query complete\n");
            *len = u.get_c->MessageLength;
            usbd_ep_write_packet(usbd_dev, 0x83, "\x01\x00\x00\x00\x00\x00\x00\x00", 8);
            return USBD_REQ_HANDLED;
        case REMOTE_NDIS_SET_MSG:
            SEGGER_RTT_printf(0, "NDIS Set oid: 0x%08x, len 0x%02x\n", u.set->Oid, u.set->InformationBufferLength);
            u.set_c->MessageType = REMOTE_NDIS_SET_CMPLT;
            u.set_c->MessageLength = sizeof(rndis_set_cmplt_t);
            u.set_c->Status = RNDIS_STATUS_SUCCESS;
            *len = u.set_c->MessageLength;
            usbd_ep_write_packet(usbd_dev, 0x83, "\x01\x00\x00\x00\x00\x00\x00\x00", 8);
            return USBD_REQ_HANDLED;
        case REMOTE_NDIS_SET_CMPLT:
            SEGGER_RTT_printf(0, "NDIS Set complete\n");
            *len = u.set_c->MessageLength;
            usbd_ep_write_packet(usbd_dev, 0x83, "\x01\x00\x00\x00\x00\x00\x00\x00", 8);
            return USBD_REQ_HANDLED;
    }
    SEGGER_RTT_printf(0, "NDIS unknown MessageType: 0x%08x, MessageLength: %d, max len: %d\n", u.header->MessageType, u.header->MessageLength, *len);
    return USBD_REQ_NEXT_CALLBACK;
}

static void rndis_data_rx_cb(usbd_device *usbd_dev, uint8_t ep)
{
    (void)ep;

    char buf[64];
    int len = usbd_ep_read_packet(usbd_dev, 0x01, buf, 64);

    if (len) {
        SEGGER_RTT_printf(0, "RX: %d\n", len);
        hexdump(buf, len, 16);
    }
}

static void rndis_set_config(usbd_device *usbd_dev, uint16_t wValue)
{
    (void)wValue;
    (void)usbd_dev;
    //SEGGER_RTT_printf(0, "USB Set config %d\n", wValue);

    usbd_ep_setup(usbd_dev, 0x01, USB_ENDPOINT_ATTR_BULK, 64, rndis_data_rx_cb);
    usbd_ep_setup(usbd_dev, 0x82, USB_ENDPOINT_ATTR_BULK, 64, NULL);
    // Interrupt endpoint is used to only send notification that some data is ready to read
    usbd_ep_setup(usbd_dev, 0x83, USB_ENDPOINT_ATTR_INTERRUPT, 8, NULL);

    usbd_register_control_callback(usbd_dev,
                                   USB_REQ_TYPE_CLASS | USB_REQ_TYPE_INTERFACE,
                                   USB_REQ_TYPE_TYPE | USB_REQ_TYPE_RECIPIENT,
                                   rndis_control_request);
}

usbd_device *start_rdnis(void) {
    usbd_device *usbd_dev;

    usbd_dev = usbd_init(
            &st_usbfs_v1_usb_driver, &dev, &config,
            usb_strings, 3,
            usbd_control_buffer, sizeof(usbd_control_buffer)
    );
    usbd_register_set_config_callback(usbd_dev, rndis_set_config);

    return usbd_dev;
}