#include <mega128.h>
#include <delay.h>
#define FLASH_PAGESIZE 256             // Flash page 256 bytes(128words) 
#define BAUD 9600                      // Baud Rate 
#define FOSC 8000000L                  // Crystal   
#define CRC_ACK  0x01
#define CRC_NACK 0x07 
#define RE_PROG 0x04
#define REPROG_DONE 0x14

// variables 
char data[FLASH_PAGESIZE];

#pragma warn-

void boot_page_ew(unsigned long p_address, char code) {      // Erase(code=0x03) and write(code=0x05)
    #asm
      ldd r26,y+3  
      out 0x3b,r26                   ; RAMPZ-Bit0 
      ldd r30,y+1                    ; Z register low byte
      ldd r31,y+2                    ; Z register high byte
      ld  r22, y 
      sts 0x68, r22
      spm
    #endasm
}

void boot_page_fill(unsigned int address, int data) {   // write a word to temporary buffer page
     #asm
       ldd r30, y+2                ;r30 = LSB of address
       ldd r31, y+3                ;r31 = MSB of address, address in the Z-register(R31:R30)
       ld r0,   y                  ;LSB of data
       ldd r1,  y+1                ;MSB of data
       ldi r22,0x01
       sts 0x68, r22
       spm
     #endasm
} 

#pragma warn+

void wait_page_rw_ok(void) {           // wait for the end of writing a flash page
     while(SPMCSR & 0x40)             // 0x40 ->RWWSB=1, page erase or page write in progress
     { 
         while(SPMCSR & 0x01);                    
         SPMCSR = 0x11;               // re-enable RWW
         #asm("spm"); 
     }    
}             

void write_one_page(unsigned long address) {        //write one flash page
    unsigned int i;

    boot_page_ew(address,0x03);                    // erase one flash page first
    wait_page_rw_ok();                             // wait for erase finish  
                                    
    for(i=0;i<FLASH_PAGESIZE;i+=2) {               // write data to flash buffer page
       boot_page_fill(i, data[i]+(((int)data[i+1])<<8));  // one word each time
    } 
    boot_page_ew(address,0x05);                    // write a flash page from buffer page 
    wait_page_rw_ok();                             // wait for write finish   
}         

void uart_putchar(char c) {                // send a char through serial port
    while(!(UCSR0A & 0x20));              // 0x20 -> UDRE0
    UDR0 = c; 
}

int uart_getchar(void) {                   // receive a char through serial port
    unsigned char status,res;
    if(!(UCSR0A & 0x80)) return -1;       // no data to be received 
    status = UCSR0A; 
    res = UDR0; 
    if (status & 0x1c) return -1;         // If error, return -1 
    return res; 
} 

char uart_waitchar(void) {                 // wait to receive a valid char
    int c;
    while((c=uart_getchar())==-1); 
    return (char)c; 
}

unsigned int calc_crc(char *ptr, int count) {
    unsigned int crc = 0;
    char i; 
     
    while (--count >= 0) 
    { 
        crc = crc ^ (int) *ptr++ << 8; 
        i = 8; 
        do 
        { 
        if (crc & 0x8000) 
            crc = crc << 1 ^ 0x1021; 
        else 
            crc = crc << 1; 
        } while(--i); 
    }  
 
    return (crc); 
} 
 
void quit(void) {  // exit Bootloader program & start application from 0x0000
    MCUCR = 0x01;
    MCUCR = 0x00;           // move interrupt vector back
    RAMPZ = 0x00;           // RAMPZ reset 
    #asm("jmp 0x0000");     // start application
} 

void main(void) {
    int i = 0;
    unsigned int inCRC;
    unsigned int numPackets, currPacket;
    unsigned long address;


    UBRR0H = 0;
    UBRR0L = 51;                // baudrate 9600 bps 
 
    
    UCSR0B = 0x18;              // Enable Receiver and Transmitter 
 
    uart_putchar(RE_PROG);
    numPackets = ((int)uart_waitchar())<<8;
    numPackets += uart_waitchar();

    address = 0;
    currPacket = 0;
    
    while (currPacket < numPackets) { 
        for(i=0; i<FLASH_PAGESIZE; i++) {    // receive 256 bytes data
            data[i] = uart_waitchar(); 
 
        }
                        
        inCRC = ((int)uart_waitchar())<<8;   // CRC high byte
        inCRC += uart_waitchar();            // CRC low byte 
 
        if (calc_crc(data, 256) == inCRC) {
            write_one_page(address);
            address += FLASH_PAGESIZE;
            currPacket++;
            uart_putchar(CRC_ACK); 
 
        } else {
            uart_putchar(CRC_NACK);
        }
    }
    delay_ms(2);
    uart_putchar(REPROG_DONE);
    quit();
}