#include "usb.h"

#include <memory.h>

#include "usb_callback_event.h"
#include "usb_common.h"
#include "usb_driver_common.h"

#include FLATUSB_DESCRIPTOR_HEADER

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

/** \cond false */
#define MAX(a, b) (((a) > (b)) ? (a) : (b))
#define MIN(a, b) (((a) < (b)) ? (a) : (b))
/** \endcond */

#define USBDRV_ARM_IN_ZLP(ep) drv->arm_IN((ep), NULL, 0) ///< Arm a ZLP IN message

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

static uint8_t tx_assembly_buf[USB_RX_BUF_SIZE] DWORD_ALIGN; ///< Buffer for assembling packets
static Usb_SetupTransferState stups; ///< Setup transfer state.
static UsbDrv_DrvIntf *drv; ///< Mutable pointer to the driver interface

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

/**
 * @fn void usb_event_callback(USB_CallbackEvent *cbevt)
 * Prototype for the USB event callback.
 * 
 * @param cbevt pointer to the event data
 */
__attribute__((weak)) void usb_event_callback(Usb_CallbackEvent *cbevt) {
    return;
}

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

/**
 * Callback function for bus issued Reset is done notification.
 */
static void usbcore_rst_notify() {
    usbcore_reset();
}

/**
 * Callback for Enumeration Done notification.
 */
static void usbcore_enum_notify(uint8_t spd) {
    (void)spd;
}

void usbcore_init(UsbDrv_DrvIntf *drvIntf) {
    usbcore_reset(); // reset USB Core

    drv = drvIntf;                          // store USB driver interface assignments
    drv->rst_notify = usbcore_rst_notify;   // assign Reset Done callback
    drv->enum_notify = usbcore_enum_notify; // assign Enumeration Done callback
}

void usbcore_reset() {
    // init state
    memset(&stups, 0, sizeof(Usb_SetupTransferState));
}

/**
 * Process SETUP packets.
 * 
 * @param data pointer to the SETUP packet
 * @param size size of the packet
 * @param stage stage of the SETUP transaction
 */
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
    drv->arm_OUT(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:
            drv->stall(0, USB_IN, true); // stall IN, since request in unsupported
            break;
        }

        // arm data transmission
        drv->arm_IN(0, data, sz);
        break;
    }

    case UREQ_SetAddress: { // SET ADDRESS
        uint8_t addr = stups.setup_req.wValue & 0x7F;
        drv->set_address(addr);
        drv->arm_OUT(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;
        drv->set_config(config_id - 1);
        USBDRV_ARM_IN_ZLP(0);
        break;
    }
    case UREQ_SetInterface: // SET INTERFACE
    {
        // TODO: set interface
        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) {
            drv->arm_IN(0, cbevt.reply_data, cbevt.reply_size);
        }
        break;
    }
    }

    stups.out_complete = false;
}

/**
 * Process the non-SETUP packets.
 * 
 * @param cbcpd pointer to callback compound
 */
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) {
                drv->autoarm(cbcpd->ep);
            }
        }
    }
}

void usbcore_wake_up_endpoint(uint8_t ep, uint8_t dir) {
    drv->en_ep_irq(ep, dir, true);
}

void usbcore_register_IN_callback(uint8_t ep, UsbDrv_IN_cb cb) {
    drv->reg_IN_cb(ep, cb);
}

uint32_t usbcore_schedule_transmission(uint8_t ep, const uint8_t *data, uint16_t size) {
    drv->en_ep_irq(ep, USB_IN, true);
    return drv->arm_IN(ep, data, size);
}

uint32_t usbcore_schedule_reception(uint8_t ep, uint16_t size) {
    return drv->arm_OUT(ep, size);
}