#include <avr/io.h>
#include <avr/interrupt.h>
#include <avr/sleep.h>
#define F_CPU 8000000UL
#include <util/delay.h>

#define USI_PULSE_0 (1<<USIWM0)|(1<<USITC)
#define USI_PULSE_1 (1<<USIWM0)|(1<<USICLK)|(1<<USITC)
#define SAMPLE_DEF_SIZE 6

register struct {
	uint8_t tick: 1;
	uint8_t play: 1;
} flags asm("r15");
register uint8_t sample_buffer asm("r14");

uint8_t n_samples;
uint8_t button_state;
uint8_t channels;
uint16_t button_debounce;
uint32_t sample_count;

void usi_init() {
	// Set CS high initially
	PORTA = 0b10000000;
	// Set PA7 (CS) PA5 (DO) and PA4 (USCK) to outputs
	DDRA = 0b10110000;
}

void usi_deinit() {
	PORTA = 0;
	DDRA = 0;
}

#define spi_cs_enable() (PORTA &= ~0b10000000)
#define spi_cs_disable() (PORTA |= 0b10000000)

uint8_t spi_transfer(uint8_t d) {
	cli();
	USIDR = d;
	USICR = USI_PULSE_0;
	USICR = USI_PULSE_1;
	USICR = USI_PULSE_0;
	USICR = USI_PULSE_1;
	USICR = USI_PULSE_0;
	USICR = USI_PULSE_1;
	USICR = USI_PULSE_0;
	USICR = USI_PULSE_1;
	USICR = USI_PULSE_0;
	USICR = USI_PULSE_1;
	USICR = USI_PULSE_0;
	USICR = USI_PULSE_1;
	USICR = USI_PULSE_0;
	USICR = USI_PULSE_1;
	USICR = USI_PULSE_0;
	USICR = USI_PULSE_1;
	sei();
	return USIDR;
}

void flash_id() {
	spi_cs_enable();
	spi_transfer(0x9F);
	spi_transfer(0x00);
	spi_transfer(0x00);
	spi_transfer(0x00);
	spi_cs_disable();
}

void begin_playback(uint8_t sample) {
	sample_buffer = 0x7F;
	uint8_t addr[3];
	uint16_t dir_addr = sample * SAMPLE_DEF_SIZE + 1;
	uint8_t* sc = (uint8_t*) &sample_count;
	sc[3] = 0;

	usi_init();

	// Fetch directory information
	spi_cs_enable();
	spi_transfer(0x03);
	spi_transfer(0);
	spi_transfer(dir_addr >> 8);
	spi_transfer(dir_addr & 0xFF);
	addr[0] = spi_transfer(0);
	addr[1] = spi_transfer(0);
	addr[2] = spi_transfer(0);
	sc[2] = spi_transfer(0);
	sc[1] = spi_transfer(0);
	sc[0] = spi_transfer(0);
	spi_cs_disable();

	// Begin SPI transfer
	spi_cs_enable();
	spi_transfer(0x03);
	spi_transfer(addr[0]);
	spi_transfer(addr[1]);
	spi_transfer(addr[2]);

	if (channels == 0) {
		// Wait for the amp to power up
		_delay_ms(200);
	}
	channels = 1;

	// Clear OC0A on compare match, Fast PWM
	TCCR0A = 0b10000011;
	// Reset counter
	TCNT0 = 0;
	// Enable overflow interrupt
	TIMSK0 = _BV(TOIE0);
	// Match at sample
	OCR0A = sample_buffer;
	// Enable timer - Fast PWM, no prescaler
	TCCR0B = 0b00000001;

	// Switch to idle sleep because we need the timer/counter to function
	set_sleep_mode(SLEEP_MODE_IDLE);
}

void end_playback() {
	// Disable overflow interrupt
	TIMSK0 = 0;
	// disable timer
	TCCR0B = 0;
	// End SPI transfer
	spi_cs_disable();
	// Return USI pins to Hi-Z
	usi_deinit();

	// Return to power-off sleep
	set_sleep_mode(SLEEP_MODE_PWR_DOWN);

	channels = 0;
	button_debounce = 0;
}

uint8_t adc_read() {
	// VCC reference, select PA0
	ADMUX = 0;
	// Enable, Start conversion, Prescaler divisor 16
	ADCSRA = _BV(ADEN) | _BV(ADSC) | 0b100;

	// Wait for conversion to finish
	while (!(ADCSRA & _BV(ADIF)));
	// Clear interrupt flag
	ADCSRA = _BV(ADIF);
	// Both registers must be read
	uint8_t tmp = ADCL;
	ADCH;
	// Shut off the ADC
	ADCSRA = 0;
	return tmp;
}

uint16_t lfsr_state;

void lfsr_init() {
	lfsr_state = adc_read();
}

uint8_t lfsr_read() {
	lfsr_state ^= lfsr_state >> 7;
	lfsr_state ^= lfsr_state << 9;
	lfsr_state ^= lfsr_state >> 13;
	return lfsr_state & 0xFF;
}

void main() {
	// Set PB2 to output
	DDRB = 0b00000100;
	// Set pull-up on PB0 (button)
	PORTB = 0b00000001;
	// Enable interrupts for pin changes on PORTB
	GIMSK = _BV(PCIE1);
	// Enable pin change interrupt on PB0
	PCMSK1 = _BV(PCINT8);

	button_state = 0b00000001;
	button_debounce = 0;
	channels = 0;
	flags.tick = flags.play = 0;

	lfsr_init();

	// Grab the number of samples from the first byte of flash
	usi_init();
	spi_cs_enable();
	spi_transfer(0x03);
	spi_transfer(0x00);
	spi_transfer(0x00);
	spi_transfer(0x00);
	n_samples = spi_transfer(0);
	spi_cs_disable();
	usi_deinit();

	// Enable interrupts
	sei();

	set_sleep_mode(SLEEP_MODE_PWR_DOWN);
	while (1) {
		sleep_mode();

		if (button_debounce > 0) {
			button_debounce--;
		}

		if (flags.tick) {
			// fetch the next sample
			sample_buffer = spi_transfer(0x00);

			if (--sample_count == 0) {
				end_playback();
			}
			flags.tick = 0;
		}

		if (flags.play) {
			spi_cs_disable();

			uint8_t n = lfsr_read() % n_samples;
			begin_playback(n);

			flags.play = 0;
		}
	}
}

ISR(TIM0_OVF_vect) {
	OCR0A = sample_buffer;
	flags.tick = 1;
}

ISR(PCINT1_vect) {
	uint8_t current_state = PINB & 0b00000001;
	if (current_state == button_state) {
		return;
	}

	uint8_t pressed = button_state & ~current_state;
	if (button_debounce == 0 && (pressed & _BV(PB0))) {
		flags.play = 1;
		button_debounce = 4000;
	}

	button_state = current_state;
}