#include "usb.h"
#include "usb_descriptors.h"
#include "usb_util.h"
#include "usb_cdc.h"
#include "usb_vendor.h"
#include "led.h"

#include "ownership.h"
MODULE_OWNS_PERIPHERAL(USB);
MODULE_OWNS_PIN(GPIOA, PIN_USB_DM);
MODULE_OWNS_PIN(GPIOA, PIN_USB_DP);
MODULE_OWNS_PIN(GPIOA, PIN_USB_PULLUP);

USB_PacketHandler_t USB_EP0OutPacketHandler;

uint8_t USB_DeviceStatus[2] = {0x00, 0x01};
volatile unsigned int USB_ResetCount = 0;
volatile unsigned int USB_Address = 0;

void USB_Init(void)
{
    // GPIOA clock
    RCC->APB2ENR |= RCC_APB2ENR_IOPAEN;
    RCC->APB2ENR |= RCC_APB2ENR_AFIOEN;
    RCC->APB1ENR |= RCC_APB1ENR_USBEN;

    GPIOA->BRR = (1 << PIN_USB_DM) | (1 << PIN_USB_DP) | (1 << PIN_USB_PULLUP);
    GPIOA->CRH = (GPIOA->CRH
        & ~(0xf << (4 * PIN_USB_DM - 32))
        & ~(0xf << (4 * PIN_USB_DP - 32))
        & ~(0xf << (4 * PIN_USB_PULLUP - 32)))
        | (0x3 << (4 * PIN_USB_DM - 32))        // PP output, 50 MHz
        | (0x3 << (4 * PIN_USB_DP - 32))        // PP output, 50 MHz
        | (0x3 << (4 * PIN_USB_PULLUP - 32))    // PP output, 2 MHz
        ;

    // Enable D- pull-up
    GPIOA->BSRR = (1 << PIN_USB_PULLUP);

    // GPIOA->CRH &= ~(GPIO_CRH_CNF11 | GPIO_CRH_MODE11
    //     | GPIO_CRH_CNF12 | GPIO_CRH_MODE12);
    // GPIOA->CRH |= GPIO_CRH_MODE11 | GPIO_CRH_MODE12;
    // GPIOA->ODR &= ~(GPIO_CRH_MODE11 | GPIO_CRH_MODE12);
    USB_Delay(100000);

    // Enable reset and correct transfer interrupts
    NVIC_EnableIRQ(USB_LP_IRQn);
    NVIC_SetPriority(USB_LP_IRQn, 7);

    // Analog power up
    USB->CNTR = (uint16_t)USB_CNTR_FRES;
    // Minimum delay: 1 µs
    USB_Delay(3000);

    USB->CNTR = (uint16_t)0;
    USB->ISTR = (uint16_t)0;
    USB->CNTR = (uint16_t)(USB_CNTR_RESETM | USB_CNTR_CTRM);

    // Configure USB pins (PA11 and PA12 in AF mode, 50 MHz push-pull)
    GPIOA->CRH |= GPIO_CRH_CNF11_1 | GPIO_CRH_MODE11 |\
        GPIO_CRH_CNF12_1 | GPIO_CRH_MODE12;

}

void USB_Disable(void)
{
    // Disable interrupts, force reset
    USB->CNTR = USB_CNTR_FRES;

    // Clear interrupt flags
    USB->ISTR = 0x00;

    // Power down
    USB->CNTR = USB_CNTR_FRES | USB_CNTR_PDWN;

    // Disable D- pull-up
    GPIOA->BRR = (1 << PIN_USB_PULLUP);
}

static inline void USB_HandleReset(void)
{
    // Remove reset flag
    USB->ISTR = (uint16_t)~(USB_ISTR_RESET);
    USB_ResetCount++;

    // Set buffer table origin
    USB->BTABLE = USB_BTABLE_OFFSET;

    // Set buffer table contents for control endpoint 0
    USB_BTABLE_ENTRIES[0].COUNT_RX = USB_EP_RXCOUNT_BL_SIZE | (1 << 10);
    USB_BTABLE_ENTRIES[0].ADDR_RX = 0x40;
    USB_BTABLE_ENTRIES[0].COUNT_TX = 0;
    USB_BTABLE_ENTRIES[0].ADDR_TX = 0x80;

    // Configure endpoint 0
    USB_SetEPR(&(USB->EP0R), USB_EPR_EP_TYPE_CONTROL
        | USB_EPR_STAT_TX_NAK | USB_EPR_STAT_RX_VALID);

    USBCDC_InitEndpoints();

    // Enable
    USB->DADDR = USB_DADDR_EF;

    LED_Off(LED_Orange);
}

static inline void USB_HandleIn(void)
{
    if((USB->DADDR & USB_DADDR_ADD) != USB_Address)
    {
        USB->DADDR = USB_Address | USB_DADDR_EF;
    }

    // Ready for next packet
    USB_SetEPRXStatus(&USB->EP0R, USB_EP_RX_VALID);
}

static inline void USB_Handle0Out(void)
{
    int length = USB_BTABLE_ENTRIES[0].COUNT_RX & 0x3ff;
    if(length != 0)
    {
        if(USB_EP0OutPacketHandler != NULL)
        {
            USB_EP0OutPacketHandler(length);
        }
    }

    // Ready for next packet
    USB_SetEPRXStatus(&USB->EP0R, USB_EP_RX_VALID);
}

static inline void USB_HandleSetup(void)
{
    USB_SetupPacket_t sp;
    USB_PMAToMemory((uint8_t*)&sp, USB_BTABLE_ENTRIES[0].ADDR_RX,
        sizeof(USB_SetupPacket_t));

    // Delete this: Sanity check for abnormal setup package lengths
    unsigned int packet_length = USB_BTABLE_ENTRIES[0].COUNT_RX & 0x3ff;
    if(packet_length > sizeof(USB_SetupPacket_t))
    {
        __asm__ volatile("bkpt");
    }

    const void *reply_data = NULL;
    int reply_length = 0;
    uint8_t reply_response = USB_TOKEN_ACK;

    if((sp.bmRequestType & (USB_REQUEST_TYPE | USB_REQUEST_RECIPIENT))
        == (USB_REQUEST_TYPE_STANDARD | USB_REQUEST_RECIPIENT_DEVICE))
    {
        switch(sp.bRequest)
        {
            case USB_REQUEST_GET_STATUS:;
                if(sp.wValue == 0 && sp.wIndex == 0 && sp.wLength == 2)
                {
                    reply_length = 2;
                    reply_data = USB_DeviceStatus;
                }
                break;

            case USB_REQUEST_GET_DESCRIPTOR:;
                USB_DescriptorType_t descriptor_type = sp.wValue >> 8;
                int descriptor_index = sp.wValue & 0xff;
                reply_length = sp.wLength;
                USB_HandleGetDescriptor(descriptor_type, descriptor_index,
                    &reply_data, &reply_length, &reply_response);
                break;

            case USB_REQUEST_SET_ADDRESS:
                USB_Address = sp.wValue & USB_DADDR_ADD;
                reply_response = USB_EP_TX_VALID;
                break;

            case USB_REQUEST_SET_CONFIGURATION:
                reply_response = USB_EP_TX_VALID;
                break;

            default:;
                __asm__ volatile("nop");
                reply_response = USB_EP_TX_STALL;
                break;
        }
    }
    else if((sp.bmRequestType & (USB_REQUEST_TYPE | USB_REQUEST_RECIPIENT))
        == (USB_REQUEST_TYPE_CLASS | USB_REQUEST_RECIPIENT_INTERFACE))
    {
        USBCDC_HandleInterfaceSetup(&sp, &reply_data, &reply_length,
            &reply_response);
    }
    else if((sp.bmRequestType & (USB_REQUEST_TYPE | USB_REQUEST_RECIPIENT))
        == (USB_REQUEST_TYPE_VENDOR | USB_REQUEST_RECIPIENT_DEVICE))
    {
        USBVendor_HandleDeviceSetup(&sp, &reply_data, &reply_length,
            &reply_response);
    }
    else
    {
        // Unknown request
        __asm__ volatile("bkpt");
    }

    if(reply_data)
    {
        // Reply with data
        USB_MemoryToPMA(USB_BTABLE_ENTRIES[0].ADDR_TX, reply_data,
            reply_length);
        USB_BTABLE_ENTRIES[0].COUNT_TX = reply_length;
        USB_SetEPTXStatus(&(USB->EP0R), USB_EP_TX_VALID);
    }
    else
    {
        // Send response
        USB_BTABLE_ENTRIES[0].COUNT_TX = 0;
        USB_SetEPTXStatus(&(USB->EP0R), reply_response);
    }

    // Ready for the next packet
    USB_SetEPRXStatus(&(USB->EP0R), USB_EP_RX_VALID);
}

// Interrupt service routine
void USB_LP_IRQHandler(void)
{
    if(USB->ISTR & USB_ISTR_RESET)
    {
        // Reset happened
        USB_HandleReset();
        return;
    }
    uint16_t istr;
    while((istr = USB->ISTR) & (USB_ISTR_CTR))
    {
        if(istr & USB_ISTR_CTR)
        {
            // Correct transfer
            int ep = istr & USB_ISTR_EP_ID;
            switch(ep)
            {
                case 0:
                    // Determine transfer direction
                    if(istr & USB_ISTR_DIR)
                    {
                        // Out transfer
                        if(USB->EP0R & USB_EP0R_SETUP)
                        {
                            // Clear CTR_RX
                            USB->EP0R = USB->EP0R & (USB_EP0R_SETUP
                                | USB_EP0R_EP_TYPE | USB_EP0R_EP_KIND
                                | USB_EP0R_CTR_TX | USB_EP0R_EA);

                            // Setup packed received
                            USB_HandleSetup();
                        }
                        else
                        {
                            // Clear CTR_RX
                            USB->EP0R = USB->EP0R & (USB_EP0R_SETUP
                                | USB_EP0R_EP_TYPE | USB_EP0R_EP_KIND
                                | USB_EP0R_CTR_TX | USB_EP0R_EA);

                            USB_Handle0Out();
                        }
                    }
                    else
                    {
                        // In transfer

                        // Clear CTR_TX
                        USB->EP0R = USB->EP0R & (USB_EP0R_CTR_RX
                            | USB_EP0R_SETUP | USB_EP0R_EP_TYPE
                            | USB_EP0R_EP_KIND | USB_EP0R_EA);

                        // In transfer
                        USB_HandleIn();
                    }
                    break;

                case 2:
                    // Data in transfer finished.
                    USB_BTABLE_ENTRIES[2].COUNT_TX = 0;
                    USB_ClearEPCTTX(&USB->EP2R);
                    // Start the next in transfer, if there is data left in the
                    // buffer
                    USBCDC_HandleDataInComplete();
                    break;

                case 3:
                    // A successful out transfer on endpoint 3 means new data
                    // arrived on the virtual serial port
                    USBCDC_HandleDataOut();
                    USB_ClearEPCTRX(&USB->EP3R);
                    break;

                default:
                    // Other endpoints are not implemented (most importantly the
                    // notification interrupt endpoint 1)
                    __asm__ volatile("bkpt");
                    break;
            }
        }
    }
}

