//
// Example of firmware update via UART
// By PEK '2004
//
// Compiler: GCC
// MCU: ATtiny2313
// Clock: 8 MHz
//
// Include the following line to the linker, enabling the new
// memory segment at address 0x0200:
// -Wl,--section-start=.memsection=0x0200
//
// Remember to set the SELFPRGEN fuse bit, to enable self
// programming.
//
//
// Firmware data to the MCU consists of the following parts,
// where <LAST> is the first byte received:
//
// Data:  <LAST><PAGE><DATA><CHECKSUM>
// Bytes:  -1-   -1-   -32-    -1-
//
// LAST: 0x00 = not the last message
//       0x01 = last message in update sequence
// PAGE: Page in memory = 0x00-0x3F (tiny2313)
// DATA: Data to write
// CHECKSUM: Checksum of message = 2 complement of all bytes (except the checksum)
//
//
// Answer to firmware data from the MCU:
//
// Data:  <STATUS><CHECKSUM>
// Bytes:   -1-      -1-
//
// STATUS: 0x00 = wrong checksum, send again
//         0x01 = no errors, ready for next message of firmware data
// CHECKSUM: Checksum of message = 2 complement of the <STATUS> byte
//
//
// Example of programming sequence:
// 1. Send 0x55 to enable "Firmware update mode"
// 2. Wait for message <0x01><0xFF>, i.e. the MCU is ready
// 3. Send <0x00><0x00><32 bytes><checksum> to program page 0x00
// 4. Wait for message <0x01><0xFF>, i.e. the MCU is ready for next page to program
// 5. Send <0x00><0x01><32 bytes><checksum> to program page 0x01
// 6. Wait for message <0x01><0xFF>, i.e. the MCU is ready for next page to program
// 7. Send <0x01><0x02><32 bytes><checksum> to program page 0x02 and quit
// 8. Wait for message <0x01><0xFF>, everything went OK
//
//
// Notes:
// If you receive the message of checksum error in the last programming
// message (the last page to download), the program will still quit from
// programming mode. To try reprogram the last page you need to enter programming
// mode again and program the last page. To go around this you can always send
// <0x00> as the first byte in the programming message. When you have finished
// all programming you send 0x01 as the first byte and be assured of that the
// checksum is wrong. Then the program will quit from programming mode.
//


#include <avr/io.h>
#include <avr/interrupt.h>
#include <avr/signal.h>

#ifndef BYTE
#define BYTE uint8_t
#endif

#ifndef WORD
#define WORD uint16_t
#endif

// Problems with COFF generation to simulate other memory sections in
// AVRStudio. Due to this problem the memory section below is not
// used during simulation, to have all code in the same segment.

// Use this row during binary generation for the circuit 
#define MEM_SECTION __attribute__ ((section (".memsection")))

// Use this row during simulation with AVRStudio
//#define MEM_SECTION

// Message for firmware update
#define UART_UPDATE_FIRMWARE 0x55

#define MEMORY_PAGE_SIZE 32	// Size of a memory page in tiny2313

#define FIRMWARE_MSG_SIZE MEMORY_PAGE_SIZE + 3
#define FIRMWARE_LASTMSG_BYTE 0
#define FIRMWARE_PAGE_BYTE 1
#define FIRMWARE_DATA_BYTES 2 
#define FIRMWARE_CHECKSUM_BYTE FIRMWARE_MSG_SIZE - 1


// Declaration of functions
void page_erase(BYTE page) MEM_SECTION;
void page_fill(BYTE address, uint16_t data) MEM_SECTION;
void page_write(BYTE page) MEM_SECTION;
void update_firmware(void) MEM_SECTION;
void send_byte(BYTE data) MEM_SECTION;
BYTE calc_checksum(BYTE* data, BYTE size) MEM_SECTION;


// ---------------------------------
// Start of firmware
// ---------------------------------

SIGNAL(SIG_USART0_RX)
{
	BYTE data;
	
	data = UDR;	// Store received byte

	if(data == UART_UPDATE_FIRMWARE)
	{
		PORTB = 0x01;	// Enter programming mode
		update_firmware();
		PORTB = 0x00;	// Programming finnished
	}
	
	PORTB = data;
}

int main(void)
{
	// Initiate UART
	UBRRH = 0;	// 19200 bps at 8 MHz
	UBRRL = 25;
	
	UCSRB = _BV(RXEN) | _BV(TXEN) | _BV(RXCIE);	// Rx Complete Interrupt & Enable Rx/Tx
	UCSRC = _BV(UCSZ1) | _BV(UCSZ0);	// 8N1

	// Initiate Ports
	DDRB = 0xFF;	// Port B as output
	PORTB = 0;

	sei();	// Enable global interrupts
	
	while(1);	// Loop forever
	
	return 1;
}


// ---------------------------------
// Start of updating software
// ---------------------------------

void update_firmware(void)
{	
	BYTE i;
	BYTE data[FIRMWARE_MSG_SIZE];
	BYTE* dataPntr;
	
	// Ready for new message
	send_byte(0x01);
	send_byte((~0x01)+1);	// Checksum
	
	do
	{
		for(i = 0; i < FIRMWARE_MSG_SIZE; i++)
		{
			// Wait for byte
			while (!(UCSRA & _BV(RXC)));
		
			data[i] = UDR;
		}

		if(calc_checksum(data, FIRMWARE_MSG_SIZE - 1) == data[FIRMWARE_CHECKSUM_BYTE])
		{
			// Correct checksum, erase and write
		
			dataPntr = &(data[FIRMWARE_DATA_BYTES]);
		
			// Erase page
			page_erase(data[FIRMWARE_PAGE_BYTE]);
	
			for(i = 0; i < MEMORY_PAGE_SIZE; i+=2)
			{
				// Fill temporary buffer with data
				page_fill(i, (WORD)((dataPntr[i+1] << 8) | dataPntr[i]));
		
				// Write to page
				page_write(data[FIRMWARE_PAGE_BYTE]);
			}
			
			// Ready for new message
			send_byte(0x01);
			send_byte((~0x01)+1);	// Checksum
		}
		else
		{
			// Wrong checksum, send again
			send_byte(0x00);
			send_byte((~0x00)+1);	// Checksum
		}
	} while(!data[FIRMWARE_LASTMSG_BYTE]);
	
	// Wait for restart
	//while(1);
}

// Erase page (64 pages in ATtiny2313)
void page_erase(BYTE page)
{
	uint16_t address = (WORD)page << 5;
	
	asm volatile
	(
		"movw r30, %2\n\t"
		"sts %0, %1\n\t"
		"spm\n\t"
		: "=m" (SPMCSR)
		: "r" ((BYTE)0x03), "r" ((WORD)address)
		: "r30", "r31"
	);
}

// Fill temporary buffer with data (16 words / page in ATtiny2313)
void page_fill(BYTE address, WORD data)
{
	asm volatile
	(
		"movw  r0, %3\n\t"
		"movw r30, %2\n\t"
		"sts %0, %1\n\t"
		"spm\n\t"
		"clr  r1\n\t"
		: "=m" (SPMCSR)
		: "r" ((BYTE)0x01), "r" ((WORD)address), "r" ((WORD)data)
		: "r0", "r30", "r31"
	);
}

// Write word in temporary buffer to Flash
void page_write(BYTE page)
{
	WORD address = (WORD)page << 5;
	
	asm volatile
	(
		"movw r30, %2\n\t"
		"sts %0, %1\n\t"
		"spm\n\t"
		: "=m" (SPMCSR)
		: "r" ((BYTE)0x05), "r" ((WORD)address)
		: "r30", "r31"
	);
}

// Calculate checksum
BYTE calc_checksum(BYTE* data, BYTE size)
{
	BYTE checksum = 0;
	
	while(size)
	{
		checksum += data[--size];
	}

	checksum = (~checksum) + 1;	// 2 complement

	return checksum;
}

// Send byte
void send_byte(BYTE data)
{
	while(!(UCSRA & _BV(UDRE)));	// Wait for empty transmit buffer
	
	UDR = data;	// Send byte
}