#include "usb.h"

#include <memory.h>

#include "usb_common.h"
#include "usb_core_types.h"

// #include "utils/gen_queue.h"

#include "desc/usb_desc.h"

#include "usb_driver.h"

// ---------------

#define MAX(a, b) (((a) > (b)) ? (a) : (b))
#define MIN(a, b) (((a) < (b)) ? (a) : (b))

#ifdef USBDBGMSG
#include "../cli/stdio_uart.h"
#define USBMSG(...) MSG(__VA_ARGS__)
#else
#define USBMSG(...)
#endif

// ---------------

#ifndef USB_HIGH_SPEED
#define USB_RX_BUF_SIZE (USB_MAX_FS_PCKT_SIZE_NON_ISOCHRONOUS)
#else
#define USB_RX_BUF_SIZE (USB_MAX_HS_PCKT_SIZE_NON_ISOCHRONOUS)
#endif


static uint8_t tx_assembly_buf[USB_RX_BUF_SIZE] DWORD_ALIGN; // buffer for assembling packets

// ---------------

// // state of changeing address
// typedef enum {
//     USB_ADDRCHG_IDLE = 0, // idle state
//     USB_ADDRCHG_NEW_ADDR_RECVD, // new address received
// } USB_AddressChangeState;

// ----------------

__weak void usb_event_callback(USB_CallbackEvent *cbevt) {
    return;
}

// ----------------

// prepare descriptor for transmission, return with final transmission size
// static uint8_t usb_prepare_descriptor(const uint8_t *desc, uint8_t desc_size, uint8_t req_size) {
//     uint8_t sz = (uint8_t)(MIN(req_size, desc_size));     // determine transmission size
//     memcpy(gs.tx_assembly_buf, desc, sz);                 // copy that portion to the assembly buffer
//     USB_DescHdr *hdr = (USB_DescHdr *)gs.tx_assembly_buf; // obtain header
//     hdr->bLength = MIN(sz, hdr->bLength);                 // write transmit size
//     return sz;
// }

// #define USB_SETUP_STREAM_BUFFER (72)

typedef struct {
    USB_SetupRequest setup_req; // setup request
    uint8_t next_stage;         // next expected stage
    bool out_complete;          // signals if transfer's OUT direction is complete, no later data reception is expected
} USB_SetupTransferState;

static USB_SetupTransferState stups;

void usbcore_init() {
    // init state
    memset(&stups, 0, sizeof(USB_SetupTransferState));
}

void usbcore_process_setup_pckt(const uint8_t *data, uint16_t size, uint8_t stage) {

    // MSG("STUP: %u, %s\n", size, stage == UST_SETUP ? "SETUP" : "DATA");

    // #define USB_PREPARE_DESC_VAR(var, size) sz = usb_prepare_descriptor((const uint8_t *)&(var), sizeof((var)), (size))

    switch (stups.next_stage) {
    case UST_SETUP: {             // expecting SETUP stage
        if (stage == UST_SETUP) { // SETUP received
            // save setup request
            stups.setup_req = *((USB_SetupRequest *)data);

            // if data direction is IN, then don't expect an OUT transaction
            if ((stups.setup_req.bmRequestType.fields.dir == USB_IN) ||
                ((stups.setup_req.bmRequestType.fields.dir == USB_OUT) && (stups.setup_req.wLength == 0))) {
                stups.out_complete = true; // DATA stage is IN
            } else {
                stups.out_complete = false;  // OUT transaction are not exhausted yet (i.e. DATA stage is OUT)
                stups.next_stage = UST_DATA; // next stage is DATA OUT stage
            }
        }
        break;
    }

    case UST_DATA: {             // expecting DATA stage
        if (stage == UST_DATA) { // DATA received
            stups.out_complete = true;
        }

        // step into SETUP state
        stups.next_stage = UST_SETUP;
        break;
    }

    default:
        break;
    }

    // expect status transaction to follow the data phase
    usbdrv_arm_OUT_endpoint(0, 64);

    // only continue processing if the transaction has concluded
    if (!stups.out_complete) {
        return;
    }

#define DETERMINE_TRANSFER_SIZE(desc_size) sz = (MIN((stups.setup_req.wLength), (desc_size)))
#define SET_TRANSMISSION_POINTER(ptr) data = (const uint8_t *)(ptr)

    switch (stups.setup_req.bRequest) {
    case UREQ_GetDescriptor: // GET DESCRIPTOR
    {
        // which descriptor?
        uint8_t desc_type = (stups.setup_req.wValue >> 8);    // get descriptor type
        uint8_t desc_index = (stups.setup_req.wValue & 0xFF); // get descriptor index
        uint8_t sz = 0;
        const uint8_t *data = NULL;

        switch (desc_type) {
        case UD_Device: // DEVICE DESCRIPTOR
            DETERMINE_TRANSFER_SIZE(devDesc.bLength);
            SET_TRANSMISSION_POINTER(&devDesc);
            break;
        case UD_Configuration: // CONFIGURATION DESCRIPTOR
        case UD_OtherSpeedConfiguration: // OTHER SPEED CONFIGURATION DESCRIPTOR
            DETERMINE_TRANSFER_SIZE(confDescs[desc_index]->wTotalLength);
            SET_TRANSMISSION_POINTER(confDescs[desc_index]);

            if (desc_type == UD_OtherSpeedConfiguration) { // change bDescriptorType to Other_Speed_Configuration
                confDescs[desc_index]->bDescriptorType = UD_OtherSpeedConfiguration;
            }
            break;
        case UD_String: // STRING DESCRIPTOR
            DETERMINE_TRANSFER_SIZE(strDescs[desc_index]->bLength);
            SET_TRANSMISSION_POINTER(strDescs[desc_index]);
            break;
        case UD_DeviceQualifier:
            DETERMINE_TRANSFER_SIZE(devQualDesc.bLength);
            SET_TRANSMISSION_POINTER(&devQualDesc);
            break;
        // Descriptors not implemented
        default:
            usbdrv_stall_endpoint(0, USB_IN, true); // stall IN, since request in unsupported
            break;
        }

        // arm data transmission
        usbdrv_arm_IN_endpoint(0, data, sz);
        break;
    }

    case UREQ_SetAddress: { // SET ADDRESS
        uint8_t addr = stups.setup_req.wValue & 0x7F;
        usbdrv_set_address(addr);
        usbdrv_arm_OUT_endpoint(0, 64); // prepare for data OUT stage
        USBDRV_ARM_IN_ZLP(0);           // ZLP IN
        break;
    }
    case UREQ_SetConfiguration: // SET CONFIGURATION
    {
        uint8_t config_id = stups.setup_req.wValue & 0xFF;
        usbdrv_fetch_endpoint_configuration(config_id - 1);
        USBDRV_ARM_IN_ZLP(0);
        break;
    }
    case UREQ_SetInterface: // SET INTERFACE
    {
        // TODO: set configuration
        USBDRV_ARM_IN_ZLP(0);
        break;
    }
    default: { // UNKNOWN REQUEST, pass processing to user application
        USB_CallbackEvent cbevt = {0};
        cbevt.type = USB_CBEVT_UNKNOWN_REQ;
        cbevt.setup_request = &stups.setup_req;
        cbevt.data = (const uint8_t *)data;
        cbevt.size = size;
        cbevt.ep = 0;
        cbevt.dir = USB_OUT;
        cbevt.reply_valid = false;
        usb_event_callback(&cbevt);

        if (cbevt.reply_valid) {
            usbdrv_arm_IN_endpoint(0, cbevt.reply_data, cbevt.reply_size);
        }
        break;
    }
    }

    stups.out_complete = false;
}

void usbcore_process_nonsetup_event(USBDRV_CallbackCompound *cbcpd) {
    USB_CallbackEvent cbevt = {0}; // allocate and clear structure...
    cbevt.ep = cbcpd->ep;          // later only fill nonzero fields

    bool discard_event = false; // only discard if event does not fit any category above
    switch (cbcpd->code) {
    case USB_CBC_OUT: {
        cbevt.type = USB_CBEVT_OUT;
        cbevt.data = cbcpd->data;
        cbevt.size = cbcpd->size;
        cbevt.dir = USB_OUT;
        cbevt.arm_out_endpoint = true;
        break;
    }
    case USB_CBC_IN_FIFOEMPTY:
    case USB_CBC_IN_DONE: {
        cbevt.type = USB_CBEVT_IN;
        if (cbcpd->code == USB_CBC_IN_FIFOEMPTY) { // signals a failed in transfer
            cbevt.subtype = USB_CBEVST_IN_REQ;
        }
        cbevt.dir = USB_IN;
        break;
    }

    default:
        discard_event = true;
        break;
    }

    // if event is not ment to be discarded, then invoke event callback
    if (!discard_event) {
        usb_event_callback(&cbevt);

        // for OUT events, check the autoarm flag
        if (cbevt.dir == USB_OUT) {
            if (cbevt.arm_out_endpoint) {
                usbdrv_autoarm_OUT_endpoint(cbcpd->ep);
            }
        }
    }
}

void usbcore_wake_up_endpoint(uint8_t ep, uint8_t dir) {
    usbdrv_enable_endpoint_interrupt(ep, dir, true);
}

uint32_t usbcore_schedule_transmission(uint8_t ep, const uint8_t *data, uint16_t size) {
    usbdrv_enable_endpoint_interrupt(ep, USB_IN, true);
    return usbdrv_arm_IN_endpoint(ep, data, size);
}

uint32_t usbcore_schedule_reception(uint8_t ep, uint16_t size) {
    return usbdrv_arm_OUT_endpoint(ep, size);
}