#include <avr/io.h>
#include <avr/wdt.h>

#include "bits.h"

#include <avr/io.h>
#include <avr/interrupt.h>
#include <avr/wdt.h>
#include <avr/pgmspace.h>

#include "bits.h"

// All of port B is address.
// All of port D is data.
#define PORTC_ICONT     2
#define PORTC_nDBE      4
#define PORTC_POLL      5
#define PORTC_nASR      6
#define PORTC_ADV       7

#define EP0_SIZE        64

void setup_delay()
{
    for (volatile uint16_t i = 0; i < 9; ++i)
    {
        asm volatile("nop");
    }
}

void write_reg(uint8_t addr, uint8_t data)
{
    // Take control of instrument bus and enable data bus drivers.
    // This also drives ADV low at the same time.
    PORTC = 0x00;
    DDRC = _BV(PORTC_ICONT) | _BV(PORTC_nDBE) | _BV(PORTC_ADV);
    
    // Set up address.  This will cause the data buffer to drive into the instrument.
    DDRB = 0xff;
    PORTB = addr & 0x7f;
    
    // Set up data.
    DDRD = 0xff;
    PORTD = data;
    
    setup_delay();
    
    // Assert ADV
    set(PORTC, PORTC_ADV);
    
    // 800ns
    asm volatile("nop");
    asm volatile("nop");
    asm volatile("nop");
    asm volatile("nop");
    asm volatile("nop");
    
    // Deassert ADV.
    clear(PORTC, PORTC_ADV);
    
    setup_delay();
    
    // Stop driving the data lines.
    // This will cause the address buffer to drive out of the instrument because AD7 will be pulled high.
    PORTD = 0xff;
    DDRD = 0x00;
    
    // Stop driving the address lines.
    PORTB = 0xff;
    DDRB = 0x00;
    
    // Return control to the instrument computer.
    DDRC = _BV(PORTC_ICONT) | _BV(PORTC_nDBE);
    PORTC = _BV(PORTC_ICONT) | _BV(PORTC_nDBE);
}

uint8_t read_reg(uint8_t addr)
{
    // Take control of instrument bus and enable data bus drivers.
    // This also drives ADV low at the same time.
    PORTC = 0x00;
    DDRC = _BV(PORTC_ICONT) | _BV(PORTC_nDBE) | _BV(PORTC_ADV);
    
    // Set up address.  This will cause the data buffer to drive out of the instrument.
    DDRB = 0xff;
    PORTB = addr | 0x80;
    
    setup_delay();
    
    // Assert ADV
    set(PORTC, PORTC_ADV);
    
    //FIXME - 15.40us
    setup_delay();
    
    uint8_t value = PIND;
    
    // Deassert ADV.
    clear(PORTC, PORTC_ADV);
    
    setup_delay();
    
    // Stop driving the address lines.
    PORTB = 0xff;
    DDRB = 0x00;
    
    // Return control to the instrument computer.
    DDRC = _BV(PORTC_ICONT) | _BV(PORTC_nDBE);
    PORTC = _BV(PORTC_ICONT) | _BV(PORTC_nDBE);
    
    return value;
}

typedef struct
{
    uint8_t bmRequestType;
    uint8_t bRequest;
    uint16_t wValue;
    uint16_t wIndex;
    uint16_t wLength;
} USB_Setup_Request;

// Used for putting words in descriptor byte arrays
#define WORD(x) ((x) & 0xff), (((x) >> 8) & 0xff)

const uint8_t device_desc[] PROGMEM =
{
    18,             // bLength
    1,              // bDescriptorType
    WORD(0x200),    // bcdUSB
    0xff,           // bDeviceClass
    0,              // bDeviceSubclass
    0,              // bDeviceProtocol
    EP0_SIZE,       // bMaxPacketSize0
    WORD(0x3141),   // idVendor
    WORD(0x0005),   // idProduct
    WORD(0),        // bcdDevice
    0,              // iManufacturer
    0,              // iProduct
    0,              // iSerialNumber
    1               // bNumConfigurations
};

#define CONFIG_SIZE     9
#define INTF_SIZE       9
#define EP_SIZE         7
const uint8_t config_desc[] PROGMEM =
{
    // Configuration
    CONFIG_SIZE,    // bLength
    2,              // bDescriptorType
    WORD(           // wTotalLength
        CONFIG_SIZE +
        INTF_SIZE),
    1,              // bNumInterfaces
    1,              // bConfigurationValue
    0,              // iConfiguration
    0xc0,           // bmAttributes
    0,              // MaxPower
    
    // Interface 0
    INTF_SIZE,      // bLength
    4,              // bDescriptorType
    0,              // bInterfaceNumber
    0,              // bAlternateSetting
    0,              // bNumEndpoints (not counting endpoint 0)
    0xff,           // bInterfaceClass
    0xff,           // bInterfaceSubclass
    0,              // bInterfaceProtocol
    0               // iInterface
};

// Read n bytes from the current endpoint
void usb_read(void *buf, uint8_t n)
{
    while (n--)
    {
        *(uint8_t *)buf++ = UEDATX;
    }
}

// Sends the data currently in the FIFO (typically used to send a zero-length packet)
void usb_write_packet()
{
    clear(UEINTX, TXINI);
    loop_until_bit_is_set(UEINTX, TXINI);
}

// Send data from program memory
void usb_write_P(PGM_VOID_P buf, uint8_t len, uint16_t requested)
{
    // Only send up to the requested amount of data.
    // Otherwise we will be waiting forever for another IN token from the host.
    uint8_t n = len;
    if (n > requested)
    {
        n = requested;
    }
    
    uint8_t sent = 0;
    while (n--)
    {
        UEDATX = pgm_read_byte(buf++);
        ++sent;
        
        if (n == 0 || sent == 64)       //FIXME - Endpoint size
        {
            sent = 0;
            clear(UEINTX, TXINI);
            while (!(UEINTX & ((1 << TXINI) | (1 << RXOUTI))))
                ;
            if (bit_is_set(UEINTX, RXOUTI))
            {
                // Aborted by host
                return;
            }
        }
    }
    
    // Wait for status
    loop_until_bit_is_set(UEINTX, RXOUTI);
    clear(UEINTX, RXOUTI);
}

// Send data from RAM
void usb_write(void *buf, uint8_t len, uint16_t requested)
{
    // Only send up to the requested amount of data.
    // Otherwise we will be waiting forever for another IN token from the host.
    uint8_t n = len;
    if (n > requested)
    {
        n = requested;
    }
    
    uint8_t sent = 0;
    while (n--)
    {
        UEDATX = *(uint8_t *)buf++;
        ++sent;
        
        if (n == 0 || sent == 64)       //FIXME - Endpoint size
        {
            sent = 0;
            clear(UEINTX, TXINI);
            while (!(UEINTX & ((1 << TXINI) | (1 << RXOUTI))))
                ;
            if (bit_is_set(UEINTX, RXOUTI))
            {
                // Aborted by host
                return;
            }
        }
    }
    
    // Wait for status
    loop_until_bit_is_set(UEINTX, RXOUTI);
    clear(UEINTX, RXOUTI);
}

// Starts the bootloader
void bootloader()
{
    cli();
    set(UDCON, DETACH);
    ((void (*)(void))0x1800)();
}

// Disables and frees memory for all endpoints except 0.
void usb_free_endpoints()
{
    uint8_t old_ep = UENUM;
    
    uint8_t i;
    for (i = 1; i <= 4; ++i)
    {
        UENUM = i;
        clear(UECONX, EPEN);
        clear(UECFG1X, ALLOC);
    }
    
    UENUM = old_ep;
}

// Called when a SETUP packet is received on endpoint 0.
void usb_handle_setup()
{
    // Read the packet
    USB_Setup_Request req;
    usb_read(&req, 8);
    
    // Clear the FIFO
    //FIXME - Do this later if this is an OUT transfer (host will be sending in the DATA stage)?
    clear(UEINTX, RXSTPI);
    
    // Every request must end with either an IN packet (usb_write_packet or similar) or
    // a STALL, which is sent if valid == 0.
    uint8_t valid = 1;
    
    if ((req.bmRequestType & 0x60) == 0)
    {
        // Device request
        switch (req.bRequest)
        {
            case 0x05:  // Set address
                UDADDR = (UDADDR & (1 << ADDEN)) | (req.wValue & 0x7f);
                usb_write_packet();
                set(UDADDR, ADDEN);
                break;
                
            case 0x06:  // Get descriptor
                if (req.wValue == 0x0100)
                {
                    // Get device descriptor
                    usb_write_P(device_desc, sizeof(device_desc), req.wLength);
                } else if (req.wValue == 0x0200)
                {
                    // Get configuration descriptor
                    usb_write_P(config_desc, sizeof(config_desc), req.wLength);
                } else {
                    valid = 0;
                }
                break;
            
            case 0x09:  // Set configuration
                usb_write_packet();
                break;
            
            default:
                valid = 0;
                break;
        }
    } else if ((req.bmRequestType & 0x7f) == 0x40)
    {
        // Vendor request
        switch (req.bRequest)
        {
            case 0:
                write_reg(req.wIndex, req.wValue);
                usb_write_packet();
                break;
            
            case 1:
            {
                // This is a copy of usb_write which doesn't require an intermediate buffer.
                uint8_t sent = 0;
                uint16_t n = req.wLength;
                while (n--)
                {
                    UEDATX = read_reg(req.wIndex);
                    ++sent;
                    
                    if (n == 0 || sent == EP0_SIZE)
                    {
                        sent = 0;
                        clear(UEINTX, TXINI);
                        while (!(UEINTX & ((1 << TXINI) | (1 << RXOUTI))))
                            ;
                        if (bit_is_set(UEINTX, RXOUTI))
                        {
                            // Aborted by host
                            return;
                        }
                    }
                }
                
                // Wait for status
                loop_until_bit_is_set(UEINTX, RXOUTI);
                clear(UEINTX, RXOUTI);
                break;
            }
            
            case 0xff:  // Bootloader
                usb_write_packet();
                bootloader();
                break;
            
            default:
                valid = 0;
                break;
        }
    } else {
        valid = 0;
    }
    
    if (!valid)
    {
        // Unsupported request: stall EP0
        set(UECONX, STALLRQ);
    }
}

ISR(USB_GEN_vect)
{
    if (bit_is_set(UDINT, WAKEUPI))
    {
        // Wakeup (USB bus activity is detected)
        
        // We have to turn off the interrupt or it will trigger frequently when
        // the device is active.  This causes all sorts of weird problems.
        clear(UDIEN, WAKEUPE);
        
        // Start the 48MHz PLL
        set(PLLCSR, PLLE);
        loop_until_bit_is_set(PLLCSR, PLOCK);
        
        // Enable the USB clock
        clear(USBCON, FRZCLK);
        
        // Clear the interrupt flag after starting the clock
        clear(UDINT, WAKEUPI);
        
        // The USB controller will indicate end-of-reset soon.
        // The end-of-reset handler will set up endpoints.
    }
    
    if (bit_is_set(UDINT, EORSTI))
    {
        // USB reset
        clear(UDINT, EORSTI);
        
        usb_free_endpoints();
        
        // Endpoint 0: control
        //   64-byte packets (so we can send the device and config descriptors in one packet)
        //   Single buffer
        UENUM = 0;
        clear(UECONX, EPEN);
        clear(UECFG1X, ALLOC);
        UECONX = (1 << EPEN);       // Do this before other configuration
        UEIENX = (1 << RXSTPE);
        UECFG0X = 0;
        UECFG1X = (3 << EPSIZE0);
        set(UECFG1X, ALLOC);
    }
}

ISR(USB_COM_vect)
{
    UENUM = 0;
    if (bit_is_set(UEINTX, RXSTPI))
    {
        usb_handle_setup();
    }
}

int main()
{
    // Disable the watchdog timer.
    // This must be done in case we are started from the clear_bit,
    // which uses the WDT to reset.
    wdt_reset();
    clear(MCUSR, WDRF);
    WDTCSR |= (1 << WDCE) | (1 << WDE);
    WDTCSR = 0;
    
    // Run at full clock speed
    CLKPR = 1 << CLKPCE;
    CLKPR = 0;
    
    // Set up I/O ports
    DDRB = 0x00;
    PORTB = 0xff;
    
    DDRC = _BV(PORTC_ICONT) | _BV(PORTC_nDBE);
    PORTC = _BV(PORTC_ICONT) | _BV(PORTC_nDBE);
    
    DDRD = 0x00;
    PORTD = 0xff;
    
    // Reset the USB controller.
    // This will disable all interrupts and endpoints too.
    USBCON = (1 << FRZCLK);
    UDCON = (1 << DETACH);
    
    // Enable the USB controller
    set(USBCON, USBE);
    
    // Attach device
    clear(UDCON, DETACH);
    
    // Enable interrupts
    UDIEN = (1 << WAKEUPE) | (1 << EORSTE);
    
    sei();
    
    while (1)
    {
    }
}