/* This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */
/* Copyright (c) 2015 King Kévin <kingkevin@cuvoodoo.info> */
/* This library allows to communicate with a nordic semiconductor nRF24L01 2.4GHz single chip transceiver
 * this library uses the SPI bus and requires interrupts to handle the IRQ
 */
#include <stdint.h> // Standard Integer Types
#include <stdlib.h> // General utilities
#include <stdbool.h> // Boolean
#include <string.h> // Memory utilities

#include <avr/io.h>  // AVR device-specific IO definitions
#include <avr/interrupt.h> // Interrupts

#include <spi.h> // SPI functions
#include <nrf24.h> // nRF24L01 configuration

#include <stdio.h> // Standard IO facilities
#include <avr/pgmspace.h> // Program Space Utilities

#if defined(IRQ_DDR) && defined(IRQ_IO)
/* if the nRF24L01 is initialized with interrupts, watch for nrf_activity to go true
 * once true, check if transmission or reception happened using nrf24_activity(), then nrf24_tx_activity() or nrf24_rx_activity()
 * you can then check is transmission succeeded using nrf24_tx_succeeded()
 * or get the received data using nrf24_rx_data()
 * set it to false after checks are done
 */
volatile bool nrf24_flag = false;
#endif
/* stay in RX mode to receive data */
bool rx = false;

/* initialize the nRF24L01 */
void nrf24_init()
{
	// configure SPI
	spi_init(); // SPI is used to communicate with the nRF24L01
	// configure IO
	CE_DDR |= (1<<CE_IO); // set CE as output
	CE_PORT &= ~(1<<CE_IO); // disable CE (TR/RX)
	// initialise variables
	rx = false;
	// power up device
	// use spi_transfer_blocking in case global interrupt has not been enabled
	uint8_t cmd[2]; // to send commands
	cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
	spi_transfer_blocking(cmd,sizeof(cmd)); // read configuration
#if defined(IRQ_DDR) && defined(IRQ_IO)
	cmd[1] &= ~(0x70); // enable interrupt by clearing MASK_RX_DR, MASK_TX_DS, MASK_MAX_RT
#else 
	cmd[1] |= (0x70); // disable interrupt by setting MASK_RX_DR, MASK_TX_DS, MASK_MAX_RT
#endif
	cmd[1] |= (1<<1); // set PWR_UP to power up
	cmd[1] &= ~(1<<2); // set CRC to 1 byte
	//cmd[1] |= (1<<2); // set CRC to 2 byte
	cmd[0] = 0x20+0x00; // W_REGISTER command + CONFIG register
	spi_transfer_blocking(cmd,sizeof(cmd)); // write configuration
	// flush TX FIFO
	cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
	spi_transfer_blocking(cmd,sizeof(cmd)); // read config
	cmd[0] = 0x20+0x00; // W_REGISTER command + CONFIG register
	cmd[1] &= ~(1<<0); // set PRIM_RX to 0 (PTX)
	spi_transfer_blocking(cmd,sizeof(cmd)); // write config to got in PTX mode
	cmd[0] = 0xe1; // FLUSH_TX
	spi_transfer_blocking(cmd,1); // flush TX FIFO
	// flush RX FIFO
	cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
	spi_transfer_blocking(cmd,sizeof(cmd)); // read config
	cmd[0] = 0x20+0x00; // W_REGISTER command + CONFIG register
	cmd[1] |= (1<<0); // set PRIM_RX to 1 (PRX)
	spi_transfer_blocking(cmd,sizeof(cmd)); // write config to got in PRX mode
	cmd[0] = 0xe2; // FLUSH_RX
	spi_transfer_blocking(cmd,1); // flush RX FIFO
	// clear interrupts
	cmd[0] = 0x20+0x07; // W_REGISTER command + STATUS register
	cmd[1] = 0x70; // clear interrupts
	spi_transfer_blocking(cmd,sizeof(cmd)); // clear interrupts
	// enable auto acknowledge
	cmd[0] = 0x20+0x01; // W_REGISTER command + EN_AA register
	cmd[1] = 0x3f; // enable auto for all channels
	spi_transfer_blocking(cmd,sizeof(cmd)); // write auto acknowledge
	// enable data pipes 0 (to enable acknowledgement) and 1 (to receive)
	cmd[0] = 0x20+0x02; // W_REGISTER command + EN_RXADDR register
	cmd[1] = 0x03; // enable RX pipes 0 and 1
	spi_transfer_blocking(cmd,sizeof(cmd)); // enable pipes
	// enable dynamic payload
	cmd[0] = 0x20+0x1d; // W_REGISTER command + FEATURE register
	cmd[1] = 0x04; // set EN_DPL
	spi_transfer_blocking(cmd,sizeof(cmd)); // enable dynamic payload
	cmd[0] = 0x20+0x1c; // W_REGISTER command + DYNPD register
	cmd[1] = 0x3f; // set DPL_Px
	spi_transfer_blocking(cmd,sizeof(cmd)); // enable dynamic payload on all pipes
	// set address width to 5 bytes
	cmd[0] = 0x20+0x03; // W_REGISTER command + SETUP_AW register
	cmd[1] = 0x03; // set to 5 bytes
	spi_transfer_blocking(cmd,sizeof(cmd)); // set address width
	// set RF: data rate, power
	cmd[0] = 0x00+0x06; // R_REGISTER command + RF_SETUP register
	spi_transfer_blocking(cmd,sizeof(cmd)); // read RF setup
	cmd[0] = 0x20+0x06; // W_REGISTER command + RF_SETUP register
	cmd[1] &= ~(1<<3); // set RF_DR to 1 Mbps
	//cmd[1] |= (1<<3); // set RF_DR to 2 Mbps
	cmd[1] |= (1<<2)|(1<<1); // set RF_PWR to 0 dBm (max)
	spi_transfer_blocking(cmd,sizeof(cmd)); // set RF: data rate, power
	// set re-transmit properties
	cmd[0] = 0x20+0x04; // W_REGISTER command + SETUP_RETR register
	cmd[1] = 0x1f; // set ARD to 500uS and ARC to 15
	spi_transfer_blocking(cmd,sizeof(cmd)); // set re-transmit properties
	// enable interrupt
#if defined(IRQ_IO)
	IRQ_DDR &= ~(1<<IRQ_IO); // set IRQ as input
#if defined(IRQ_INT) && (IRQ_INT==1) && (IRQ_IO==PD2) // INT0 interrupt
	EIMSK |= (1<<INT0); // enable interrupt for IRQ pin
	EICRA &= ~(1<<ISC00); // interrupt on falling edge
	EICRA |= (1<<ISC01); // interrupt on falling edge
#elif defined(IRQ_INT) && (IRQ_INT==1) && (IRQ_IO==PD3) // INT1 interrupt
	EIMSK |= (1<<INT1); // enable interrupt for IRQ pin
	EICRA &= ~(1<<ISC10); // interrupt on falling edge
	EICRA |= (1<<ISC11); // interrupt on falling edge
#else // PCINT interrupt
	PCICR |= (1<<IRQ_PCIE); // enable interrupt for IRQ port
	IRQ_PCMSK |= (1<<IRQ_IO); // enable interrupt for IRQ pin
#endif
	nrf24_flag = false;
#endif
}

/* register names */
const char reg_00[] PROGMEM = "CONFIG";
const char reg_01[] PROGMEM = "EN_AA";
const char reg_02[] PROGMEM = "EN_RXADDR";
const char reg_03[] PROGMEM = "SETUP_AW";
const char reg_04[] PROGMEM = "SETUP_RETR";
const char reg_05[] PROGMEM = "RF_CH";
const char reg_06[] PROGMEM = "RF_SETUP";
const char reg_07[] PROGMEM = "STATUS";
const char reg_08[] PROGMEM = "OBSERVE_TX";
const char reg_09[] PROGMEM = "CD";
const char reg_10[] PROGMEM = "RX_ADDR_P0";
const char reg_11[] PROGMEM = "RX_ADDR_P1";
const char reg_12[] PROGMEM = "RX_ADDR_P2";
const char reg_13[] PROGMEM = "RX_ADDR_P3";
const char reg_14[] PROGMEM = "RX_ADDR_P4";
const char reg_15[] PROGMEM = "RX_ADDR_P5";
const char reg_16[] PROGMEM = "TX_ADDR";
const char reg_17[] PROGMEM = "RX_PW_P0";
const char reg_18[] PROGMEM = "RX_PW_P1";
const char reg_19[] PROGMEM = "RX_PW_P2";
const char reg_20[] PROGMEM = "RX_PW_P3";
const char reg_21[] PROGMEM = "RX_PW_P4";
const char reg_22[] PROGMEM = "RX_PW_P5";
const char reg_23[] PROGMEM = "FIFO_STATUS";
/* register names table */
PGM_P const reg_names[] PROGMEM = {
    reg_00,
    reg_01,
    reg_02,
    reg_03,
    reg_04,
    reg_05,
    reg_06,
    reg_07,
    reg_08,
    reg_09,
    reg_10,
    reg_11,
    reg_12,
    reg_13,
    reg_14,
    reg_15,
    reg_16,
    reg_17,
    reg_18,
    reg_19,
    reg_20,
    reg_21,
    reg_22,
    reg_23
};
/* register size */
const uint8_t reg_sizes[] PROGMEM = {1,1,1,1,1,1,1,1,1,1,5,5,1,1,1,1,5,1,1,1,1,1,1,1};

/* dump config */
void nrf24_dump()
{
	uint8_t cmd[6]; // SPI nFR24L01 command
	uint8_t reg_size; // the register size
	char* str; // to print strings
	for (uint8_t i=0; i<sizeof(reg_sizes); i++) { // go through registers
		cmd[0] = 0x00+i; // R_REGISTER command + register
		reg_size = pgm_read_byte(&(reg_sizes[i]));
		spi_transfer_wait(cmd,1+reg_size); // read register
		str = malloc(strlen_PF((uint_farptr_t)pgm_read_word(&(reg_names[i]))));
		strcpy_PF(str, (uint_farptr_t)pgm_read_word(&(reg_names[i])));
		printf(str);
		free(str);
		printf(": ");
		for (uint8_t j=0; j<reg_size; j++) {
			printf("%02x ",cmd[1+j]);
		}
		printf("\n");
	}
}

/* set the 5 bytes TX address
 * it will be used as the destination address of the data to transmit
 * this address will also be used (in RX_ADDR_P0) to receive the ack
 */
void nrf24_set_tx_addr(uint8_t* addr)
{
	uint8_t cmd[6]; // command
	cmd[0] = 0x20+0x10; // W_REGISTER command + TX_ADDR register
	memcpy(&cmd[1],addr,sizeof(cmd)-1); // copy address
	spi_transfer_wait(cmd,sizeof(cmd)); // write TX_ADDR
	cmd[0] = 0x20+0x0A; // W_REGISTER command + RX_ADDR_P0 register
	memcpy(&cmd[1],addr,sizeof(cmd)-1); // copy address (since it has been overwritten)
	spi_transfer_wait(cmd,sizeof(cmd)); // write RX_ADDR_P0
}

/* set the 5 bytes RX address
 * it will be used as the address of the receiving device
 * this address will be used in RX_ADDR_P1
 */
void nrf24_set_rx_addr(uint8_t* addr)
{
	uint8_t cmd[6]; // command
	cmd[0] = 0x20+0x0B; // W_REGISTER command + RX_ADDR_P1 register
	memcpy(&cmd[1],addr,sizeof(cmd)-1); // copy address
	spi_transfer_wait(cmd,sizeof(cmd)); // write RX_ADDR_P1
}

/* set the RF channel to transmit on */
void nrf24_set_rf_channel(uint8_t channel)
{
	if (channel>125) {
		return;
	}
	/* verify there is space in the FIFO */
	uint8_t cmd[2]; // command
	cmd[0] = 0x20+0x05; // W_REGISTER command + RF_CH register
	cmd[1] = channel; // set channel
	spi_transfer_wait(cmd,sizeof(cmd)); // write RF_CH
}

/* transmit <payload> with of <length> bytes (max 32 bytes)
 * returns false if input is corrupted or FIFO is already full
 * returns true after payload is written
 */
bool nrf24_transmit(uint8_t* payload, uint8_t length)
{
	if (payload==NULL || length==0 || length>32) { // ensure data and size is valid
		return false;
	}
	/* check if TX FIFO is full */
	uint8_t cmd[2]; // command
	cmd[0] = 0x00+0x17; // R_REGISTER command + FIFO_STATUS register
	spi_transfer_wait(cmd,sizeof(cmd)); // read FIFO status
	if (cmd[1]&(1<<5)) { // TX_FULL
		return false;
	}
	/* go into TX mode */
	cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
	spi_transfer_wait(cmd,sizeof(cmd)); // read config
	cmd[0] = 0x20+0x00; // W_REGISTER command + CONFIG register
	cmd[1] &= ~(1<<0); // set PRIM_RX to 0 (PTX)
	spi_transfer_wait(cmd,sizeof(cmd)); // write config to got in PTX mode
	/* write TX payload */
	uint8_t* tx = calloc(length+1,sizeof(uint8_t));
	tx[0] = 0xa0; // W_TX_PAYLOAD command
	memcpy(&tx[1],payload,length); // copy data
	spi_transfer_wait(tx,length+1); // write TX payload
	/* start transmission */
	CE_PORT |= (1<<CE_IO); // enable CE to start transmission
	// clean memory (the payload transfer is complete at this point)
	free(tx);
	tx = NULL;

	return true;
}

/* put device into RX mode or disable it */
void nrf24_enable_rx(bool enable)
{
	rx = enable; // remember to stay in RX mode
	if (rx) {
		uint8_t cmd[2]; // SPI command
		cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
		spi_transfer_wait(cmd,sizeof(cmd)); // read config
		cmd[0] = 0x20+0x00; // W_REGISTER command + CONFIG register
		cmd[1] |= (1<<0); // set PRIM_RX to 1 (PRX)
		spi_transfer_wait(cmd,sizeof(cmd)); // write config to got in PRX mode
		CE_PORT |= (1<<CE_IO); // enable CE to start receiving
	} else {
		uint8_t cmd[2]; // SPI command
		cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
		spi_transfer_wait(cmd,sizeof(cmd)); // read config
		if (cmd[1]&(1<<0)) { // currently in RX mode
			CE_PORT &= ~(1<<CE_IO); // disable CE to go to sleep
		} // do not disable if in TX mode
	}
}

/* handle IRQ interrupt */
#if defined(IRQ_IO)
#if defined(IRQ_INT) && (IRQ_INT==1) && (IRQ_IO==PD2) // INT0 interrupt
ISR(INT0_vect)
#elif defined(IRQ_INT) && (IRQ_INT==1) && (IRQ_IO==PD3) // INT1 interrupt
ISR(INT1_vect)
#else // PCINT interrupt
ISR(IRQ_PCINT_vect)
#endif
{ /* RX or TX succeeded or failed */
	nrf24_flag = true; // warn user
}
#endif

/* check the current status of the device
 * return nrf24_activity_t bits correspoinding to the activies
 * call it after an interrupt
 * puts device back in PRX mode if transmission completed and rx has been enabled
 * puts device back to sleep if transmission completed and rx has not been enabled
 */
uint8_t nrf24_activity()
{
	uint8_t to_return = 0; // is there activity
	uint8_t cmd[2]; // SPI nFR24L01 command
	// get status
	cmd[0] = 0xff; // NOP to read STATUS
	spi_transfer_wait(cmd,1); // read STATUS register
	// clear receive interrupt
	to_return = cmd[0]&(0x70); // copy RX_DR, TS_DS, MAX_RT
	if ((cmd[0]&0x0e)<0x0b) { // payload is available based on RX_P_NO
		to_return |= RX_AVAILABLE; // warn user
	}
	if (cmd[0]&(0x70)) { // activity cause by interrupt
		cmd[0] = 0x20+0x07; // W_REGISTER command + STATUS register
		cmd[1] = 0x70; // clear interrupts
		spi_transfer_wait(cmd,2); // clear interrupt
	}
	if (to_return&TX_FAILED) { // transmission failed based on MAX_RT
		cmd[0] = 0xe1; // FLUSH_TX command
		spi_transfer_wait(cmd,1); // flush failed TX_FIFO
	}
	// verify is there is still data to transmit
	cmd[0] = 0x00+0x17; // R_REGISTER command + FIFO_STATUS register
	spi_transfer_wait(cmd,2); // read FIFO_STATUS register
	if (!(cmd[1]&(1<<4))) { // check TX_EMPTY, if not empty put is TX mode
		cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
		spi_transfer_wait(cmd,2); // read config
		cmd[0] = 0x20+0x00; // W_REGISTER command + CONFIG register
		cmd[1] &= ~(1<<0); // set PRIM_RX to 0 (PTX)
		spi_transfer_wait(cmd,2); // write config to got in PTX mode
		CE_PORT |= (1<<CE_IO); // enable CE to start transmission
	} else if (rx) { // go in RX mode if enables (and TX FIFO is empty)
		cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
		spi_transfer_wait(cmd,2); // read config
		cmd[0] = 0x20+0x00; // W_REGISTER command + CONFIG register
		cmd[1] |= (1<<0); // set PRIM_RX to 1 (PRX)
		spi_transfer_wait(cmd,2); // write config to got in PRX mode
		CE_PORT |= (1<<CE_IO); // enable CE to start receiving
	} else { // nothing to send, and RX is not enabled
		CE_PORT &= ~(1<<CE_IO); // disable CE to go to sleep
	}
	
	return to_return;
}

/* verify if data is available in the RX FIFO */
bool nrf24_data_available()
{
	uint8_t cmd[2]; // SPI nFR24L01 command
	cmd[0] = 0x00+0x17; // R_REGISTER command + FIFO_STATUS register
	spi_transfer_wait(cmd,sizeof(cmd)); // read STATUS register
	return !(cmd[1]&0x01); // not RX FIFO empty
}	

/* writes to received payload in provided <payload> pointer
 * returns size of received data
 * returns 0 if no data is available
 * write maximum <size> in <payload> (max payload size is 32)
 */
uint8_t nrf24_rx_payload(uint8_t* payload, uint8_t size)
{
	uint8_t to_return = 0; // the size of the data
	uint8_t cmd[2]; // SPI nFR24L01 command
	cmd[0] = 0x60; // R_RX_PL_WID command
	spi_transfer_wait(cmd,2); // get size (and STATUS)
	if ((cmd[0]&0x0e)==0x0e) { // RX_FIFO empty based on RX_P_NO
		to_return = 0; // no data is available
	} else if (cmd[1]>32) { // something is wrong
		cmd[0] = 0xe2; // FLUSH_RX command
		spi_transfer_wait(cmd,1); // flush, as suggested by the note
		to_return = 0;
	} else { // data is available
		to_return = cmd[1];
		if (to_return>size) { // avoid buffer overflow
			to_return = size;
		}
		// go into RX mode
		cmd[0] = 0x00+0x00; // R_REGISTER command + CONFIG register
		spi_transfer_wait(cmd,2); // read config
		cmd[0] = 0x20+0x00; // W_REGISTER command + CONFIG register
		cmd[1] |= (1<<0); // set PRIM_RX to 1 (PRX)
		spi_transfer_wait(cmd,2); // write config to got in PRX mode
		// read RX payload
		uint8_t* data = calloc(to_return+1,sizeof(uint8_t)); // re-allocated memory to get data
		data[0] = 0x61; // R_RX_PAYLOAD command
		spi_transfer_wait(data,to_return+1); // read FIFO (automatically deleted)
		memcpy(payload,&data[1],to_return); // copy payload
		free(data); // clean memory
		data = NULL; // clean memory
	}

	return to_return;
}