#include "eem.h"

#include <memory.h>

#include "../usb.h"

#include <etherlib/dynmem.h>
#include <etherlib/eth_interface.h>

#include <standard_output/standard_output.h>

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

static USB_EemState eems = {0};

// static uint8_t netbBuf[USB_EEM_FIFO_SIZE];  // Network bound FIFO memory
static uint8_t hostbBuf[USB_EEM_FIFO_SIZE]; // Host bound FIFO memory

#define EEM_EVENT_QUEUE_LENGTH (16)
#define EEM_FRAME_QUEUE_LENGTH (8)
#define EEM_PCKT_SIZE (64)
#define EEM_TRANSFER_UNIT (128)

void usb_eem_thread(void *);

void usb_eem_init(uint8_t data_ep) {
    // clear the state
    memset(&eems, 0, sizeof(USB_EemState));

    // store data endpoint number
    eems.data_ep = data_ep;

    // create event queue
    eems.eventQueue = osMessageQueueNew(EEM_EVENT_QUEUE_LENGTH, sizeof(uint8_t), NULL);

    // create network bound and host bound FIFOs
    // bfifo_create(&(eems.hostbFifo), netbBuf, USB_EEM_FIFO_SIZE);
    bfifo_create(&(eems.netbFifo), hostbBuf, USB_EEM_FIFO_SIZE);

    // create queue for hostbound frames
    eems.hostbFrameQ = osMessageQueueNew(EEM_FRAME_QUEUE_LENGTH, sizeof(USB_EemFrame *), NULL);

    // no transmission or reception is in progress
    eems.hostbState = EEM_IDLE;
    eems.netbState = EEM_IDLE;

    // turn network bound autoarm on
    eems.netbAutoArm = true;

    // no EchoResponse is queued
    eems.echoRespQueued = false;

    // create thread
    osThreadAttr_t attr;
    bzero(&attr, sizeof(osThreadAttr_t));
    attr.priority = osPriorityNormal;
    attr.stack_size = 2048;
    attr.name = "eem";
    eems.th = osThreadNew(usb_eem_thread, NULL, &attr);

    // EEM is operational
    eems.initialized = true;
}

void usb_eem_push_event(USB_EemEvent evt) {
    osMessageQueuePut(eems.eventQueue, &evt, 0, 0);
}

USB_EemEvent usb_eem_pop_event() {
    USB_EemEvent evt;
    osMessageQueueGet(eems.eventQueue, &evt, 0, osWaitForever);
    return evt;
}

int usb_eem_process_and_return(USB_CallbackEvent *cbevt) {
    switch (cbevt->type) {
    case USB_CBEVT_UNKNOWN_REQ:
        break;

    case USB_CBEVT_OUT:
        if ((cbevt->ep == eems.data_ep)) {                          // verify endpoint
            bfifo_push(&(eems.netbFifo), cbevt->data, cbevt->size); // store portion of netbound packet
            usb_eem_push_event(EEM_EVT_NETBOUND_PCKT_RECEIVED);     // push notification

            if (bfifo_get_free(&eems.netbFifo) < EEM_PCKT_SIZE) {   // don't autoarm if OUT no more packets can be stored
                eems.netbAutoArm = false;
                cbevt->arm_out_endpoint = false;
            }
        }
        break;

    case USB_CBEVT_IN:
        if (cbevt->ep == eems.data_ep) {                                                                // verify endpoint number
            if ((eems.hostbFrame != NULL) && (eems.hostbFrameIndex < eems.hostbFrameLength)) {          // if there's something in the FIFO
                uint8_t *readPtr = ((uint8_t *)eems.hostbFrame) + eems.hostbFrameIndex;                 // calculate read pointer
                uint16_t bytesLeft = eems.hostbFrameLength - eems.hostbFrameIndex;                      // calculate bytes left
                uint32_t writeSize = usbcore_schedule_transmission(eems.data_ep, readPtr, bytesLeft);   // attempt to schedule transmission
                
                //bool enableZLP = (bytesLeft < EEM_TRANSFER_UNIT) && ((bytesLeft % EEM_PCKT_SIZE) == 0); // ZLP transmission should be enabled if this was the last write and transfer size is integer multiple of the USB packet size
                //cbevt->enable_autozlp = enableZLP;

                eems.hostbFrameIndex += writeSize; // advance frame index
                //MSG("U: %u T: %u\n", writeSize, bytesLeft);

                // if (eems.hostbFrameIndex >= eems.hostbFrameLength) {
                //     MSG("---\n");
                // }
            } else {                                                  // no data is waiting for transmission
                //usbcore_schedule_transmission(eems.data_ep, NULL, 0); // send ZLP
            }

            // if FIFO is empty or flushed, then send notification
            if ((eems.hostbFrame == NULL) || (eems.hostbFrameIndex >= eems.hostbFrameLength)) {
                usb_eem_push_event(EEM_EVT_HOSTBOUND_TRANSFER_DONE); // push notification
            }
        }
        break;

    default:
        break;
    }

    return 0;
}

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

#define EEM_HEADER_SIZE (2)         // EEM header size in bytes
#define EEM_HEADER_TYPE_DEFAULT (0) // default header type code
#define EEM_HEADER_TYPE_CMD (1)     // CMD header type code

// TODO rename: swap GET and the next word
#define EEM_IS_CMD_PCKT(h) ((((h)[1] & 0x80) >> 7) == EEM_HEADER_TYPE_CMD)              // checks if header signals a following command packet or not
#define EEM_GET_CMD(h) ((h[1] >> 3) & 0b111)                                            // get the actual command code from the header
#define EEM_GET_CMD_PARAM(h) (((uint16_t)(h)[0]) + ((((uint16_t)(h)[1]) & 0b111) << 8)) // get command arguments from the header

#define EEM_GET_DEFAULT_CRC_PRESENT(h) (((h)[1] >> 6) & 0b1)                                      // checks if CRC is present
#define EEM_GET_DEFAULT_FRAME_LENGTH(h) (((uint16_t)(h)[0]) + ((((uint16_t)(h)[1]) & 0x3F) << 8)) // get Ethernet frame length (14-bits)

#define EEM_BUILD_HEADER(type, crc_present, len) (((uint16_t)(((type) & 0b1) << 15)) | ((uint16_t)(((crc_present) & 0b1) << 14)) | ((len) & 0x3FFF))
#define EEM_HEADER_GET_LENGTH(h) (((h) & 0x3FFF))

void usb_eem_process_networkbound() {
    // extract netbound FIFO pointer for easier handling
    BFifo *bf = &(eems.netbFifo);

    switch (eems.netbState) {
    case EEM_IDLE: { /* IDLE state, can start a new operation */
        // verify, that the packet is big enough to contain the header
        if (bfifo_get_used(bf) < EEM_HEADER_SIZE) {
            break; // if not, then break (either a single byte, or a ZLP was received)
        }

        // certainly, we can acquire the header at this point

        // get header
        uint8_t header[EEM_HEADER_SIZE];
        bfifo_read(bf, header, EEM_HEADER_SIZE); // read
        bfifo_pop(bf, 2, 0);                     // pop (with no timeout)

        // check if this packet was a command packet
        if (EEM_IS_CMD_PCKT(header)) {
            // if it is was a command, then fetch its fields
            uint8_t bmEEMCmd = EEM_GET_CMD(header);
            uint16_t bmEEMCmdParam = EEM_GET_CMD_PARAM(header);

            // process only Echo commands (only if no previous EchoResponse was queued)
            if ((bmEEMCmd == EEM_ECHO) && (!eems.echoRespQueued)) {
                // store payload length
                eems.netbSizeLeft = bmEEMCmdParam; // payload comes in the 11-bit parameter

                // signal, that we will respond at the end the Echo packet
                eems.echoRespQueued = true;

                // step into the 'transfer in progress' state
                eems.netbState = EEM_TRANSFER_IN_PROGRESS;

                MSG("ECHO!\n");
            }
        } else { // this is a regular packet
            // store Ethernet frame length
            eems.netbFrameSize = EEM_GET_DEFAULT_FRAME_LENGTH(header); // set the frame size

            // process packet with non-zero length
            if (eems.netbFrameSize != 0) {
                eems.netbSizeLeft = eems.netbFrameSize; // the full frame is ahead

                // step into the 'transfer in progress' state
                eems.netbState = EEM_TRANSFER_IN_PROGRESS;

                // MSG("DEFAULT! %u\n", eems.netbFrameSize);
            }
        }

    } break;

    case EEM_TRANSFER_IN_PROGRESS: {                    /* transfer is in progress, collect bytes */
        if (bfifo_get_used(bf) >= eems.netbFrameSize) { // if the full packet is in the buffer, then read it!
            // create raw packet if Ethernet interface is defined
            if (eems.intf != NULL) {
                RawPckt raw;
                bzero(&raw, sizeof(RawPckt));
                raw.size = eems.netbFrameSize;        // set size
                raw.payload = dynmem_alloc(raw.size); // allocate buffer
                if (raw.payload != NULL) {
                    bfifo_read(bf, raw.payload, raw.size); // copy payload to raw packet
                    ethinf_transmit(eems.intf, &raw);      // transmit packet
                } else {
                    MSG("EEM -> ETH packet allocation failed!\n");
                }
            }

            // pop data from the packet
            bfifo_pop(bf, eems.netbFrameSize, 0);

            // if free size is greater then the size of a single USB packet, then arm the endpoint
            if ((eems.netbAutoArm == false) && (bfifo_get_free(bf) > EEM_PCKT_SIZE)) {
                eems.netbAutoArm = true;
                usbcore_schedule_reception(eems.data_ep, EEM_PCKT_SIZE);
            }

            // MSG("PKT OK! %u\n", eems.netbFrameSize);

            eems.netbState = EEM_IDLE;
        }
    } break;
    }
}

void usb_eem_process_hostbound() {
    /*// check that forward buffer is indeed empty
    BFifo * bf = &eems.hostbFifo;
    if (bfifo_get_used(bf) > 0) {
        return; // we don't want to overwrite concatenate transfers
    }*/

    // check if transmission is done indeed
    if ((eems.hostbFrame != NULL) && (eems.hostbFrameIndex < eems.hostbFrameLength)) {
        return;
    }

    // invalidate frame
    USB_EemFrame * oldEemFrame = eems.hostbFrame;
    eems.hostbFrame = NULL;

    // release frame
    dynmem_free(oldEemFrame);

    // see if there are some packets waiting for transmission
    if (osMessageQueueGetCount(eems.hostbFrameQ) == 0) {
        return; // no packets
    }

    // at this point, we can safely pop from the frame queue, popping will not block,
    // and the forward queue is certainly empty

    // get EEM frame pointer
    USB_EemFrame *eemFrame;
    osMessageQueueGet(eems.hostbFrameQ, &eemFrame, 0, osWaitForever);

    // push frame into the transfer queue
    /* uint32_t frameLength = EEM_HEADER_GET_LENGTH(eemFrame->header);
    bfifo_push(bf, (const uint8_t *) eemFrame, frameLength + EEM_HEADER_SIZE); // push header and payload */

    // setup frame length and reset read index
    eems.hostbFrameLength = EEM_HEADER_GET_LENGTH(eemFrame->header) + EEM_HEADER_SIZE;
    eems.hostbFrameIndex = 0;

    // set the frame pointer
    eems.hostbFrame = eemFrame;

    // MSG("%u\n", frameLength);
}

void usb_eem_ethernet_intercept_cb(EthInterface *intf, const RawPckt *rawPckt) {
    // check for free space in the hostbound frame queue
    if (osMessageQueueGetSpace(eems.hostbFrameQ) == 0) {
        MSG("Hostbound frame queue is full, frame dropped!\n");
        return;
    }

    // copy packet
    USB_EemFrame *eemFrame = dynmem_alloc(sizeof(USB_EemFrame) + rawPckt->size);
    if (eemFrame == NULL) {
        MSG("Allocation while intercepting traffic failed!\n");
        return;
    }

    // at this point, allocation was certainly succesful

    // build header: CRC is calculated and present, size comes from the raw packet size
    eemFrame->header = EEM_BUILD_HEADER(EEM_HEADER_TYPE_DEFAULT, 1, rawPckt->size);

    // copy payload
    memcpy(eemFrame->payload, rawPckt->payload, rawPckt->size);

    // push onto hostbound queue
    osMessageQueuePut(eems.hostbFrameQ, &eemFrame, 0, osWaitForever); // will not block, since we checked, that there is space in the queue
}

void usb_eem_set_intf(EthInterface *intf) {
    eems.intf = intf;                                                   // store interface
    ethinf_set_intercept_callback(intf, usb_eem_ethernet_intercept_cb); // set intercept callback
}

void usb_eem_thread(void *) {
    for (;;) {
        // pop event
        USB_EemEvent evt = usb_eem_pop_event();

        switch (evt) {
        // process netbound packet
        case EEM_EVT_NETBOUND_PCKT_RECEIVED:
            usb_eem_process_networkbound();
            break;
        case EEM_EVT_HOSTBOUND_TRANSFER_DONE:
            usb_eem_process_hostbound();
            break;
        default:
            break;
        }
    }
}