/* ----------------------------------------------------------------------------
 * ergoDOX : controller specific code
 * ----------------------------------------------------------------------------
 * Copyright (c) 2012 Ben Blazak <benblazak.dev@gmail.com>
 * Released under The MIT License (MIT) (see "license.md")
 * Project located at <https://github.com/benblazak/ergodox-firmware>
 * ------------------------------------------------------------------------- */

#include "./controller.h"

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

/* returns
 * - success: 0
 * - error: number of the function that failed
 */
uint8_t kb_init(void) {
	if (teensy_init())    // must be first
		return 1;
	if (mcp23018_init())  // must be second
		return 2;

	return 0;  // success
}

/* returns
 * - success: 0
 * - error: number of the function that failed
 */
uint8_t kb_update_matrix(bool matrix[KB_ROWS][KB_COLUMNS]) {
	if (teensy_update_matrix(matrix))
		return 1;
	if (mcp23018_update_matrix(matrix))
		return 2;

	return 0;  // success
}

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

// check options
#if  (MCP23018__DRIVE_ROWS && MCP23018__DRIVE_COLUMNS)	\
 || !(MCP23018__DRIVE_ROWS || MCP23018__DRIVE_COLUMNS)
	#error "See 'Pin drive direction' in 'options.h'"
#endif

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

// register addresses (see "mcp23018.md")
#define IODIRA 0x00  // i/o direction register
#define IODIRB 0x01
#define GPPUA  0x0C  // GPIO pull-up resistor register
#define GPPUB  0x0D
#define GPIOA  0x12  // general purpose i/o port register (write modifies OLAT)
#define GPIOB  0x13
#define OLATA  0x14  // output latch register
#define OLATB  0x15

// TWI aliases
#define TWI_ADDR_WRITE ( (MCP23018_TWI_ADDRESS<<1) | TW_WRITE )
#define TWI_ADDR_READ  ( (MCP23018_TWI_ADDRESS<<1) | TW_READ  )

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

/* returns:
 * - success: 0
 * - failure: twi status code
 *
 * notes:
 * - `twi_stop()` must be called *exactly once* for each twi block, the way
 *   things are currently set up.  this may change in the future.
 */
uint8_t mcp23018_init(void) {
	uint8_t ret;

	// set pin direction
	// - unused  : input  : 1
	// - input   : input  : 1
	// - driving : output : 0
	twi_start();
	ret = twi_send(TWI_ADDR_WRITE);
	if (ret) goto out;  // make sure we got an ACK
	twi_send(IODIRA);
	#if MCP23018__DRIVE_ROWS
		twi_send(0b11111111);  // IODIRA
		twi_send(0b11000000);  // IODIRB
	#elif MCP23018__DRIVE_COLUMNS
		twi_send(0b10000000);  // IODIRA
		twi_send(0b11111111);  // IODIRB
	#endif
	twi_stop();

	// set pull-up
	// - unused  : on  : 1
	// - input   : on  : 1
	// - driving : off : 0
	twi_start();
	ret = twi_send(TWI_ADDR_WRITE);
	if (ret) goto out;  // make sure we got an ACK
	twi_send(GPPUA);
	#if MCP23018__DRIVE_ROWS
		twi_send(0b11111111);  // GPPUA
		twi_send(0b11000000);  // GPPUB
	#elif MCP23018__DRIVE_COLUMNS
		twi_send(0b10000000);  // GPPUA
		twi_send(0b11111111);  // GPPUB
	#endif
	twi_stop();

	// set logical value (doesn't matter on inputs)
	// - unused  : hi-Z : 1
	// - input   : hi-Z : 1
	// - driving : hi-Z : 1
	twi_start();
	ret = twi_send(TWI_ADDR_WRITE);
	if (ret) goto out;  // make sure we got an ACK
	twi_send(OLATA);
	twi_send(0b11111111);  //OLATA
	twi_send(0b11111111);  //OLATB

out:
	twi_stop();
	return ret;
}

/* returns:
 * - success: 0
 * - failure: twi status code
 */
#if KB_ROWS != 6 || KB_COLUMNS != 14
	#error "Expecting different keyboard dimensions"
#endif

uint8_t mcp23018_update_matrix(bool matrix[KB_ROWS][KB_COLUMNS]) {
	uint8_t ret, data;

	// initialize things, just to make sure
	// - it's not appreciably faster to skip this, and it takes care of the
	//   case when the i/o expander isn't plugged in during the first
	//   init()
	ret = mcp23018_init();

	// if there was an error
	if (ret) {
		// clear our part of the matrix
		for (uint8_t row=0; row<=5; row++)
			for (uint8_t col=0; col<=6; col++)
				matrix[row][col] = 0;

		return ret;
	}


	// --------------------------------------------------------------------
	// update our part of the matrix

	#if MCP23018__DRIVE_ROWS
		for (uint8_t row=0; row<=5; row++) {
			// set active row low  : 0
			// set other rows hi-Z : 1
			twi_start();
			twi_send(TWI_ADDR_WRITE);
			twi_send(GPIOB);
			twi_send( 0xFF & ~(1<<(5-row)) );
			twi_stop();

			// read column data
			twi_start();
			twi_send(TWI_ADDR_WRITE);
			twi_send(GPIOA);
			twi_start();
			twi_send(TWI_ADDR_READ);
			twi_read(&data);
			twi_stop();

			// update matrix
			for (uint8_t col=0; col<=6; col++) {
				matrix[row][col] = !( data & (1<<col) );
			}
		}

		// set all rows hi-Z : 1
		twi_start();
		twi_send(TWI_ADDR_WRITE);
		twi_send(GPIOB);
		twi_send(0xFF);
		twi_stop();

	#elif MCP23018__DRIVE_COLUMNS
		for (uint8_t col=0; col<=6; col++) {
			// set active column low  : 0
			// set other columns hi-Z : 1
			twi_start();
			twi_send(TWI_ADDR_WRITE);
			twi_send(GPIOA);
			twi_send( 0xFF & ~(1<<col) );
			twi_stop();

			// read row data
			twi_start();
			twi_send(TWI_ADDR_WRITE);
			twi_send(GPIOB);
			twi_start();
			twi_send(TWI_ADDR_READ);
			twi_read(&data);
			twi_stop();

			// update matrix
			for (uint8_t row=0; row<=5; row++) {
				matrix[row][col] = !( data & (1<<(5-row)) );
			}
		}

		// set all columns hi-Z : 1
		twi_start();
		twi_send(TWI_ADDR_WRITE);
		twi_send(GPIOA);
		twi_send(0xFF);
		twi_stop();

	#endif

	// /update our part of the matrix
	// --------------------------------------------------------------------

	return ret;  // success
}

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

// check options
#if  (TEENSY__DRIVE_ROWS && TEENSY__DRIVE_COLUMNS)	\
 || !(TEENSY__DRIVE_ROWS || TEENSY__DRIVE_COLUMNS)
	#error "See 'Pin drive direction' in 'options.h'"
#endif
// ----------------------------------------------------------------------------

// processor frequency (from <http://www.pjrc.com/teensy/prescaler.html>)
#define  CPU_PRESCALE(n)  (CLKPR = 0x80, CLKPR = (n))
#define  CPU_16MHz        0x00
#define  CPU_8MHz         0x01
#define  CPU_4MHz         0x02
#define  CPU_2MHz         0x03
#define  CPU_1MHz         0x04
#define  CPU_500kHz       0x05
#define  CPU_250kHz       0x06
#define  CPU_125kHz       0x07
#define  CPU_62kHz        0x08


/*
 * pin macros
 * - note: you can move the `UNUSED`, `ROW`, and `COLUMN` pins around, but be
 *   sure to keep the set of all the pins listed constant.  other pins are not
 *   movable, and either are referenced explicitly or have macros defined for
 *   them elsewhere.
 * - note: if you change pin assignments, please be sure to update
 *   "teensy-2-0.md", and the '.svg' circuit diagram.
 */

// --- unused
#define  UNUSED_0  C, 7
#define  UNUSED_1  D, 7
#define  UNUSED_2  D, 4  // hard to use with breadboard (on the end)
#define  UNUSED_3  D, 5  // hard to use with breadboard (on the end)
#define  UNUSED_4  E, 6  // hard to use with breadboard (internal)

// --- rows
#define  ROW_0  F, 7
#define  ROW_1  F, 6
#define  ROW_2  F, 5
#define  ROW_3  F, 4
#define  ROW_4  F, 1
#define  ROW_5  F, 0

// --- columns
#define  COLUMN_7  B, 0
#define  COLUMN_8  B, 1
#define  COLUMN_9  B, 2
#define  COLUMN_A  B, 3
#define  COLUMN_B  D, 2
#define  COLUMN_C  D, 3
#define  COLUMN_D  C, 6

// --- helpers
#define  SET    |=
#define  CLEAR  &=~

#define _teensypin_write(register, operation, pin_letter, pin_number)          \
  do {                                                                         \
    ((register##pin_letter)operation(1 << (pin_number)));                      \
    _delay_us(1); /* allow pins time to stabilize */                           \
  } while (0)
#define teensypin_write(register, operation, pin)                              \
  _teensypin_write(register, operation, pin)

#define _teensypin_read(pin_letter, pin_number)                                \
  ((PIN##pin_letter) & (1 << (pin_number)))
#define teensypin_read(pin) _teensypin_read(pin)

#define teensypin_write_all_unused(register, operation)                        \
  do {                                                                         \
    teensypin_write(register, operation, UNUSED_0);                            \
    teensypin_write(register, operation, UNUSED_1);                            \
    teensypin_write(register, operation, UNUSED_2);                            \
    teensypin_write(register, operation, UNUSED_3);                            \
    teensypin_write(register, operation, UNUSED_4);                            \
  } while (0)

#define teensypin_write_all_row(register, operation)                           \
  do {                                                                         \
    teensypin_write(register, operation, ROW_0);                               \
    teensypin_write(register, operation, ROW_1);                               \
    teensypin_write(register, operation, ROW_2);                               \
    teensypin_write(register, operation, ROW_3);                               \
    teensypin_write(register, operation, ROW_4);                               \
    teensypin_write(register, operation, ROW_5);                               \
  } while (0)

#define teensypin_write_all_column(register, operation)                        \
  do {                                                                         \
    teensypin_write(register, operation, COLUMN_7);                            \
    teensypin_write(register, operation, COLUMN_8);                            \
    teensypin_write(register, operation, COLUMN_9);                            \
    teensypin_write(register, operation, COLUMN_A);                            \
    teensypin_write(register, operation, COLUMN_B);                            \
    teensypin_write(register, operation, COLUMN_C);                            \
    teensypin_write(register, operation, COLUMN_D);                            \
  } while (0)


/*
 * update macros
 */
#define update_rows_for_column(matrix, column)                                 \
  do {                                                                         \
    /* set column low (set as output) */                                       \
    teensypin_write(DDR, SET, COLUMN_##column);                                \
    /* read rows 0..5 and update matrix */                                     \
    matrix[0x0][0x##column] = !teensypin_read(ROW_0);                          \
    matrix[0x1][0x##column] = !teensypin_read(ROW_1);                          \
    matrix[0x2][0x##column] = !teensypin_read(ROW_2);                          \
    matrix[0x3][0x##column] = !teensypin_read(ROW_3);                          \
    matrix[0x4][0x##column] = !teensypin_read(ROW_4);                          \
    matrix[0x5][0x##column] = !teensypin_read(ROW_5);                          \
    /* set column hi-Z (set as input) */                                       \
    teensypin_write(DDR, CLEAR, COLUMN_##column);                              \
  } while (0)

#define update_columns_for_row(matrix, row)                                    \
  do {                                                                         \
    /* set row low (set as output) */                                          \
    teensypin_write(DDR, SET, ROW_##row);                                      \
    /* read columns 7..D and update matrix */                                  \
    matrix[0x##row][0x7] = !teensypin_read(COLUMN_7);                          \
    matrix[0x##row][0x8] = !teensypin_read(COLUMN_8);                          \
    matrix[0x##row][0x9] = !teensypin_read(COLUMN_9);                          \
    matrix[0x##row][0xA] = !teensypin_read(COLUMN_A);                          \
    matrix[0x##row][0xB] = !teensypin_read(COLUMN_B);                          \
    matrix[0x##row][0xC] = !teensypin_read(COLUMN_C);                          \
    matrix[0x##row][0xD] = !teensypin_read(COLUMN_D);                          \
    /* set row hi-Z (set as input) */                                          \
    teensypin_write(DDR, CLEAR, ROW_##row);                                    \
  } while (0)

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

/* returns
 * - success: 0
 */
uint8_t teensy_init(void) {
  // CPU speed : should match F_CPU in makefile
#if F_CPU != 16000000
#error "Expecting different CPU frequency"
#endif
  CPU_PRESCALE(CPU_16MHz);

  // onboard LED
  // (tied to GND for hardware convenience)
  DDRD &= ~(1 << 6);  // set D(6) as input
  PORTD &= ~(1 << 6); // set D(6) internal pull-up disabled

  // (tied to Vcc for hardware convenience)
  DDRB &= ~(1 << 4);  // set B(4) as input
  PORTB &= ~(1 << 4); // set B(4) internal pull-up disabled

  // keyboard LEDs (see "PWM on ports OC1(A|B|C)" in "teensy-2-0.md")
  TCCR1A = 0b10101001; // set and configure fast PWM
  TCCR1B = 0b00001001; // set and configure fast PWM

  // I2C (TWI)
  twi_init(); // on pins D(1,0)

  // unused pins
  teensypin_write_all_unused(DDR, CLEAR); // set as input
  teensypin_write_all_unused(PORT, SET);  // set internal pull-up enabled

  // rows and columns
  teensypin_write_all_row(DDR, CLEAR);    // set as input (hi-Z)
  teensypin_write_all_column(DDR, CLEAR); // set as input (hi-Z)
#if TEENSY__DRIVE_ROWS
  teensypin_write_all_row(PORT, CLEAR);  // pull-up disabled
  teensypin_write_all_column(PORT, SET); // pull-up enabled
#elif TEENSY__DRIVE_COLUMNS
  teensypin_write_all_row(PORT, SET);      // pull-up enabled
  teensypin_write_all_column(PORT, CLEAR); // pull-up disabled
#endif

  return 0; // success
}

/* returns
 * - success: 0
 */
#if KB_ROWS != 6 || KB_COLUMNS != 14
#error "Expecting different keyboard dimensions"
#endif

uint8_t teensy_update_matrix(bool matrix[KB_ROWS][KB_COLUMNS]) {
#if TEENSY__DRIVE_ROWS
  update_columns_for_row(matrix, 0);
  update_columns_for_row(matrix, 1);
  update_columns_for_row(matrix, 2);
  update_columns_for_row(matrix, 3);
  update_columns_for_row(matrix, 4);
  update_columns_for_row(matrix, 5);
#elif TEENSY__DRIVE_COLUMNS
  update_rows_for_column(matrix, 7);
  update_rows_for_column(matrix, 8);
  update_rows_for_column(matrix, 9);
  update_rows_for_column(matrix, A);
  update_rows_for_column(matrix, B);
  update_rows_for_column(matrix, C);
  update_rows_for_column(matrix, D);
#endif

  return 0; // success
}

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

void twi_init(void) {
  // set the prescaler value to 0
  TWSR &= ~((1 << TWPS1) | (1 << TWPS0));
  // set the bit rate
  // - TWBR should be 10 or higher (datasheet section 20.5.2)
  // - TWI_FREQ should be 400000 (400kHz) max (datasheet section 20.1)
  TWBR = ((F_CPU / TWI_FREQ) - 16) / 2;
}

uint8_t twi_start(void) {
  // send start
  TWCR = (1 << TWINT) | (1 << TWEN) | (1 << TWSTA);
  // wait for transmission to complete
  while (!(TWCR & (1 << TWINT)))
    ;
  // if it didn't work, return the status code (else return 0)
  if ((TW_STATUS != TW_START) && (TW_STATUS != TW_REP_START))
    return TW_STATUS; // error
  return 0;           // success
}

void twi_stop(void) {
  // send stop
  TWCR = (1 << TWINT) | (1 << TWEN) | (1 << TWSTO);
  // wait for transmission to complete
  while (TWCR & (1 << TWSTO))
    ;
}

uint8_t twi_send(uint8_t data) {
  // load data into the data register
  TWDR = data;
  // send data
  TWCR = (1 << TWINT) | (1 << TWEN);
  // wait for transmission to complete
  while (!(TWCR & (1 << TWINT)))
    ;
  // if it didn't work, return the status code (else return 0)
  if ((TW_STATUS != TW_MT_SLA_ACK) && (TW_STATUS != TW_MT_DATA_ACK) &&
      (TW_STATUS != TW_MR_SLA_ACK))
    return TW_STATUS; // error
  return 0;           // success
}

uint8_t twi_read(uint8_t *data) {
  // read 1 byte to TWDR, send ACK
  TWCR = (1 << TWINT) | (1 << TWEN) | (1 << TWEA);
  // wait for transmission to complete
  while (!(TWCR & (1 << TWINT)))
    ;
  // set data variable
  *data = TWDR;
  // if it didn't work, return the status code (else return 0)
  if (TW_STATUS != TW_MR_DATA_ACK)
    return TW_STATUS; // error
  return 0;           // success
}

// ----------------------------------------------
// usb
// ----------------------------------------------

/**************************************************************************
 *
 *  Configurable Options
 *
 **************************************************************************/

// You can change these to give your code its own name.
#define STR_MANUFACTURER	L"FalbabeTech"
#define STR_PRODUCT     	L"ErgoDox ergonomic keyboard"


// Mac OS-X and Linux automatically load the correct drivers.  On
// Windows, even though the driver is supplied by Microsoft, an
// INF file is needed to load the driver.  These numbers need to
// match the INF file.
#define VENDOR_ID 	0x1d50  // Openmoko, Inc.
#define PRODUCT_ID	0x6028  // ErgoDox ergonomic keyboard


// USB devices are supposed to implment a halt feature, which is
// rarely (if ever) used.  If you comment this line out, the halt
// code will be removed, saving 102 bytes of space (gcc 4.3.0).
// This is not strictly USB compliant, but works with all major
// operating systems.
#define SUPPORT_ENDPOINT_HALT

/* report id */
#define REPORT_ID_SYSTEM  	2
#define REPORT_ID_CONSUMER	3

/**************************************************************************
 *
 *  Endpoint Buffer Configuration
 *
 **************************************************************************/
#define ENDPOINT0_SIZE	32

#define KEYBOARD_INTERFACE	0
#define KEYBOARD_ENDPOINT 	1
#define KEYBOARD_SIZE     	8
#define KEYBOARD_BUFFER   	EP_DOUBLE_BUFFER

#define EXTRA_INTERFACE	1
#define EXTRA_ENDPOINT 	2
#define EXTRA_SIZE     	8
#define EXTRA_BUFFER   	EP_DOUBLE_BUFFER

#define DEBUG_INTERFACE  	2
#define DEBUG_TX_ENDPOINT	3
#define DEBUG_TX_SIZE    	32
#define DEBUG_TX_BUFFER  	EP_DOUBLE_BUFFER

static const uint8_t PROGMEM endpoint_config_table[] = {
  // enable, UECFG0X(type, direction), UECFG1X(size, bank, allocation)
  1, EP_TYPE_INTERRUPT_IN,  EP_SIZE(KEYBOARD_SIZE) | KEYBOARD_BUFFER, // 1
  1, EP_TYPE_INTERRUPT_IN,  EP_SIZE(EXTRA_SIZE)    | EXTRA_BUFFER,    // 2
#ifdef KBD_DEBUG
  1, EP_TYPE_INTERRUPT_IN,  EP_SIZE(DEBUG_TX_SIZE) | DEBUG_TX_BUFFER, // 3
#endif
  0
};


/**************************************************************************
 *
 *  Descriptor Data
 *
 **************************************************************************/

// Descriptors are the data that your computer reads when it auto-detects
// this USB device (called "enumeration" in USB lingo).  The most commonly
// changed items are editable at the top of this file.  Changing things
// in here should only be done by those who've read chapter 9 of the USB
// spec and relevant portions of any USB class specifications!


static const uint8_t PROGMEM device_descriptor[] = {
  18,                              	// bLength
  1,                               	// bDescriptorType
  0x00, 0x02,                      	// bcdUSB
  0,                               	// bDeviceClass
  0,                               	// bDeviceSubClass
  0,                               	// bDeviceProtocol
  ENDPOINT0_SIZE,                  	// bMaxPacketSize0
  LSB(VENDOR_ID), MSB(VENDOR_ID),  	// idVendor
  LSB(PRODUCT_ID), MSB(PRODUCT_ID),	// idProduct
  0x00, 0x01,                      	// bcdDevice
  1,                               	// iManufacturer
  2,                               	// iProduct
  0,                               	// iSerialNumber
  1                                	// bNumConfigurations
};

// Keyboard Protocol 1, HID 1.11 spec, Appendix B, page 59-60
static const uint8_t PROGMEM keyboard_hid_report_desc[] = {
  0x05, 0x01,	// Usage Page (Generic Desktop),
  0x09, 0x06,	// Usage (Keyboard),
  0xA1, 0x01,	// Collection (Application),
  0x75, 0x01,	// Report Size (1),
  0x95, 0x08,	// Report Count (8),
  0x05, 0x07,	// Usage Page (Key Codes),
  0x19, 0xE0,	// Usage Minimum (224),
  0x29, 0xE7,	// Usage Maximum (231),
  0x15, 0x00,	// Logical Minimum (0),
  0x25, 0x01,	// Logical Maximum (1),
  0x81, 0x02,	// Input (Data, Variable, Absolute), Modifier byte
  0x95, 0x01,	// Report Count (1),
  0x75, 0x08,	// Report Size (8),
  0x81, 0x03,	// Input (Constant), Reserved byte
  0x95, 0x05,	// Report Count (5),
  0x75, 0x01,	// Report Size (1),
  0x05, 0x08,	// Usage Page (LEDs),
  0x19, 0x01,	// Usage Minimum (1),
  0x29, 0x05,	// Usage Maximum (5),
  0x91, 0x02,	// Output (Data, Variable, Absolute), LED report
  0x95, 0x01,	// Report Count (1),
  0x75, 0x03,	// Report Size (3),
  0x91, 0x03,	// Output (Constant), LED report padding
  0x95, 0x06,	// Report Count (6),
  0x75, 0x08,	// Report Size (8),
  0x15, 0x00,	// Logical Minimum (0),
  0x25, 0xFF,	// Logical Maximum(255),
  0x05, 0x07,	// Usage Page (Key Codes),
  0x19, 0x00,	// Usage Minimum (0),
  0x29, 0xFF,	// Usage Maximum (255),
  0x81, 0x00,	// Input (Data, Array),
  0xc0       	// End Collection
};

// audio controls & system controls
// http://www.microsoft.com/whdc/archive/w2kbd.mspx
static const uint8_t PROGMEM extra_hid_report_desc[] = {
  /* consumer */
  0x05, 0x0c,              	// USAGE_PAGE (Consumer Devices)
  0x09, 0x01,              	// USAGE (Consumer Control)
  0xa1, 0x01,              	// COLLECTION (Application)
  0x85, REPORT_ID_CONSUMER,	// REPORT_ID (3)
  0x15, 0x01,              	// LOGICAL_MINIMUM (0x1)
  0x26, 0x9c, 0x02,        	// LOGICAL_MAXIMUM (0x29c)
  0x19, 0x01,              	// USAGE_MINIMUM (0x1)
  0x2a, 0x9c, 0x02,        	// USAGE_MAXIMUM (0x29c)
  0x75, 0x10,              	// REPORT_SIZE (16)
  0x95, 0x01,              	// REPORT_COUNT (1)
  0x81, 0x00,              	// INPUT (Data,Array,Abs)
  0xc0,                    	// END_COLLECTION
};

// debug messages
#ifdef KBD_DEBUG
static const uint8_t PROGMEM debug_hid_report_desc[] = {
  0x06, 0x31, 0xFF,   	// Usage Page 0xFF31 (vendor defined)
  0x09, 0x74,         	// Usage 0x74
  0xA1, 0x53,         	// Collection 0x53
  0x75, 0x08,         	// report size = 8 bits
  0x15, 0x00,         	// logical minimum = 0
  0x26, 0xFF, 0x00,   	// logical maximum = 255
  0x95, DEBUG_TX_SIZE,	// report count
  0x09, 0x75,         	// usage
  0x81, 0x02,         	// Input (array)
  0xC0                	// end collection
};
#endif

#ifdef KBD_DEBUG
#define KEYBOARD_HID_DESC_NUM	0
#define EXTRA_HID_DESC_NUM   	1
#define DEBUG_HID_DESC_NUM   	2
#define NUM_INTERFACES       	3
#else
#define KEYBOARD_HID_DESC_NUM	0
#define EXTRA_HID_DESC_NUM   	1
#define NUM_INTERFACES       	2
#endif

#define KEYBOARD_HID_DESC_OFFSET	(9+(9+9+7)*KEYBOARD_HID_DESC_NUM+9)
#define EXTRA_HID_DESC_OFFSET   	(9+(9+9+7)*EXTRA_HID_DESC_NUM+9)
#define DEBUG_HID_DESC_OFFSET   	(9+(9+9+7)*DEBUG_HID_DESC_NUM+9)
#define CONFIG1_DESC_SIZE       	(9+(9+9+7)*NUM_INTERFACES)

static const uint8_t PROGMEM config1_descriptor[CONFIG1_DESC_SIZE] = {
    // configuration descriptor, USB spec 9.6.3, page 264-266, Table 9-10
    9,                      // bLength;
    2,                      // bDescriptorType;
    LSB(CONFIG1_DESC_SIZE), // wTotalLength
    MSB(CONFIG1_DESC_SIZE),
    NUM_INTERFACES, // bNumInterfaces
    1,              // bConfigurationValue
    0,              // iConfiguration
    0xA0,           // bmAttributes
    50,             // bMaxPower

    // interface descriptor, USB spec 9.6.5, page 267-269, Table 9-12
    9,                  // bLength
    4,                  // bDescriptorType
    KEYBOARD_INTERFACE, // bInterfaceNumber
    0,                  // bAlternateSetting
    1,                  // bNumEndpoints
    0x03,               // bInterfaceClass (0x03 = HID)
    0x01,               // bInterfaceSubClass (0x01 = Boot)
    0x01,               // bInterfaceProtocol (0x01 = Keyboard)
    0,                  // iInterface
    // HID descriptor, HID 1.11 spec, section 6.2.1
    9,                                // bLength
    0x21,                             // bDescriptorType
    0x11, 0x01,                       // bcdHID
    0,                                // bCountryCode
    1,                                // bNumDescriptors
    0x22,                             // bDescriptorType
    sizeof(keyboard_hid_report_desc), // wDescriptorLength
    0,
    // endpoint descriptor, USB spec 9.6.6, page 269-271, Table 9-13
    7,                        // bLength
    5,                        // bDescriptorType
    KEYBOARD_ENDPOINT | 0x80, // bEndpointAddress
    0x03,                     // bmAttributes (0x03=intr)
    KEYBOARD_SIZE, 0,         // wMaxPacketSize
    10,                       // bInterval

    // interface descriptor, USB spec 9.6.5, page 267-269, Table 9-12
    9,               // bLength
    4,               // bDescriptorType
    EXTRA_INTERFACE, // bInterfaceNumber
    0,               // bAlternateSetting
    1,               // bNumEndpoints
    0x03,            // bInterfaceClass (0x03 = HID)
    0x00,            // bInterfaceSubClass
    0x00,            // bInterfaceProtocol
    0,               // iInterface
    // HID descriptor, HID 1.11 spec, section 6.2.1
    9,                             // bLength
    0x21,                          // bDescriptorType
    0x11, 0x01,                    // bcdHID
    0,                             // bCountryCode
    1,                             // bNumDescriptors
    0x22,                          // bDescriptorType
    sizeof(extra_hid_report_desc), // wDescriptorLength
    0,
    // endpoint descriptor, USB spec 9.6.6, page 269-271, Table 9-13
    7,                     // bLength
    5,                     // bDescriptorType
    EXTRA_ENDPOINT | 0x80, // bEndpointAddress
    0x03,                  // bmAttributes (0x03=intr)
    EXTRA_SIZE, 0,         // wMaxPacketSize
    10,                    // bInterval

#ifdef KBD_DEBUG
    // interface descriptor, USB spec 9.6.5, page 267-269, Table 9-12
    9,               // bLength
    4,               // bDescriptorType
    DEBUG_INTERFACE, // bInterfaceNumber
    0,               // bAlternateSetting
    1,               // bNumEndpoints
    0x03,            // bInterfaceClass (0x03 = HID)
    0x00,            // bInterfaceSubClass
    0x00,            // bInterfaceProtocol
    0,               // iInterface
    // HID descriptor, HID 1.11 spec, section 6.2.1
    9,                             // bLength
    0x21,                          // bDescriptorType
    0x11, 0x01,                    // bcdHID
    0,                             // bCountryCode
    1,                             // bNumDescriptors
    0x22,                          // bDescriptorType
    sizeof(debug_hid_report_desc), // wDescriptorLength
    0,
    // endpoint descriptor, USB spec 9.6.6, page 269-271, Table 9-13
    7,                        // bLength
    5,                        // bDescriptorType
    DEBUG_TX_ENDPOINT | 0x80, // bEndpointAddress
    0x03,                     // bmAttributes (0x03=intr)
    DEBUG_TX_SIZE, 0,         // wMaxPacketSize
    1,                        // bInterval
#endif
};

// If you're desperate for a little extra code memory, these strings
// can be completely removed if iManufacturer, iProduct, iSerialNumber
// in the device desciptor are changed to zeros.
struct usb_string_descriptor_struct {
  uint8_t bLength;
  uint8_t bDescriptorType;
  int16_t wString[];
};
static const struct usb_string_descriptor_struct PROGMEM string0 = {
    4, 3, {0x0409}};
static const struct usb_string_descriptor_struct PROGMEM string1 = {
    sizeof(STR_MANUFACTURER), 3, STR_MANUFACTURER};
static const struct usb_string_descriptor_struct PROGMEM string2 = {
    sizeof(STR_PRODUCT), 3, STR_PRODUCT};

// This table defines which descriptor data is sent for each specific
// request from the host (in wValue and wIndex).
static struct descriptor_list_struct {
  uint16_t wValue;
  uint16_t wIndex;
  const uint8_t *addr;
  uint8_t length;
} const PROGMEM descriptor_list[] = {
    // DEVICE descriptor
    {0x0100, 0x0000, device_descriptor, sizeof(device_descriptor)},
    // CONFIGURATION descriptor
    {0x0200, 0x0000, config1_descriptor, sizeof(config1_descriptor)},
    // HID/REPORT descriptors
    {0x2100, KEYBOARD_INTERFACE, config1_descriptor + KEYBOARD_HID_DESC_OFFSET, 9},
    {0x2200, KEYBOARD_INTERFACE, keyboard_hid_report_desc, sizeof(keyboard_hid_report_desc)},
    // Extra HID Descriptor
    {0x2100, EXTRA_INTERFACE, config1_descriptor + EXTRA_HID_DESC_OFFSET, 9},
    {0x2200, EXTRA_INTERFACE, extra_hid_report_desc, sizeof(extra_hid_report_desc)},
#ifdef KBD_DEBUG
    // debug descriptors
    {0x2100, DEBUG_INTERFACE, config1_descriptor+DEBUG_HID_DESC_OFFSET, 9},
    {0x2200, DEBUG_INTERFACE, debug_hid_report_desc, sizeof(debug_hid_report_desc)},
#endif
    // STRING descriptors
    {0x0300, 0x0000, (const uint8_t *)&string0, 4},
    {0x0301, 0x0409, (const uint8_t *)&string1, sizeof(STR_MANUFACTURER)},
    {0x0302, 0x0409, (const uint8_t *)&string2, sizeof(STR_PRODUCT)}};
#define NUM_DESC_LIST (sizeof(descriptor_list)/sizeof(struct descriptor_list_struct))


/**************************************************************************
 *
 *  Variables - these are the only non-stack RAM usage
 *
 **************************************************************************/

// zero when we are not configured, non-zero when enumerated
static volatile uint8_t usb_configuration=0;

// which modifier keys are currently pressed
// 1=left ctrl,    2=left shift,   4=left alt,    8=left gui
// 16=right ctrl, 32=right shift, 64=right alt, 128=right gui
uint8_t keyboard_modifier_keys=0;

// which keys are currently pressed, up to 6 keys may be down at once
uint8_t keyboard_keys[6]={0,0,0,0,0,0};

// protocol setting from the host.  We use exactly the same report
// either way, so this variable only stores the setting since we
// are required to be able to report which setting is in use.
static uint8_t keyboard_protocol=1;

// the idle configuration, how often we send the report to the
// host (ms * 4) even when it hasn't changed
static uint8_t keyboard_idle_config=125;

// count until idle timeout
static uint8_t keyboard_idle_count=0;

// which consumer key is currently pressed
uint16_t consumer_key;
uint16_t last_consumer_key;

volatile uint8_t debug_flush_timer=0;

/**************************************************************************
 *
 *  Public Functions - these are the API intended for the user
 *
 **************************************************************************/


// initialize USB
void usb_init(void) {
  HW_CONFIG();
  USB_FREEZE(); // enable USB
  PLL_CONFIG(); // config PLL
  while (!(PLLCSR & (1 << PLOCK))) {} // wait for PLL lock
  USB_CONFIG(); // start USB clock
  UDCON = 0;    // enable attach resistor
  usb_configuration = 0;
  UDIEN = (1 << EORSTE) | (1 << SOFE);
  sei();
}
// return 0 if the USB is not configured, or the configuration
// number selected by the HOST
uint8_t usb_configured(void) {
  return usb_configuration;
}


// perform a single keystroke
int8_t usb_keyboard_press(uint8_t key, uint8_t modifier) {
  int8_t r;

  keyboard_modifier_keys = modifier;
  keyboard_keys[0] = key;
  r = usb_keyboard_send();
  if (r) { return r; }
  keyboard_modifier_keys = 0;
  keyboard_keys[0] = 0;
  return usb_keyboard_send();
}

// send the contents of keyboard_keys and keyboard_modifier_keys
int8_t usb_keyboard_send(void) {
  uint8_t i, intr_state, timeout;

  if (!usb_configuration)
    return -1;
  intr_state = SREG;
  cli();
  UENUM = KEYBOARD_ENDPOINT;
  timeout = UDFNUML + 50;
  while (1) {
    // are we ready to transmit?
    if (UEINTX & (1 << RWAL)) { break; }
    SREG = intr_state;
    // has the USB gone offline?
    if (!usb_configuration) { return -1; }
    // have we waited too long?
    if (UDFNUML == timeout) { return -1; }
    // get ready to try checking again
    intr_state = SREG;
    cli();
    UENUM = KEYBOARD_ENDPOINT;
  }
  UEDATX = keyboard_modifier_keys;
  UEDATX = 0;
  for (i = 0; i < 6; i++) {
    UEDATX = keyboard_keys[i];
  }
  UEINTX = 0x3A;
  keyboard_idle_count = 0;
  SREG = intr_state;
  return 0;
}

/**************************************************************************
 *
 *  Private Functions - not intended for general user consumption....
 *
 **************************************************************************/

// USB Device Interrupt - handle all device-level events
// the transmit buffer flushing is triggered by the start of frame
//
ISR(USB_GEN_vect) {
  uint8_t intbits, i; // used to declare a variable `t` as well, but it
                      //   wasn't used ::Ben Blazak, 2012::
  static uint8_t div4 = 0;

  intbits = UDINT;
  UDINT = 0;
  if (intbits & (1 << EORSTI)) {
    UENUM = 0;
    UECONX = 1;
    UECFG0X = EP_TYPE_CONTROL;
    UECFG1X = EP_SIZE(ENDPOINT0_SIZE) | EP_SINGLE_BUFFER;
    UEIENX = (1 << RXSTPE);
    usb_configuration = 0;
  }
  if ((intbits & (1 << SOFI)) && usb_configuration) {
    if (keyboard_idle_config && (++div4 & 3) == 0) {
      UENUM = KEYBOARD_ENDPOINT;
      if (UEINTX & (1 << RWAL)) {
        keyboard_idle_count++;
        if (keyboard_idle_count == keyboard_idle_config) {
          keyboard_idle_count = 0;
          UEDATX = keyboard_modifier_keys;
          UEDATX = 0;
          for (i = 0; i < 6; i++) {
            UEDATX = keyboard_keys[i];
          }
          UEINTX = 0x3A;
        }
      }
    }
  }
}


// Misc functions to wait for ready and send/receive packets
static inline void usb_wait_in_ready(void) {
  while (!(UEINTX & (1 << TXINI))) {}
}
static inline void usb_send_in(void) { UEINTX = ~(1 << TXINI); }
static inline void usb_wait_receive_out(void) {
  while (!(UEINTX & (1 << RXOUTI))) {}
}
static inline void usb_ack_out(void) { UEINTX = ~(1 << RXOUTI); }



// USB Endpoint Interrupt - endpoint 0 is handled here.  The
// other endpoints are manipulated by the user-callable
// functions, and the start-of-frame interrupt.
//
ISR(USB_COM_vect) {
  uint8_t intbits;
  const uint8_t *list;
  const uint8_t *cfg;
  uint8_t i, n, len, en;
  uint8_t bmRequestType;
  uint8_t bRequest;
  uint16_t wValue;
  uint16_t wIndex;
  uint16_t wLength;
  uint16_t desc_val;
  const uint8_t *desc_addr;
  uint8_t desc_length;

  UENUM = 0;
  intbits = UEINTX;
  if (intbits & (1 << RXSTPI)) {
    bmRequestType = UEDATX;
    bRequest = UEDATX;
    wValue = UEDATX;
    wValue |= (UEDATX << 8);
    wIndex = UEDATX;
    wIndex |= (UEDATX << 8);
    wLength = UEDATX;
    wLength |= (UEDATX << 8);
    UEINTX = ~((1 << RXSTPI) | (1 << RXOUTI) | (1 << TXINI));
    if (bRequest == GET_DESCRIPTOR) {
      list = (const uint8_t *)descriptor_list;
      for (i = 0;; i++) {
        if (i >= NUM_DESC_LIST) {
          UECONX = (1 << STALLRQ) | (1 << EPEN); // stall
          return;
        }
        desc_val = pgm_read_word(list);
        if (desc_val != wValue) {
          list += sizeof(struct descriptor_list_struct);
          continue;
        }
        list += 2;
        desc_val = pgm_read_word(list);
        if (desc_val != wIndex) {
          list += sizeof(struct descriptor_list_struct) - 2;
          continue;
        }
        list += 2;
        desc_addr = (const uint8_t *)pgm_read_word(list);
        list += 2;
        desc_length = pgm_read_byte(list);
        break;
      }
      len = (wLength < 256) ? wLength : 255;
      if (len > desc_length)
        len = desc_length;
      do {
        // wait for host ready for IN packet
        do {
          i = UEINTX;
        } while (!(i & ((1 << TXINI) | (1 << RXOUTI))));
        if (i & (1 << RXOUTI))
          return; // abort
        // send IN packet
        n = len < ENDPOINT0_SIZE ? len : ENDPOINT0_SIZE;
        for (i = n; i; i--) {
          UEDATX = pgm_read_byte(desc_addr++);
        }
        len -= n;
        usb_send_in();
      } while (len || n == ENDPOINT0_SIZE);
      return;
    }
    if (bRequest == SET_ADDRESS) {
      usb_send_in();
      usb_wait_in_ready();
      UDADDR = wValue | (1 << ADDEN);
      return;
    }
    if (bRequest == SET_CONFIGURATION && bmRequestType == 0) {
      usb_configuration = wValue;
      usb_send_in();
      cfg = endpoint_config_table;
      for (i = 1; i < 5; i++) {
        UENUM = i;
        en = pgm_read_byte(cfg++);
        UECONX = en;
        if (en) {
          UECFG0X = pgm_read_byte(cfg++);
          UECFG1X = pgm_read_byte(cfg++);
        }
      }
      UERST = 0x1E;
      UERST = 0;
      return;
    }
    if (bRequest == GET_CONFIGURATION && bmRequestType == 0x80) {
      usb_wait_in_ready();
      UEDATX = usb_configuration;
      usb_send_in();
      return;
    }

    if (bRequest == GET_STATUS) {
      usb_wait_in_ready();
      i = 0;
#ifdef SUPPORT_ENDPOINT_HALT
      if (bmRequestType == 0x82) {
        UENUM = wIndex;
        if (UECONX & (1 << STALLRQ))
          i = 1;
        UENUM = 0;
      }
#endif
      UEDATX = i;
      UEDATX = 0;
      usb_send_in();
      return;
    }
#ifdef SUPPORT_ENDPOINT_HALT
    if ((bRequest == CLEAR_FEATURE || bRequest == SET_FEATURE) &&
        bmRequestType == 0x02 && wValue == 0) {
      i = wIndex & 0x7F;
      if (i >= 1 && i <= MAX_ENDPOINT) {
        usb_send_in();
        UENUM = i;
        if (bRequest == SET_FEATURE) {
          UECONX = (1 << STALLRQ) | (1 << EPEN);
        } else {
          UECONX = (1 << STALLRQC) | (1 << RSTDT) | (1 << EPEN);
          UERST = (1 << i);
          UERST = 0;
        }
        return;
      }
    }
#endif
    if (wIndex == KEYBOARD_INTERFACE) {
      if (bmRequestType == 0xA1) {
        if (bRequest == HID_GET_REPORT) {
          usb_wait_in_ready();
          UEDATX = keyboard_modifier_keys;
          UEDATX = 0;
          for (i = 0; i < 6; i++) {
            UEDATX = keyboard_keys[i];
          }
          usb_send_in();
          return;
        }
        if (bRequest == HID_GET_IDLE) {
          usb_wait_in_ready();
          UEDATX = keyboard_idle_config;
          usb_send_in();
          return;
        }
        if (bRequest == HID_GET_PROTOCOL) {
          usb_wait_in_ready();
          UEDATX = keyboard_protocol;
          usb_send_in();
          return;
        }
      }
      if (bmRequestType == 0x21) {
        if (bRequest == HID_SET_REPORT) {
          usb_wait_receive_out();
          usb_ack_out();
          usb_send_in();
          return;
        }
        if (bRequest == HID_SET_IDLE) {
          keyboard_idle_config = (wValue >> 8);
          keyboard_idle_count = 0;
          usb_send_in();
          return;
        }
        if (bRequest == HID_SET_PROTOCOL) {
          keyboard_protocol = wValue;
          usb_send_in();
          return;
        }
      }
    }
  }
  UECONX = (1 << STALLRQ) | (1 << EPEN); // stall
}

int8_t usb_extra_send(uint8_t report_id, uint16_t data) {
  uint8_t intr_state, timeout;

  if (!usb_configured())
    return -1;
  intr_state = SREG;
  cli();
  UENUM = EXTRA_ENDPOINT;
  timeout = UDFNUML + 50;
  while (1) {
    // are we ready to transmit?
    if (UEINTX & (1 << RWAL)) { break; }
    SREG = intr_state;
    // has the USB gone offline?
    if (!usb_configured()) { return -1; }
    // have we waited too long?
    if (UDFNUML == timeout) { return -1; }
    // get ready to try checking again
    intr_state = SREG;
    cli();
    UENUM = EXTRA_ENDPOINT;
  }

  UEDATX = report_id;
  UEDATX = data & 0xFF;
  UEDATX = (data >> 8) & 0xFF;

  UEINTX = 0x3A;
  SREG = intr_state;
  return 0;
}

int8_t usb_extra_consumer_send() {
  int result = 0;
  // don't resend the same key repeatedly if held, only send it once.
  if (consumer_key != last_consumer_key) {
    result = usb_extra_send(REPORT_ID_CONSUMER, consumer_key);
    if (result == 0) {
      last_consumer_key = consumer_key;
    }
  }
  return result;
}

// --------------------------------------------------------
// debug
// --------------------------------------------------------

#ifdef KBD_DEBUG
// transmit a character.  0 returned on success, -1 on error
int8_t usb_debug_putchar(char c) {
  static uint8_t previous_timeout = 0;
  uint8_t timeout, intr_state;

  // if we're not online (enumerated and configured), error
  if (!usb_configuration)
    return -1;
  // interrupts are disabled so these functions can be
  // used from the main program or interrupt context,
  // even both in the same program!
  intr_state = SREG;
  cli();
  UENUM = DEBUG_TX_ENDPOINT;
  // if we gave up due to timeout before, don't wait again
  if (previous_timeout) {
    if (!(UEINTX & (1 << RWAL))) {
      SREG = intr_state;
      return -1;
    }
    previous_timeout = 0;
  }
  // wait for the FIFO to be ready to accept data
  timeout = UDFNUML + 4;
  while (1) {
    // are we ready to transmit?
    if (UEINTX & (1 << RWAL))
      break;
    SREG = intr_state;
    // have we waited too long?
    if (UDFNUML == timeout) {
      previous_timeout = 1;
      return -1;
    }
    // has the USB gone offline?
    if (!usb_configuration)
      return -1;
    // get ready to try checking again
    intr_state = SREG;
    cli();
    UENUM = DEBUG_TX_ENDPOINT;
  }
  // actually write the byte into the FIFO
  UEDATX = c;
  // if this completed a packet, transmit it now!
  if (!(UEINTX & (1 << RWAL))) {
    UEINTX = 0x3A;
    debug_flush_timer = 0;
  } else {
    debug_flush_timer = 2;
  }
  SREG = intr_state;
  return 0;
}

// immediately transmit any buffered output.
void usb_debug_flush_output(void) {
  uint8_t intr_state;

  intr_state = SREG;
  cli();
  if (debug_flush_timer) {
    UENUM = DEBUG_TX_ENDPOINT;
    while ((UEINTX & (1 << RWAL))) {
      UEDATX = 0;
    }
    UEINTX = 0x3A;
    debug_flush_timer = 0;
  }
  SREG = intr_state;
}

void usb_debug_print(const char *s) {
  char c;
  while ((c = pgm_read_byte(c))) {
    usb_debug_putchar(c);
  }
  usb_debug_flush_output();
}

void usb_debug_printf(const char *fmt, ...) {
  char buf[128];

  va_list args;
  va_start(args, fmt);
  vsnprintf_P(buf, sizeof(buf), fmt, args);
  va_end(args);

  for (int i = 0; i < sizeof(buf); i++) {
    char c = buf[i];
    if (!c) break;
    usb_debug_putchar(c);
  }
  usb_debug_flush_output();
}

void usb_debug_free_memory() {
  extern int __heap_start, *__brkval; int v, free;
  free = (int) &v - (__brkval == 0 ? (int) &__heap_start : (int) __brkval);
  debug_printf("free memory: %d bytes\n", free);
}
#endif
