/*****************************************************************************
*
* AVRPROG compatible boot-loader
* Version  : 0.75 (Feb. 2006)
* Compiler : avr-gcc 3.4.1 / avr-libc 1.0.2
* size     : depends on features and startup ( minmal features < 512 words)
* by       : Martin Thomas, Kaiserslautern, Germany
*            eversmith@heizung-thomas.de
*
* License  : Copyright (c) 2005 Martin Thomas
*            Free to use. You have to mention the copyright
*            owners in source-code and documentation of derived
*            work. No warranty.
*
* Additional code and improvements contributed by:
* - Uwe Bonnes
* - Bjoern Riemer
*
*
* Tested with ATmega8, ATmega16, ATmega32, ATmega128, AT90CAN128
*
* - based on the Butterfly Bootloader-Code
*   Copyright (C) 1996-1998 Atmel Corporation
*   Author(s)     : BBrandal, PKastnes, ARodland, LHM
*   The orignal code has been made available by ATMEL together with the
*   Butterfly application code. Since ATMEL.NO had no problem with
*   the application gcc-port they hopefully will not have any concerns about
*   publishing this port. A lot of things have been change but the ATMEL
*   "skeleton" is still in this code. Make sure to keep the copyright notice
*   in derived work to avoid trouble.
*
* - based on boot.h from the avr-libc (c) Eric Weddington
*
****************************************************************************
*
*  The boot interrupt vector is included (this bootloader is completly in
*  ".text" section). If you need this space for further functions you have to
*  add a separate section for the bootloader-functions and add an attribute
*  for this section to _all_ function prototypes of functions in the loader.
*  With this the interrupt vector will be placed at .0000 and the bootloader
*  code (without interrupt vector) at the adress you define in the linker
*  options for the newly created section. See the avr-libc FAQ, the avr-
*  libc's avr/boot.h documentation and the makefile for further details.
*
*  See the makefile for information how to adopt the linker-settings to
*  the selected Boot Size (_Bxxx below)
*
*  With BOOT_SIMPLE this bootloader has 0x3DE bytes size and should fit
*  into a 512word bootloader-section.
*
*  Set AVR clock-frequency and the baudrate below, set MCU-type in
*  makefile.
*
****************************************************************************/
/*
	Does not work reliably so far:
	- lock bits set
*/

// programmers-notepad tabsize 4
#define VERSION_HIGH '0'
#define VERSION_LOW  '7'

/* MCU frequency */
#define F_CPU 7372800

/* UART Baudrate */
#define BAUDRATE 115200

/* use "Double Speed Operation" */
//#define UART_DOUBLESPEED

/* use second UART on mega128 / can128 */
//#define UART_USE_SECOND

#include <stdint.h>
#include <avr/io.h>
#include <avr/wdt.h>
#include <avr/boot.h>
#include <avr/pgmspace.h>

/* enable/disable readout of fuse and lock-bits
  (will not work for Mega169 since not supported by AVRPROG 1.37 */
#define ENABLEREADFUSELOCK

/* enable/disable write of lock-bits
  WARNING: lock-bits can not be reseted by bootloader (as far as I know)
  Only protection no unprotection, "chip erase" from bootloader only
  clears the flash but does no real "chip erase" (this is not possible
  with a bootloader as far as I know)
  Keep this undefined!
 */
// #define WRITELOCKBITS

#include "chipdef.h"

#define UART_RX_BUFFER_SIZE SPM_PAGESIZE
uint8_t gBuffer[UART_RX_BUFFER_SIZE];

#define eeprom_is_ready()             bit_is_clear(EECR, EEWE)
#define my_eeprom_busy_wait()         do{}while(!eeprom_is_ready())

uint32_t address;
uint8_t device;

void sendchar(uint8_t data)
{
	loop_until_bit_is_set(UART_STATUS, UART_TXREADY);
	UART_DATA = data;
}

uint8_t recvchar(void)
{
	loop_until_bit_is_set(UART_STATUS, UART_RXREADY);
	return UART_DATA;
}

uint8_t BufferLoad(uint16_t size, uint8_t mem)
{
	uint16_t data, cnt;
	uint32_t tempaddress;

	for (cnt = 0; cnt < UART_RX_BUFFER_SIZE; cnt++) {
		if (cnt < size)
			gBuffer[cnt] = recvchar();
		else
			gBuffer[cnt] = 0xFF;
	}
	cnt = 0;

	tempaddress = address;					// Store address in page

	my_eeprom_busy_wait();

	if (device == DEVTYPE) {
		// Flash
		if (mem == 'F') {
			do {
				data = gBuffer[cnt++];
				data |= (gBuffer[cnt++] << 8);
				boot_page_fill(address, data);	// call asm routine.

				address = address +2;  		// Select next word in memory
				size -= 2;			// Reduce number of bytes to write by two
			} while(size);				// Loop until all bytes written

			/* commented out since not compatible with mega8 -
			   secondary benefit: saves memory
			tempaddress &= 0xFF80;			// Ensure the address points to the first byte in the page
			*/

			boot_page_write(tempaddress);
			boot_spm_busy_wait();
			boot_rww_enable();			// Re-enable the RWW section

			/* commented out since not compatible with mega8
			if (address != (address & 0xFF80))
			{					// Ensure that the address points to the beginning of the next page
				address &= 0xFF80;
				address += SPM_PAGESIZE;
			}
			*/
		}						// End FLASH

		// Start EEPROM
		if (mem == 'E') {
			address >>= 1;
			do {
				EEARL = address;		// Setup EEPROM address
				EEARH = (address >> 8);
				address++;			// Select next byte
				EEDR = gBuffer[cnt++];

				EECR |= (1<<EEMWE);		// Write data into EEPROM
				EECR |= (1<<EEWE);

				while (EECR & (1<<EEWE));	// Wait for EEPROM write to finish
				size--;				// Decreas number of bytes to write
			} while(size);				// Loop until all bytes written
		}
		return '\r';					// Report programming OK
	}
	return 0;						// Report programming failed
}

void BlockRead(uint16_t size, uint16_t mem)
{
	uint16_t data;

	my_eeprom_busy_wait();

	// Read EEPROM
	if (mem == 'E') {
		// address>>=1; // not needed here - hmm, somehow inconsistant TODO
		do {
			EEARL = address;			// Setup EEPROM address
			EEARH = (address >> 8);
			address++;				// Select next EEPROM byte
			EECR |= (1<<EERE);			// Read EEPROM
			sendchar(EEDR);				// Transmit EEPROM data to PC
			size--;					// Decrease number of bytes to read
		} while (size);					// Repeat until all block has been read
	// Read Flash
	} else {
		do {
#if defined(RAMPZ)
			data = pgm_read_word_far(address);
#else
			data = pgm_read_word_near((uint16_t)address);
#endif
			sendchar((unsigned char)data);		//send LSB
			sendchar((unsigned char)(data >> 8));	//send MSB
			address += 2;				// Select next word in memory
			size -= 2;				// Subtract two bytes from number of bytes to read
		} while (size);					// Repeat until all block has been read
	}
}

uint8_t read_fuse_lock(uint16_t addr, uint8_t mode)
{
	uint8_t retval;

	asm volatile
	(
		"movw r30, %3\n\t" /* Z to addr */ \
		"sts %0, %2\n\t" /* set mode in SPM_REG */ \
		"lpm\n\t"  /* load fuse/lock value into r0 */ \
		"mov %1,r0\n\t" /* save return value */ \
		: "=m" (SPM_REG),
		  "=r" (retval)
		: "r" (mode),
		  "r" (addr)
		: "r30", "r31", "r0"
	);
	return retval;
}

void send_boot(void)
{
	sendchar('A');
	sendchar('V');
	sendchar('R');
	sendchar('B');
	sendchar('O');
	sendchar('O');
	sendchar('T');
}

void (*jump_to_app)(void) = 0x0000;

int main(void)
{
	uint8_t val;

#ifdef START_POWERSAVE
	uint8_t OK = 1;
#endif

	BLDDR  &= ~(1<<BLPNUM);		// set as Input
	BLPORT |= (1<<BLPNUM);		// Enable pullup

	// Set baud rate
	UART_BAUD_HIGH = (UART_CALC_BAUDRATE(BAUDRATE)>>8) & 0xFF;
	UART_BAUD_LOW = (UART_CALC_BAUDRATE(BAUDRATE) & 0xFF);

#ifdef UART_DOUBLESPEED
	UART_STATUS = UART_DOUBLE;
#endif

	UART_CTRL = UART_CTRL_DATA;
	UART_CTRL2 = UART_CTRL2_DATA;

#ifdef START_POWERSAVE
	/*
		This is an adoption of the Butterfly Bootloader startup-sequence.
		It may look a little strange but separating the login-loop from
		the main parser-loop gives a lot a possibilities (timeout, sleep-modes
	    etc.).
	*/
	for(;OK;) {
		if ((BLPIN & (1<<BLPNUM))) {
			// jump to main app if pin is not grounded
			BLPORT &= ~(1<<BLPNUM);	// set to default
			jump_to_app();		// Jump to application sector

		} else {
			val = recchar();
			/* ESC */
			if (val == 0x1B) {
				// AVRPROG connection
				// Wait for signon
				while (val != 'S')
					val = recchar();

				send_boot();			// Report signon
				OK = 0;

			} else {
				sendchar('?');
			}
	        }
		// Power-Save code here
	}

#elif defined(START_SIMPLE)

	if ((BLPIN & (1<<BLPNUM))) {
		// jump to main app if pin is not grounded
		BLPORT &= ~(1<<BLPNUM);		// set to default
		jump_to_app();			// Jump to application sector
	}

#elif defined(START_WAIT)

	// Timer-Setup for ATmega8
	// - verify that the configuration is valid for the target AVR

	#define MY_WAIT 900
	// wait ca 1 sec (900ms)
	TCCR1A = 0; // timer setup
	// F_OSC / 8 / 1000 -> 1ms
	#if (((F_CPU / 8 / 1000)*MY_WAIT) < 65535)
		#warning Information: setting prescaler to 8
		#define WAIT_VALUE ((F_CPU / 8 / 1000)*MY_WAIT)
		TCCR1B |= _BV(CS01);
	#elif ((((F_CPU / 64 / 1000)*MY_WAIT) < 65535))
		#warning Information: setting prescaler to 64
		#define WAIT_VALUE ((F_CPU / 64 / 1000)*MY_WAIT)
		TCCR1B |= _BV(CS01)| _BV(CS00);
	#elif ((((F_CPU / 256 / 1000)*MY_WAIT) < 65535))
		#warning Information: setting prescaler to 256
		#define WAIT_VALUE ((F_CPU / 256 / 1000)*MY_WAIT)
		TCCR1B |= _BV(CS02);
	#else //((((F_CPU / 1024 / 1000)*MY_WAIT) < 65535))
		#warning Information: setting prescaler to 1024
		#define WAIT_VALUE ((F_CPU / 1024 / 1000)*MY_WAIT)
		TCCR1B |= _BV(CS00) |_BV(CS02); //1024 prescaler
	#endif

	while (1) {
		if (UART_STATUS & (1<<UART_RXREADY)) {
			if (UART_DATA == 'S')
				break;
		}
		if (TCNT1 >= WAIT_VALUE){
			BLPORT &= ~(1<<BLPNUM);		// set to default
			TCCR1B = 0;			// timer off
			jump_to_app();			// Jump to application sector
		}
	}
	TCCR1B = 0; // timer off
	send_boot();

#elif defined(START_BOOTICE)
#warning "BOOTICE mode - no startup-condition"

#else
#error "Select START_ condition for bootloader in main.c"
#endif

	for(;;) {
		val = recvchar();
		// Autoincrement?
		if (val == 'a') {
			sendchar('Y');			// Autoincrement is quicker

		//write address
        	} else if (val == 'A') {
			address = recvchar();		//read address 8 MSB
			address = (address<<8) | recvchar();

			address = address<<1;		// !! convert from word address to byte address
			sendchar('\r');

		// Buffer load support
		} else if (val == 'b') {
			sendchar('Y');					// Report buffer load supported
			sendchar((UART_RX_BUFFER_SIZE >> 8) & 0xFF);	// Report buffer size in bytes
			sendchar(UART_RX_BUFFER_SIZE & 0xFF);

		// Start buffer load
		} else if (val == 'B') {
			uint16_t size;
			size = recvchar() << 8;				// Load high byte of buffersize
			size |= recvchar();				// Load low byte of buffersize
			val = recvchar();				// Load memory type ('E' or 'F')
			sendchar(BufferLoad(size, val));		// Start downloading of buffer

		// Block read
		} else if (val == 'g') {
			uint16_t size;
			size = recvchar() << 8;				// Load high byte of buffersize
			size |= recvchar();				// Load low byte of buffersize
			val = recvchar();				// Get memtype
			BlockRead(size, val);				// Perform the block read

		// Chip erase
 		} else if (val == 'e') {
			if (device == DEVTYPE) {
				// erase only main section (bootloader protection)
				uint32_t addr = 0;
				while (APP_END > addr) {
					boot_page_erase(addr);		// Perform page erase
					boot_spm_busy_wait();		// Wait until the memory is erased.
					addr += SPM_PAGESIZE;
				}
			}
			boot_rww_enable();
			sendchar('\r');

		// Exit upgrade
		} else if (val == 'E') {
			wdt_enable(WDTO_15MS); // Enable Watchdog Timer to give reset
			sendchar('\r');

#ifdef WRITELOCKBITS
#warning "Extension 'WriteLockBits' enabled"
		// TODO: does not work reliably
		// write lockbits
		} else if (val == 'l') {
			if (device == DEVTYPE) {
				// write_lock_bits(recchar());
				boot_lock_bits_set(recchar());	// boot.h takes care of mask
				boot_spm_busy_wait();
			}
			sendchar('\r');
#endif
		// Enter programming mode
		} else if (val == 'P') {
			sendchar('\r');

		// Leave programming mode
		} else if (val == 'L') {
			sendchar('\r');

		// return programmer type
		} else if (val == 'p') {
			sendchar('S');		// always serial programmer

#ifdef ENABLEREADFUSELOCK
#warning "Extension 'ReadFuseLock' enabled"
		// read "low" fuse bits
		} else if (val == 'F') {
			sendchar(read_fuse_lock(0x0000, _BV(BLBSET) | _BV(SPMEN)));

		// read lock bits
		} else if (val == 'r') {
			sendchar(read_fuse_lock(0x0001, _BV(BLBSET) | _BV(SPMEN)));

		// read high fuse bits
		} else if (val == 'N') {
			sendchar(read_fuse_lock(0x0003, _BV(BLBSET) | _BV(SPMEN)));

		// read extended fuse bits
		} else if (val == 'Q') {
			sendchar(read_fuse_lock(0x0002, _BV(BLBSET) | _BV(SPMEN)));
#endif

		// Return device type
		} else if (val == 't') {
			sendchar(DEVTYPE);
			sendchar(0);

		// clear and set LED ignored
        	} else if ((val == 'x') || (val == 'y')) {
			recvchar();
			sendchar('\r');

		// set device
		} else if (val == 'T') {
			device = recvchar();
			sendchar('\r');

		// Return software identifier
		} else if (val == 'S') {
			send_boot();

		// Return Software Version
		} else if (val == 'V') {
			sendchar(VERSION_HIGH);
			sendchar(VERSION_LOW);

		// Return Signature Byte
		} else if (val == 's') {
			sendchar(SIG_BYTE1);
			sendchar(SIG_BYTE2);
			sendchar(SIG_BYTE3);

		/* ESC */
		} else if(val != 0x1b) {
			sendchar('?');
		}
	}
	return 0;
}