#include <stdlib.h>
#include "core.h"
#include "CH58x_common.h"

static volatile uint8_t *ep_t_len_regs[] = {
    &R8_UEP0_T_LEN,
    &R8_UEP1_T_LEN,
    &R8_UEP2_T_LEN,
    &R8_UEP3_T_LEN,
    &R8_UEP4_T_LEN,
    &R8_UEP5_T_LEN,
    &R8_UEP6_T_LEN,
    &R8_UEP7_T_LEN,
};

static volatile uint8_t *ep_ctrl_regs[] = {
    &R8_UEP0_CTRL,
    &R8_UEP1_CTRL,
    &R8_UEP2_CTRL,
    &R8_UEP3_CTRL,
    &R8_UEP4_CTRL,
    &R8_UEP5_CTRL,
    &R8_UEP6_CTRL,
    &R8_UEP7_CTRL,
};

static uint8_t recvlens[8];
uint8_t bConfigurationValue = 0;

// The hardware requires EP0 and EP4 to share a contiguous DMA buffer.
// Fuck if I know why.
__attribute__((aligned(4))) volatile uint8_t ep04_buf[64 + 64 + 64];

__attribute__((aligned(4))) volatile uint8_t epbuf[(8 - 2) * 2 * 64];

volatile uint8_t *buf_for_ep(uint8_t ep, int dir_in) {
    if (ep == 0) return ep04_buf;
    if (ep == 4) return ep04_buf + (dir_in ? 128 : 64);

    // EP0 and EP4 aren't in here. so skip them.
    ep -= 1;
    if (ep > 3) ep -= 1;
    return epbuf + (ep * 128 + (dir_in ? 64 : 0));
}

// Configure the response to set for the next token for the endpoint.
// The interrupt handler will reconfigure this to return NAK whenever
// the next token has been received.
enum usb_endpoint_state set_endpoint_state(uint8_t endpoint, enum usb_endpoint_state state) {
    int dir_in = endpoint & 0x80; // 0 -> OUT endpoint, 1 -> IN endpoint
    int ep_num = endpoint & 0x7F;

    if (dir_in) {
        uint8_t prev_state = *ep_ctrl_regs[ep_num] & MASK_UEP_T_RES;
        *ep_ctrl_regs[ep_num] = (*ep_ctrl_regs[ep_num] & ~MASK_UEP_T_RES) | (state & 0x03);
        return prev_state;
    } else {
        uint8_t prev_state = (*ep_ctrl_regs[ep_num] & MASK_UEP_R_RES) >> 2;
        *ep_ctrl_regs[ep_num] = (*ep_ctrl_regs[ep_num] & ~MASK_UEP_R_RES) | ((state & 0x03) << 2);
        return prev_state;
    }
}


// Toggle the DATA0/DATA1 toggle bit, indicating which of
// DATA0 and DATA1 to use fo rthe next receiving or sending of data.
// These bits are only toggled once data has been successfully transferred;
// a la USB_EP_STATE_ACK.
// 
// The SETUP control transfers on endoint 0 _always_ use DATA0.
static void flip_endpoint_data_toggle(uint8_t endpoint) {
    int dir_in = endpoint & 0x80;
    int ep_num = endpoint & 0x7F;

    if (dir_in) {
        *ep_ctrl_regs[ep_num] ^= RB_UEP_T_TOG;
    } else {
        *ep_ctrl_regs[ep_num] ^= RB_UEP_R_TOG;
    }
}

static void handle_bus_reset() {
    R8_USB_INT_FG = RB_UIF_BUS_RST;
    R8_USB_DEV_AD = 0;
    bConfigurationValue = 0;
}

static void handle_suspend() {
    R8_USB_INT_FG = RB_UIF_SUSPEND;

    // TODO
}

static void handle_sof() {
    R8_USB_INT_FG = RB_UIF_HST_SOF;
}


__HIGH_CODE
static void handle_transfer(uint8_t ep_num, int is_in) {
    uint8_t endpoint = ep_num | (is_in ? 0x80 : 0);

    // We may have received a transfer.
    // If we expected this transfer, toggle the DATA0/DATA1 bits.
    enum usb_endpoint_state prev_state = set_endpoint_state(endpoint, USB_EP_STATE_NAK);
    if (prev_state == USB_EP_STATE_ACK) {
        flip_endpoint_data_toggle(endpoint);
    }

    if (ep_num == 0) {
        handle_ctrl_transfer(is_in);
        return;
    }

    if (!is_in) {
        recvlens[endpoint] = R8_USB_RX_LEN;
    }
}

void usb_init() {

	R8_USB_CTRL = 0x00;

	R8_UEP4_1_MOD = RB_UEP4_RX_EN | RB_UEP4_TX_EN |
				RB_UEP1_RX_EN | RB_UEP1_TX_EN;
	R8_UEP2_3_MOD = RB_UEP2_RX_EN | RB_UEP2_TX_EN |
				RB_UEP3_RX_EN | RB_UEP3_TX_EN;
	R8_UEP567_MOD = RB_UEP7_RX_EN | RB_UEP7_TX_EN |
				RB_UEP6_RX_EN | RB_UEP6_TX_EN |
				RB_UEP5_RX_EN | RB_UEP5_TX_EN;

	R16_PIN_ANALOG_IE |= RB_PIN_USB_IE | RB_PIN_USB_DP_PU;
	R8_UDEV_CTRL = RB_UD_PD_DIS | RB_UD_PORT_EN;

	R8_USB_DEV_AD = 0x00;
	R8_USB_INT_FG = 0xFF;

	R16_UEP0_DMA = (uint16_t) (uintptr_t) ep04_buf;
	R16_UEP1_DMA = (uint16_t) (uintptr_t) (epbuf + 0 * 128);
	R16_UEP2_DMA = (uint16_t) (uintptr_t) (epbuf + 1 * 128);
	R16_UEP3_DMA = (uint16_t) (uintptr_t) (epbuf + 2 * 128);
	R16_UEP5_DMA = (uint16_t) (uintptr_t) (epbuf + 3 * 128);
	R16_UEP6_DMA = (uint16_t) (uintptr_t) (epbuf + 4 * 128);
	R16_UEP7_DMA = (uint16_t) (uintptr_t) (epbuf + 5 * 128);

	for (int i = 0; i < 8; i++) {
		set_endpoint_state(i, USB_EP_STATE_NAK);
		set_endpoint_state(i | 0x80, USB_EP_STATE_NAK);
	}

	R8_USB_CTRL = RB_UC_DEV_PU_EN | RB_UC_INT_BUSY | RB_UC_DMA_EN;
	R8_USB_INT_EN = RB_UIE_SUSPEND | RB_UIE_BUS_RST | RB_UIE_TRANSFER;

	PFIC_EnableIRQ(USB_IRQn);
}

__INTERRUPT
__HIGH_CODE
void USB_IRQHandler(void) {
    uint8_t interrupt = R8_USB_INT_FG;

    if (interrupt & RB_UIF_BUS_RST) {
        handle_bus_reset();
    }

    if (interrupt & RB_UIF_SUSPEND) {
        handle_suspend();
    }

    if (interrupt & RB_UIF_TRANSFER) {
        uint8_t status = R8_USB_INT_ST;
        uint8_t ep_num = status & MASK_UIS_ENDP;
        uint8_t token = ((status & MASK_UIS_TOKEN) >> 4) & 0x03;

        if (token == 0x00) { // OUT
            if (status & RB_UIS_TOG_OK) handle_transfer(ep_num, 0);
            R8_USB_INT_FG = RB_UIF_TRANSFER;
        } else if (token == 0x01) { // SOF
            handle_sof();
            R8_USB_INT_FG = RB_UIF_TRANSFER;
        } else if (token == 0x02) { // IN
            handle_transfer(ep_num, 1);
            R8_USB_INT_FG = RB_UIF_TRANSFER;
        }
        if ((status & RB_UIS_SETUP_ACT)) {
            handle_setup_request();
            R8_USB_INT_FG = RB_UIF_TRANSFER;
        }
    }
}

// Receive up to buflen bytes from selected OUT endpoint.
// If data is not yet available, returns -1.
int16_t usb_recv(uint8_t endpoint, void *buffer, size_t buflen) {
    if (recvlens[endpoint] == 0xFF) {
        set_endpoint_state(endpoint, USB_EP_STATE_ACK);
        return -1;
    }

    if (recvlens[endpoint] < buflen) {
        buflen = recvlens[endpoint];
    }

    recvlens[endpoint] = 0xFF;

    memcpy(buffer, (uint8_t *)buf_for_ep(endpoint, 0), buflen);
    return buflen;
}

// Sends up to buflen bytes to selected IN endpoint.
// If the buffer is still busy, returns -1.
int16_t usb_xmit(uint8_t endpoint, void *buffer, size_t buflen) {
    if ((*ep_ctrl_regs[endpoint] & MASK_UEP_T_RES) != USB_EP_STATE_NAK) {
        return -1;
    }

    if (buflen > 64) buflen = 64;

    memcpy((uint8_t *)buf_for_ep(endpoint, 1), buffer, buflen);
    *ep_t_len_regs[endpoint] = buflen;
    set_endpoint_state(endpoint | 0x80, USB_EP_STATE_ACK);

    return buflen;
}

int usb_can_xmit(uint8_t endpoint) {
    return (*ep_ctrl_regs[endpoint] & MASK_UEP_T_RES) == USB_EP_STATE_NAK;
}

int usb_connected(void) {
    return !!(R8_USB_MIS_ST & RB_UMS_SUSPEND);
}