#include <avr/io.h>
#include <avr/interrupt.h>
#include <util/delay.h>

#if defined(__AVR_ATmega32__)
#define UART_CR UCSRB
#define UART_DR UDR
#elif defined(__AVR_AT90S8515__)
#define UART_CR UCR
#define UART_DR UDR
#elif defined(__AVR_ATmega8__)
#define UART_CR UCSRB
#define UART_DR UDR
#elif defined(__AVR_ATmega48__)
#define UART_CR UCSR0B
#define UART_DR UDR0
#endif

#include "radio.h"
#include "radio_def.h"
#include "radio_addr.h"

#define RADIO_EN_RSSI 1

unsigned int radio_tx_timer=RADIO_TIMER_OFF;
unsigned int radio_rx_timer=RADIO_TIMER_OFF;

unsigned char tx_pri_buffer[RADIO_BUFFER_SIZE+RADIO_OVERHEAD];
unsigned char tx_pri_len=0;
unsigned char tx_pri_pending=0,tx_pri_sending=0;
unsigned char tx_pri_prev_tx=0;

unsigned char tx_buffer[TX_NUM_MESSAGES][RADIO_BUFFER_SIZE+RADIO_OVERHEAD];
unsigned char tx_len[TX_NUM_MESSAGES]={RADIO_OVERHEAD};
unsigned char tx_retx[TX_NUM_MESSAGES]={0};
unsigned char tx_pos=0;
unsigned char tx_head=0, tx_tail=0, tx_msg_count=0;
unsigned char tx_state=RADIO_TX_STATE_IDLE;
unsigned char tx_retx_timeout=0;

unsigned char rx_buffer[RX_NUM_MESSAGES][RADIO_BUFFER_SIZE+RADIO_OVERHEAD];
unsigned char rx_len[RX_NUM_MESSAGES]={0};
signed char rx_ret_val[RX_NUM_MESSAGES]={0};
unsigned char rx_pos=RADIO_LEN_START_SEQ;
unsigned char rx_state=RADIO_RX_CHAR_STATE_0;
unsigned char rx_head=0, rx_exec_tail=0, rx_tail=0;

unsigned char radio_handshake=RADIO_HS_NONE;
unsigned char current_mode=RADIO_MODE_TX;

extern unsigned int led_rate;

void radio_init(void)
{
    //TX/RX enables are active low.
    RADIO_DDR|=1<<RADIO_RX_EN;
    RADIO_DDR|=1<<RADIO_TX_EN;
    radio_switch_mode(RADIO_MODE_RX,0);

    //Setup UART
    #if defined(__AVR_ATmega32__)
    UBRRL=51;//19.2 kbps, no parity, 1 stop bit.
    #elif defined(__AVR_AT90S8515__)
    UBRR=12;//19.2 kbps, no parity, 1 stop bit.
    #elif defined(__AVR_ATmega8__)
    UBRRL=12;//19.2 kbps, no parity, 1 stop bit.
    #elif defined(__AVR_ATmega48__)
    UBRR0L=12;//19.2 kbps, no parity, 1 stop bit.
    #endif
    UART_CR=1<<RXEN0|1<<TXEN0;//RX interrupt enabled. RX/TX enabled.

    //Are we using RSSI triggering?
    #if RADIO_EN_RSSI
    ACSR|=1<<ACIE;//Trigger on toggle.
    DIDR1|=1<<AIN1D|1<<AIN0D;//Disable inputs.
    _delay_us(1);
    if(!(ACSR&ACO))
    {
        UART_CR|=1<<RXCIE0;//Check initial RSSI.
        led_rate=500;
    }
    else led_rate=250;
    #else
    UART_CR|=1<<RXCIE0;
    #endif

    tx_retx_timeout=(RADIO_DEV_ADDR&0x0F) ^ ((RADIO_DEV_ADDR>>4)&0x0F);//Initialize LFSR.
    if(tx_retx_timeout==0)tx_retx_timeout=1;
}

void radio_switch_mode(unsigned char mode,unsigned char pri)
{
    if(mode==RADIO_MODE_TX)
    {
        RADIO_PORT|=RADIO_RX_EN;
        if(current_mode==RADIO_MODE_RX)_delay_us(100);
        RADIO_PORT&=~(RADIO_TX_EN);

        tx_state=RADIO_TX_STATE_SEND;
        radio_handshake=RADIO_HS_NONE;
        //tx_pos=1;
        tx_pos=0;
        if(current_mode==RADIO_MODE_RX)_delay_ms(8);
        //Send first byte to start transmit.
        if(pri)
        {
            //UART_DR=tx_pri_buffer[0];
            UART_DR=0xBB;
            tx_pri_sending=1;
        }
        else UART_DR=0xAA;//UART_DR=tx_buffer[tx_tail][0];
        UART_CR|=1<<TXCIE0;//Enable TXC interrupt.
        UART_CR|=1<<UDRIE0;//Enable DRE interrupt.
    }
    else if(mode==RADIO_MODE_RX)
    {
        if(current_mode==RADIO_MODE_TX)_delay_ms(1);
        RADIO_PORT|=1<<RADIO_TX_EN;
        if(current_mode==RADIO_MODE_TX)_delay_us(100);
        RADIO_PORT&=~(1<<RADIO_RX_EN);
        UART_CR&=~TXCIE0;//Disable TXC interrupt.
        UART_CR&=~UDRIE0;//Disable DRE interrupt.
    }
    current_mode=mode;
}

signed char radio_send(unsigned char dest, unsigned char retransmit, unsigned char len, unsigned char *payload)
{
    return radio_send_msg(dest,retransmit,RADIO_MSG_TYPE_DATA,len,payload);
}

signed char radio_send_msg(unsigned char dest, unsigned char retransmit, unsigned char type, unsigned char len, unsigned char *payload)
{
    unsigned char i;
    unsigned int crc;

    if(tx_msg_count>=TX_NUM_MESSAGES)return RET_TX_BUFFER_FULL;//Is the TX buffer full?
    if(len>RADIO_BUFFER_SIZE)return RET_TX_ERROR;//Payload is too long.

    //Message Parameters
    tx_len[tx_head]=len+RADIO_OVERHEAD;
    tx_retx[tx_head]=retransmit;

    //Message Header
    tx_buffer[tx_head][RADIO_POS_SEQ0]=RADIO_START_BYTE0;
    tx_buffer[tx_head][RADIO_POS_SEQ1]=RADIO_START_BYTE1;
    tx_buffer[tx_head][RADIO_POS_SEQ2]=RADIO_START_BYTE2;
    tx_buffer[tx_head][RADIO_POS_SEQ3]=RADIO_START_BYTE3;
    tx_buffer[tx_head][RADIO_POS_DEST]=dest;
    tx_buffer[tx_head][RADIO_POS_SRC]=RADIO_DEV_ADDR;
    tx_buffer[tx_head][RADIO_POS_TYPE]=type;
    tx_buffer[tx_head][RADIO_POS_LEN]=len;

    //Message Payload
    for(i=0;i<len;i++)tx_buffer[tx_head][(unsigned char)(i+RADIO_POS_DATA)]=payload[i];

    crc=(len==0)?0:crc16(len,payload);
    tx_buffer[tx_head][(unsigned char)(RADIO_LEN_HEADER+len)]=(crc>>8)&0xFF;
    tx_buffer[tx_head][(unsigned char)(RADIO_LEN_HEADER+len+1)]=crc&0xFF;

    RADIO_INC_TX(tx_head);//Update head for next outbound message.

    tx_msg_count++;

    //Is the system idle? Make sure we are not waiting to retransmit.
    if(tx_state==RADIO_TX_STATE_IDLE&&
       radio_tx_timer==RADIO_TIMER_OFF)
    {
        radio_switch_mode(RADIO_MODE_TX,0);
    }

    return RET_TX_SUCCESS;
}

signed char radio_send_pri(unsigned char dest, unsigned char type, unsigned char len, unsigned char *payload)
{
    unsigned char i;
    unsigned int crc;

    if(tx_pri_pending)return RET_TX_BUFFER_FULL;//Only one priority message at a time.
    if(len>RADIO_BUFFER_SIZE)return RET_TX_ERROR;//Payload is too long.

    //Message Parameters
    tx_pri_len=len+RADIO_OVERHEAD;

    //Message Header
    tx_pri_buffer[RADIO_POS_SEQ0]=RADIO_START_BYTE0;
    tx_pri_buffer[RADIO_POS_SEQ1]=RADIO_START_BYTE1;
    tx_pri_buffer[RADIO_POS_SEQ2]=RADIO_START_BYTE2;
    tx_pri_buffer[RADIO_POS_SEQ3]=RADIO_START_BYTE3;
    tx_pri_buffer[RADIO_POS_DEST]=dest;
    tx_pri_buffer[RADIO_POS_SRC]=RADIO_DEV_ADDR;
    tx_pri_buffer[RADIO_POS_TYPE]=type;
    tx_pri_buffer[RADIO_POS_LEN]=len;

    //Message Payload
    for(i=0;i<len;i++)tx_pri_buffer[(unsigned char)(i+RADIO_POS_DATA)]=payload[i];

    crc=(len==0)?0:crc16(len,payload);
    tx_pri_buffer[(unsigned char)(RADIO_LEN_HEADER+len)]=(crc>>8)&0xFF;
    tx_pri_buffer[(unsigned char)(RADIO_LEN_HEADER+len+1)]=crc&0xFF;

    tx_pri_pending=1;
    tx_pri_sending=0;

    //Is the system currently transmitting? If not, start priority transmit.
    if(tx_state!=RADIO_TX_STATE_SEND)
    {
        tx_pri_prev_tx=(tx_state==RADIO_TX_STATE_IDLE)?(radio_tx_timer!=RADIO_TIMER_OFF):1;//Store current transmit/retransmit state.
        radio_tx_timer=RADIO_TIMER_OFF;
        radio_switch_mode(RADIO_MODE_TX,1);
    }

    return RET_TX_SUCCESS;
}

/*void radio_cancel_tx(void)
{
    if(tx_head==tx_tail)return;
    //FIXME Add a cancel pending flag instead of cancelling when in SEND state (will screw with length).
    radio_tx_timer=RADIO_TIMER_OFF;//Cancel timer.
    tx_msg_count--;//Decrement message count.
    RADIO_INC_TX(tx_tail);//Increment tail.
}*/

void radio_tx_exec(void)
{
    switch(tx_state)
    {
    case RADIO_TX_STATE_SEND:
        //Do nothing in this state.
        //The transmit interrupt will switch to the next state
        //when the transmission has finished.
        radio_tx_timer=RADIO_TIMER_OFF;
        break;
    case RADIO_TX_STATE_HS:
        radio_switch_mode(RADIO_MODE_RX,0);
        //ACK received?
        if(radio_handshake==RADIO_HS_ACK)
        {
            tx_state=RADIO_TX_STATE_COMP;
            radio_handshake=RADIO_HS_NONE;
        }
        //NACK received?
        else if(radio_handshake==RADIO_HS_NACK)
        {
            tx_state=RADIO_TX_STATE_FAIL;
            radio_handshake=RADIO_HS_NONE;
        }
        //Timed out?
        else if(radio_tx_timer==0)
        {
            radio_tx_timer=RADIO_TIMER_OFF;
            tx_state=RADIO_TX_STATE_FAIL;
            radio_handshake=RADIO_HS_NONE;
        }
        break;
    case RADIO_TX_STATE_COMP:
        radio_switch_mode(RADIO_MODE_RX,0);
        RADIO_INC_TX(tx_tail);//Go to the next available message.
        tx_msg_count--;//Decrement message count.
        //Is there a priority message pending?
        if(tx_pri_pending)radio_switch_mode(RADIO_MODE_TX,1);//Start message and go to SEND state.
        //Is there another message pending?
        else if(tx_msg_count>0)radio_switch_mode(RADIO_MODE_TX,0);//Start message and go to SEND state.
        else tx_state=RADIO_TX_STATE_IDLE;
        radio_tx_timer=RADIO_TIMER_OFF;//Clear timer.
        break;
    case RADIO_TX_STATE_FAIL:
        //Are we retransmitting?
        tx_state=RADIO_TX_STATE_IDLE;
        if(tx_retx[tx_tail])radio_tx_timer=RADIO_RETX_TIMEOUT;//FIXME next_retx_timeout();
        else
        {
            radio_tx_timer=RADIO_TIMER_OFF;
            tx_msg_count--;//Decrement message count.
            RADIO_INC_TX(tx_tail);
            //Is there another message pending?
            if(tx_msg_count>0)radio_switch_mode(RADIO_MODE_TX,0);//Start message and go to SEND state.
        }

        //Is there a priority message pending?
        if(tx_pri_pending)radio_switch_mode(RADIO_MODE_TX,1);//Start message and go to SEND state.
        break;
    case RADIO_TX_STATE_IDLE:
        radio_switch_mode(RADIO_MODE_RX,0);
        //Is there a priority message pending?
        if(tx_pri_pending)radio_switch_mode(RADIO_MODE_TX,1);//Start message and go to SEND state.
        //Is a retransmit pending/ready?
        if(radio_tx_timer==0)radio_switch_mode(RADIO_MODE_TX,0);//Start message and go to SEND state.
        break;
    default:
        tx_state=RADIO_TX_STATE_IDLE;
        break;
    }
}

signed char radio_receive(unsigned char *src,unsigned char *len,unsigned char *payload)
{
    unsigned char old_tail;
    if(rx_tail==rx_exec_tail)return RET_RX_QUEUE_EMPTY;//No messages available

    //Is the message verified?
    if(rx_ret_val[rx_tail]==RET_RX_PACKET_REC)
    {
        for(*len=0;*len<rx_len[rx_tail];(*len)++)payload[*len]=rx_buffer[rx_tail][(unsigned char)(*len+RADIO_POS_DATA)];
        *src=rx_buffer[rx_tail][RADIO_POS_SRC];
    }

    old_tail=rx_tail;
    RADIO_INC_RX(rx_tail);
    return rx_ret_val[old_tail];
}

void radio_rx_exec(void)
{
    unsigned int crc_calc,crc_rec;
    unsigned char broadcast=0;

    //Check RX timer for data timeout.
    if(rx_state!=RADIO_RX_CHAR_STATE_0/*radio_rx_timer!=RADIO_TIMER_OFF*/&&radio_rx_timer==0)
    {
        rx_ret_val[rx_head]=RET_RX_LEN_ERROR;//Flag as length error.
        rx_len[rx_head]=rx_pos;//Set received message length.
        //radio_rx_timer=RADIO_TIMER_OFF;//Stop timer.
        //Prepare for next message.
        rx_state=RADIO_RX_CHAR_STATE_0;
        RADIO_INC_RX(rx_head);
        rx_pos=RADIO_LEN_START_SEQ;
    }

    if(rx_head==rx_exec_tail)return;//No messages available

    crc_calc=crc16(rx_len[rx_exec_tail],&rx_buffer[rx_exec_tail][RADIO_POS_DATA]);
    crc_rec=((((unsigned int)(rx_buffer[rx_exec_tail][(unsigned char)(RADIO_POS_DATA+rx_len[rx_exec_tail])]))<<8)|
            rx_buffer[rx_exec_tail][(unsigned char)(RADIO_POS_DATA+rx_len[rx_exec_tail]+1)]);
    broadcast=(rx_buffer[rx_exec_tail][RADIO_POS_DEST]==RADIO_ADDR_BROADCAST);

    //Was the packet sent to NULL?
    if(rx_buffer[rx_exec_tail][RADIO_POS_DEST]==RADIO_ADDR_NULL)
    {
        rx_ret_val[rx_exec_tail]=RET_RX_QUEUE_EMPTY;
    }
    //Does the destination match?
    else if(rx_buffer[rx_exec_tail][RADIO_POS_DEST]!=RADIO_DEV_ADDR&&!broadcast)
    {
        rx_ret_val[rx_exec_tail]=RET_RX_DEST_MISMATCH;
    }
    //Does the length match?
    else if(rx_ret_val[rx_exec_tail]==RET_RX_LEN_ERROR)
    {
        #if !(defined(__AVR_ATmega8__)||defined(__AVR_ATmega48__))
        PORTA^=1<<2;
        #endif
        //Did we get a source to send a NACK to?
        if(rx_len[rx_exec_tail]>RADIO_POS_SRC&&!broadcast)radio_send_hs(rx_buffer[rx_exec_tail][RADIO_POS_SRC],RADIO_MSG_TYPE_NACK);
    }
    //Does the CRC match?
    else if(crc_calc!=crc_rec)
    {
        #if !(defined(__AVR_ATmega8__)||defined(__AVR_ATmega48__))
        PORTA^=1<<1;
        #endif
        if(!broadcast)radio_send_hs(rx_buffer[rx_exec_tail][RADIO_POS_SRC],RADIO_MSG_TYPE_NACK);
        rx_ret_val[rx_exec_tail]=RET_RX_CRC_ERROR;
    }
    //Good message. Check for handshake, otherwise
    //return message to user.
    else if(rx_buffer[rx_tail][RADIO_POS_TYPE]==RADIO_MSG_TYPE_ACK||
            rx_buffer[rx_tail][RADIO_POS_TYPE]==RADIO_MSG_TYPE_NACK)
    {
        radio_handshake=(rx_buffer[rx_exec_tail][RADIO_POS_TYPE]==RADIO_MSG_TYPE_ACK)?RADIO_HS_ACK:RADIO_HS_NACK;
        rx_ret_val[rx_exec_tail]=RET_RX_QUEUE_EMPTY;
    }
    else
    {
        rx_ret_val[rx_exec_tail]=RET_RX_PACKET_REC;
        if(!broadcast)radio_send_hs(rx_buffer[rx_exec_tail][RADIO_POS_SRC],RADIO_MSG_TYPE_ACK);
        #if !(defined(__AVR_ATmega8__)||defined(__AVR_ATmega48__))
        PORTA^=1<<3;
        #endif
    }
    RADIO_INC_RX(rx_exec_tail);
}

inline void radio_exec(void){radio_tx_exec();radio_rx_exec();}

unsigned int crc16(unsigned char len, unsigned char *payload)
{
    unsigned int crc;
    unsigned char data[2];
    unsigned int uLen;
    uLen=len;
    crc = 0;
    data[0] = 0;
    while (uLen-- ){
        data[1] = data[0];
        data[0] = *payload++;
        if(crc & 0x8000){
            crc = (crc & 0x7fff) << 1;
            crc ^= 0x8005;
        }
        else {
            crc <<= 1;
        }
        crc ^= (((int)data[1])<<8 | data[0]);
    }
    return crc;
}

/*unsigned char next_retx_timeout(void)
{
    unsigned char ret,shift_in;
    ret=tx_retx_timeout;
    shift_in=((tx_retx_timeout>>2)&0x01)^((tx_retx_timeout>>3)&0x01);
    tx_retx_timeout=(tx_retx_timeout<<1)&0x0F|shift_in;
    return ret;
}*/

void radio_send_hs(unsigned char src,unsigned char type)
{
    radio_send_pri(src,type,0,0);
}

/*void radio_send_hs_debug(unsigned char src,unsigned char type,unsigned char len,unsigned char *payload)
{
    radio_send_msg(src,0,type,len,payload);
}*/

#if defined(__AVR_ATmega32__)||defined(__AVR_ATmega8__)||defined(__AVR_ATmega48__)
ISR(USART_UDRE_vect)
#elif defined(__AVR_AT90S8515__)
ISR(UART_UDRE_vect)
#endif
{
    if(tx_pri_sending)
    {
        if(tx_pos<tx_pri_len)UART_DR=tx_pri_buffer[tx_pos++];//Send next byte.

        //Are we done sending?
        if(tx_pos==tx_pri_len)
        {
            UART_CR&=~UDRIE0;//Disable DRE interrupt.
        }
    }
    else
    {
        if(tx_pos<tx_len[tx_tail])UART_DR=tx_buffer[tx_tail][tx_pos++];//Send next byte.

        //Are we done sending?
        if(tx_pos==tx_len[tx_tail])
        {
            UART_CR&=~UDRIE0;//Disable DRE interrupt.
        }
    }
}

#if defined(__AVR_ATmega32__)||defined(__AVR_ATmega8__)||defined(__AVR_ATmega48__)
ISR(USART_TX_vect)
#elif defined(__AVR_AT90S8515__)
ISR(UART_TX__vect)
#endif
{
    //Was this a priority message?
    if(tx_pri_sending)
    {
        tx_pri_sending=0;
        tx_pri_pending=0;
        //Were we transmitting earlier?
        if(tx_pri_prev_tx)
        {
            radio_switch_mode(RADIO_MODE_TX,0);//Re-send message.
        }
        //Are there messages pending transmission?
        else if(tx_msg_count>0)
        {
            //RADIO_INC_TX(tx_tail);//Go to the next available message.
            radio_switch_mode(RADIO_MODE_TX,0);//Send message.
        }
        else tx_state=RADIO_TX_STATE_IDLE;
    }
    //Was this a handshake packet?
    else if(tx_buffer[tx_tail][RADIO_POS_TYPE]==RADIO_MSG_TYPE_ACK||
       tx_buffer[tx_tail][RADIO_POS_TYPE]==RADIO_MSG_TYPE_NACK)
    {
        tx_state=RADIO_TX_STATE_COMP;//Enter message complete state.
    }
    else
    {
        if(tx_buffer[tx_tail][RADIO_POS_DEST]!=RADIO_ADDR_BROADCAST)
        {
            tx_state=RADIO_TX_STATE_HS;//Enter handshake state.
            radio_tx_timer=RADIO_HS_TIMEOUT;//Set handshake timer.
        }
        else tx_state=RADIO_TX_STATE_COMP;
    }
}

#if defined(__AVR_ATmega32__)||defined(__AVR_ATmega8__)||defined(__AVR_ATmega48__)
ISR(USART_RX_vect)
#elif defined(__AVR_AT90S8515__)
ISR(UART_RX__vect)
#endif
{
    unsigned char c;
    c=UART_DR;

    switch(rx_state)
    {
    case RADIO_RX_CHAR_STATE_0:
        if(c==RADIO_START_BYTE0)rx_state=RADIO_RX_CHAR_STATE_1;
        break;
    case RADIO_RX_CHAR_STATE_1:
        if(c==RADIO_START_BYTE1)rx_state=RADIO_RX_CHAR_STATE_2;
        else rx_state=RADIO_RX_CHAR_STATE_0;
        break;
    case RADIO_RX_CHAR_STATE_2:
        if(c==RADIO_START_BYTE2)rx_state=RADIO_RX_CHAR_STATE_3;
        else rx_state=RADIO_RX_CHAR_STATE_0;
        break;
    case RADIO_RX_CHAR_STATE_3:
        if(c==RADIO_START_BYTE3)rx_state=RADIO_RX_CHAR_STATE_OK;
        else rx_state=RADIO_RX_CHAR_STATE_0;
        break;
    case RADIO_RX_CHAR_STATE_OK:
        if(rx_pos==RADIO_POS_LEN)
        {
            rx_len[rx_head]=c;
            //Is the length invalid? This could happen if noise occurs after start sequence.
            if(c>RADIO_BUFFER_SIZE)
            {
                rx_state=RADIO_RX_CHAR_STATE_0;
                rx_pos=RADIO_LEN_START_SEQ;
                //radio_rx_timer=RADIO_TIMER_OFF;
                break;
            }
        }
        rx_buffer[rx_head][(unsigned char)(rx_pos++)]=c;

        //Is the message complete?
        if(rx_pos>RADIO_POS_LEN&&rx_pos==rx_len[rx_head]+RADIO_POS_LEN+RADIO_LEN_CRC+1)
        {
            rx_state=RADIO_RX_CHAR_STATE_0;
            RADIO_INC_RX(rx_head);
            rx_pos=RADIO_LEN_START_SEQ;
            //radio_rx_timer=RADIO_TIMER_OFF;//Stop timer.
        }
        break;
    default:
        rx_state=RADIO_RX_CHAR_STATE_0;
        break;
    }
    radio_rx_timer=RADIO_CHAR_TIMEOUT;//Start/reset character timeout.
}

ISR(ANALOG_COMP_vect)
{
    #if RADIO_EN_RSSI
    if(!(ACSR&ACO))UART_CR|=RXCIE0;
    else UART_CR&=~RXCIE0;
    led_rate=(UART_CR&RXCIE0)?500:250;
    #endif
}