/* 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> */
/* the spark counter transmits electricity measurements over radio
 * the electricity measurements (voltage, current, power, energy) are query everry second from a peacefair PZEM-004 power meter
 * they are then transmitted using an nRF2L01+ transceiver
 */

#include <stdint.h> // Standard Integer Types
#include <stdio.h> // Standard IO facilities
#include <stdlib.h> // General utilities
#include <stdbool.h> // Boolean
#include <string.h> // Strings

#include <avr/io.h> // AVR device-specific IO definitions
#include <util/delay.h> // Convenience functions for busy-wait delay loops
#include <avr/interrupt.h> // Interrupts
#include <avr/wdt.h> // Watchdog timer handling
#include <avr/pgmspace.h> // Program Space Utilities
#include <avr/sleep.h> // Power Management and Sleep Modes
#include <util/crc16.h> // CRC Computations

#include "main.h" // main definitions
#include "usart.h" // basic USART functions
#include "nrf24.h" // nRF24L01 functions

// AES encryption functions
#include "aes_types.h"
#include "aes128_enc.h"
#include "aes_keyschedule.h"

/* variables */
bool query_pzem004 = false; // flag to query the PZEM-004 power meter
volatile uint8_t timer_seconds = 0; // how many seconds have passed
volatile bool watchdog = true; // set when watchdog interrupt occurred

/* disable watchdog when booting */
void wdt_init(void) __attribute__((naked)) __attribute__((section(".init3")));
void wdt_init(void)
{
	MCUSR = 0;
	wdt_disable();
}

/* enable watchdog with interrupt and system reset (modified wdt_enable) */
#define wdt_set(value) \
__asm__ __volatile__ ( \
    "in __tmp_reg__,__SREG__" "\n\t" \
    "cli" "\n\t" \
    "wdr" "\n\t" \
    "sts %0,%1" "\n\t" \
    "out __SREG__,__tmp_reg__" "\n\t" \
    "sts %0,%2" "\n\t" \
    : /* no outputs */ \
    : "M" (_SFR_MEM_ADDR(_WD_CONTROL_REG)), \
    "r" (_BV(_WD_CHANGE_BIT) | _BV(WDE) | _BV(WDIE)), \
    "r" ((uint8_t) ((value & 0x08 ? _WD_PS3_MASK : 0x00) | _BV(WDE)  | _BV(WDIE) | (value & 0x07)) ) \
    : "r0"  \
)

/* initialize GPIO */
void io_init(void)
{
	/* use UART as terminal */
	usart_init();
	stdout = &usart_output;
	stdin  = &usart_input;
	/* use nRF24L01 */
	nrf24_init();
	/* set up timer 1 to regularly wake up and request data from the power meter */
	TCCR1B |= (0<<WGM13)|(1<<WGM12); // use timer 1 in CTC mode
	TCCR1A |= (0<<WGM11)|(0<<WGM10); // use timer 1 in CTC mode
	TCNT1 = 0; // reset timer 1 counter
	OCR1A = F_CPU/256UL; // the timer to trigger every second with a prescale of 256 (the max is 1.05s)
	TIMSK1 |= (1<<OCIE1A); // enable interrupt
	
	sei(); // enable interrupts
}

int main(void)
{
	_delay_ms(1000); // wait a bit until power meter started
	io_init(); // initialize IOs
	uint8_t nrf24_payload[1+4*4+1] = {0}; // measurement values to be sent: ID byte + encrypted 4x4 float values (voltage, current, power, energy) + 1 CRC8 Mixim/Dallas/iButton/1Wire of the plain text (to know if the key/IV is correct)
	
	//printf(PSTR("welcome to the spar counter power meter monitor\n"));
	
	/* configure nRF24 transceiver */
	nrf24_set_rx_addr(conf.rx_addr); // set device receiving address
	nrf24_set_tx_addr(conf.tx_addr); // set transmission destination address
	nrf24_set_rf_channel(conf.channel); // set transceiver channel
	
	/* PZEM-004 IP address (it won't reply to other requests before it is set) */
	uint8_t cmd[7] = {0xB4,0xC0,0xA8,0x01,0x01,0x00,0x1E}; // set the address 192.168.1.1 (keep this address for the other commands)

	/* start encryption */
	aes128_ctx_t ctx; // the context where the round keys are stored
	aes128_init(conf.key, &ctx); // generating the round keys from the 128 bit key
	// set data to invalid
	memset(&nrf24_payload[1],0xff,sizeof(conf.iv));
	// set ID
	nrf24_payload[0] = conf.id;
	// calculate CRC over values
	nrf24_payload[17] = 0;
	for (uint8_t i = 0; i < sizeof(conf.iv); i++) {
		nrf24_payload[17] = _crc_ibutton_update(nrf24_payload[17], nrf24_payload[1+i]);
	}
	// XOR data with IV for AES CBC mode
	for (uint8_t i=0; i<sizeof(nrf24_payload)-1 && i<sizeof(conf.iv); i++) {
		nrf24_payload[i+1] ^= conf.iv[i];
	}
	aes128_enc(&nrf24_payload[1], &ctx); // encrypt data
	memcpy(conf.iv,&nrf24_payload[1],sizeof(conf.iv)); // save IV back to continue CBC mode
	nrf24_transmit(nrf24_payload,sizeof(nrf24_payload)); // transmit empty packet to sync CBC

	/* start timer to periodically query the power meter */
	timer_seconds = 0; // restart seconds counter
	TCCR1B |= (1<<CS12)|(0<<CS11)|(0<<CS10); // set prescale to 256 (starting the timer)
	query_pzem004 = false; // immediately start querying the power meter

	/* PZEM-004 power meter values */
	float voltage = 0, current = 0, power = 0, energy = 0; // the read values
	bool send_values = false; // set to true to send out the values
	
	bool action = false; // to know if we performed any king of action during which some other activity could have happened

	while (true) { // endless loop for micro-controller
		action = false; // new cycle of actions
		if (watchdog) { // watchdog interrupt trigger. reset watchdog before system reset
			wdt_set(WDTO_2S); // set to 4s (interrupt and reset mode)
			sei(); // re-enable interrupts
			watchdog = false; // wait for next reset
		}
		if (timer_seconds!=0 && timer_seconds==conf.request_period) {
			action = true; // an action is performed
			timer_seconds = 0; // restart timer
			query_pzem004 = true; // remember to query power meter
		}
		if (query_pzem004) {
			action = true; // an action is performed
			query_pzem004 = false; // clear flag (do not wait until previous is finished)
			// query all values (voltage, current, power, energy)
			// start with setting the address, the values will be queried after the power meter answered
			cmd[0] = 0xb4; // set address
			cmd[6] = 0x1e; // update checksum  
			for (uint8_t i=0; i<sizeof(cmd); i++) {
				usart_putchar_nonblocking(cmd[i],NULL);
			}
		}
		if (usart_incoming>=7) { // wait for an answer to be available
			action = true; // an action is performed
			//printf(PSTR("got value\n"));
			// store answer
			uint8_t answer[7];
			for (uint8_t i=0; i<sizeof(answer); i++) {
				answer[i] = usart_getchar(NULL); // save input (this updates usart_incoming)
			}
			// verify checksum
			uint8_t checksum = 0;
			for (uint8_t i=0; i<sizeof(answer)-1; i++) {
				checksum += answer[i];
			}
			if (checksum==answer[sizeof(answer)-1]) { // checksum is correct
				char str[5+1+3+1]; // store the value as string
				switch (answer[0]) {
					case 0xa0: // voltage
						snprintf(str,sizeof(str),"%u.%u",((uint16_t)(answer[1])<<8)+answer[2],answer[3]);
						voltage = atof(str);
						cmd[0] = 0xb1; // query current
						cmd[6] = 0x1b; // update checksum  
						for (uint8_t i=0; i<sizeof(cmd); i++) {
							usart_putchar_nonblocking(cmd[i],NULL);
						}
						break;
					case 0xa1: // current
						snprintf(str,sizeof(str),"%u.%u",((uint16_t)(answer[1])<<8)+answer[2],answer[3]);
						current = atof(str);
						cmd[0] = 0xb2; // query power
						cmd[6] = 0x1c; // update checksum  
						for (uint8_t i=0; i<sizeof(cmd); i++) {
							usart_putchar_nonblocking(cmd[i],NULL);
						}
						break;
					case 0xa2: // power
						snprintf(str,sizeof(str),"%u.%u",((uint16_t)(answer[1])<<8)+answer[2],answer[3]);
						power = atof(str);
						cmd[0] = 0xb3; // query energy
						cmd[6] = 0x1d; // update checksum  
						for (uint8_t i=0; i<sizeof(cmd); i++) {
							usart_putchar_nonblocking(cmd[i],NULL);
						}
						break;
					case 0xa3: // energy
						energy = (((uint32_t)answer[1])<<16)+((uint32_t)(answer[2])<<8)+answer[3];
						send_values = true; // this should be the last of the 4 values we requested. now send them
						break;
					case 0xa4: // address
						cmd[0] = 0xb0; // query voltage
						cmd[6] = 0x1a; // update checksum
						for (uint8_t i=0; i<sizeof(cmd); i++) {
							usart_putchar_nonblocking(cmd[i],NULL);
						}
						break;
				}
			}
		}
		if (send_values) { // send the stored values from the power meter using the nRF24L01
			action = true; // an action is performed
			send_values = false; // clear flag
			memcpy(&nrf24_payload[1+0],&voltage,4); // populate voltage
			memcpy(&nrf24_payload[1+4],&current,4); // populate current
			memcpy(&nrf24_payload[1+8],&power,4); // populate power
			memcpy(&nrf24_payload[1+12],&energy,4); // populate energy
			// set ID
			nrf24_payload[0] = conf.id;
			// calculate CRC over values
			nrf24_payload[17] = 0;
			for (uint8_t i = 0; i < sizeof(conf.iv); i++) {
			nrf24_payload[17] = _crc_ibutton_update(nrf24_payload[17], nrf24_payload[1+i]);
			}
			// XOR data with last data for AES CBC mode
			for (uint8_t i=0; i<sizeof(nrf24_payload)-1 && i<sizeof(conf.iv); i++) {
				nrf24_payload[i+1] ^= conf.iv[i];
			}
			aes128_enc(&nrf24_payload[1], &ctx); // encrypt data
			memcpy(conf.iv,&nrf24_payload[1],sizeof(conf.iv)); // save data back to continue CBC mode
			nrf24_transmit(nrf24_payload,sizeof(nrf24_payload)); // transmit empty packet to sync CBC
		}
		if (nrf24_flag) {
			action = true; // an action was performed 
			nrf24_flag = false; // reset flag
			uint8_t activity = nrf24_activity();
			if (activity&TX_SUCCEEDED) {
				//printf(PSTR("transmission succeeded\n"));
			}
			if (activity&TX_FAILED) {
				//printf(PSTR("transmission failed\n"));
			}
			if (activity&RX_RECEIVED) {
				//printf(PSTR("received data\n"));
			}
			if (activity&RX_AVAILABLE) {
				//printf(PSTR("data available\n"));
				// read all RX FIFO to clear them
				while (nrf24_data_available()) {
					uint8_t data[32];
					nrf24_rx_payload(data,sizeof(data));
					//uint8_t size = nrf24_rx_payload(data,sizeof(data));
					//printf("got %d bytes\n",size);
				}
			}
		}
		/* go to sleep and wait for next interrupt */
		if (!action) { // only go to sleep if no action had to be performed
			set_sleep_mode(SLEEP_MODE_IDLE);
			sleep_mode();
		}
	}
	return 0;
}

/* timer 1 triggered */
ISR(TIMER1_COMPA_vect)
{
	timer_seconds++; // count seconds
}

/* watchdog interrupt */
ISR(WDT_vect)
{
	watchdog = true;
}