#include <avr/io.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include "serial.h"
#include "writepin.h"
#include "train.h"
#include "eeprom.h"
#include "bitrep.h"
#include "ringbuffer.h"
#include "staticvector.h"
#include "turnout.h"

#define COMMAND_BUFFER_SIZE 64
#define SNPRINTF_BUFFER_SIZE 128
#define EPPROM_SIZE 1024

char buffer[SNPRINTF_BUFFER_SIZE];

SVector<Train, 16> trains;
SVector<Turnout, 32> turnouts;

bool autoff = true;
bool powerIsOn = true;

#define TICK_MAX 255
volatile uint8_t tick = 0;

volatile bool resendEvent = false;

constexpr bool USE_PULSES = false;

ISR(TIMER0_OVF_vect) 
{
	++tick;
	if(tick  == 0)
		resendEvent = true;
	if constexpr(USE_PULSES)
	{
		if((tick & 0b00000111) < 1)
			Train::setOutput(Train::HIGH);
		else Train::setOutput(Train::LOW);
	}
}

void timer0InterruptEnable(const bool enable)
{
	if(enable)
		TIMSK0 = 0b00000001;
	else
		TIMSK0 = 0b00000000;
}

void save_state()
{
	cli();
	EEPROM_write_char( 0, trains.count());
	EEPROM_write_char( 1, autoff);
	EEPROM_write_char( 2, trains.maxSize());
	EEPROM_write_char( 3, turnouts.count());
	EEPROM_write_char( 4, turnouts.maxSize());
	for(uint16_t i = 0; i < trains.count(); i++)
	{
		EEPROM_write_char( i*2+32, trains[i].getAddress());
		EEPROM_write_char( i*2+32+1, trains[i].getFunctionMask());
	}
	for(uint16_t i = 0; i < turnouts.count(); i++)
	{
		EEPROM_write_char( i*2+32+trains.maxSize()*2, turnouts[i].getAddress());
		EEPROM_write_char( i*2+32+1+trains.maxSize()*2, turnouts[i].getSubaddress());
	}
	sei();
}

void restore_state()
{
	uint8_t trainCount = EEPROM_read_char(0);
	uint8_t turnoutCount = EEPROM_read_char(3);
	autoff = EEPROM_read_char(1);
	
	trains.clear();
	if(trainCount > trains.maxSize() || trains.maxSize() != EEPROM_read_char(2) ) 
	{
		for(uint16_t i = 0; i < EPPROM_SIZE; i++)
			EEPROM_write_char(i, 0);
	}
	else 
	{
		for(uint8_t i = 0; i < trainCount; i++)
			trains.push_back(Train(EEPROM_read_char(32+i*2), EEPROM_read_char(32+1+i*2)));
		for(uint8_t i = 0; i < turnoutCount; i++)
			turnouts.push_back(Turnout(EEPROM_read_char(32+i*2+trains.maxSize()*2), EEPROM_read_char(32+1+i*2+trains.maxSize()*2)));
	}
}

inline static void printHelp(Serial* serial)
{
	serial->write_p(PSTR("Available Commands: \n\
				help                                     :  Show this prompt.\n\
				train add [address] [functionmask]       :  Add train.\n\
				train [nn] delete                        :  Delete last train.\n\
				train list                               :  Print list of saved trains.\n\
				train [nn] s(top)                        :  Stop nth train.\n\
				train [nn] s(peed) [sp]                  :  Set nth train speed.\n\
				train [nn] function [x]                  :  Toggle x'ed fuction on train n.\n\
				train [nn] r(everse)                     :  Reverse train n.\n\
				turnout add [address] [subaddress]       :  Add a turnout\n\
				turnout list                             :  List turnouts\n\
				turnout set [left/right]                 :  Set turnout direction\n\
				turnout delete [nn]                      :  Delete Turnout\n\
				stop                                     :  stop all trains\n\
				power off                                :  power off the rail\n\
				power on                                 :  power on the rail\n\
				power auto                               :  power off the rail when no trains are moveing\n\
				dump                                     :  prints epprom contence\n\
				erase                                    :  Erase epprom.\n"));
}

void setPower(bool on)
{
	if(powerIsOn != on)
	{
		powerIsOn = on;
		
		if(on)
		{
			Train::setOutput(Train::LOW);
			_delay_ms(100);
			timer0InterruptEnable(true);
		}
		else
		{
			timer0InterruptEnable(false);
			Train::setOutput(Train::OFF);
		}
	}
}

int turnoutDispatch(char* inBuffer, Serial* serial)
{
	if( strcmp(inBuffer, "add") == 0 )
	{
		char* token = strtok(NULL, " ");
		uint8_t address = 0;
		if(token != NULL)
			address = strtol(token, nullptr, 10);
		if(address != 0 && turnouts.remainingCapacity() > 0) 
		{
			uint8_t subaddress = 0;
			
			token = strtok(NULL, " ");
			if(token != NULL) 
				subaddress = strtol(token, nullptr, 10);

			turnouts.push_back(Turnout(address, subaddress));

			uint8_t size = snprintf(buffer, SNPRINTF_BUFFER_SIZE, "TUNROUT saved! NUMBER: %u ADDRESS: %u SUBADDRESS: %u\n", 
									turnouts.count()-1,
									address,
									turnouts.back().getSubaddress());
			serial->write(buffer, size);

			save_state();
			return 0;
		}
		serial->write_p(PSTR("Usage: turnout add [address] [subaddress]"));
	}
	else if( strcmp(inBuffer, "list") == 0 )
	{
		serial->write_p(PSTR("Turnouts:\n"));
		for(uint8_t i = 0; i < turnouts.count(); i++)
		{
			snprintf(buffer, SNPRINTF_BUFFER_SIZE, "NUMBER: %u ADDRESS: %u SUBADDRESS: %u CURRENT PACKET: %x DIRECTION: %u\n", 
						i, turnouts[i].getAddress(), turnouts[i].getSubaddress(),
						turnouts[i].getPacket(), turnouts[i].getDirection());
			serial->write(buffer, SNPRINTF_BUFFER_SIZE);
		}
		serial->putChar('\n');
		return 0;
	}
	else
	{
		uint8_t id = strtol(inBuffer, nullptr, 10);
		if(id < turnouts.count() )
		{
			setPower(true);
			char* token = strtok(NULL, " ");
			if(token != NULL && strcmp(token, "set") == 0 ) 
			{
				char* boolToken = strtok(NULL, " ");
				if(token != NULL && boolToken !=  NULL)
				{
					bool direction = (strcmp(boolToken, "right") == 0);
					turnouts[id].setDirection(direction);
					serial->write_p(PSTR("Set turnout direction "));
					serial->write(direction ? "right\n" : "left\n");
					return 0;
				}
			}
			else if(token != NULL && strcmp(token, "raw") == 0 ) 
			{
				token = strtok(NULL, " ");
				serial->flush();
				if(token != NULL)
				{
					cli();
					uint16_t i = strtol(token, nullptr, 16 );
					snprintf(buffer, SNPRINTF_BUFFER_SIZE, "SENDING: %x to %x\n", i, turnouts[id].getAddress());
					serial->write(buffer, strlen(buffer));
					while(!serial->dataIsWaiting())
					{
						turnouts[id].sendRaw(i);
						sei();
						_delay_ms(20);
						cli();
					}
					serial->write_p(PSTR("Finished\n"));
					serial->flush();
					sei();
					return 0;
				}
			}
			else if(token != NULL && strcmp(token, "delete") == 0)
			{
				turnouts.erase(id);
				serial->write_p(PSTR("Turnout: "));
				serial->write(id);
				serial->write_p(PSTR(" deleted\n"));
				save_state();
				return 0;
			}
			else 
			{
				serial->write_p(PSTR("Not a valid command\n"));
				return -1;
			}
		}
		else
		{
			serial->write_p(PSTR("Id out of range.\n"));
			return -2;
		}
	}
	return -3;
}

int trainDispatch(char* inBuffer, Serial* serial)
{
	if( strcmp(inBuffer, "add") == 0 )
	{
		char* token = strtok(NULL, " ");
		uint8_t address = 0;
		if(token != NULL)
			address = strtol(token, nullptr, 10);
		if(address != 0 && trains.remainingCapacity() > 0) 
		{
			uint8_t functionMask = 0;
			
			token = strtok(NULL, " ");
			if(token != NULL) 
				functionMask = strtol(token, nullptr, 2);

			trains.push_back(Train(address, functionMask));

			uint8_t size = snprintf(buffer, SNPRINTF_BUFFER_SIZE, "TRAIN saved! NUMBER: %u ADDRESS: %u FUNCTIONS: %s FUNCTIONMASK: %s\n", 
									trains.count(),
									address,
									bit_rep[trains.back().getFunctions() & 0x0F],
									bit_rep[trains.back().getFunctionMask() & 0x0F]);
			serial->write(buffer, size);

			save_state();
			return 0;
		}
		serial->write_p(PSTR("Usage: train add [address] [functionmask]"));
	}
	else if( strcmp(inBuffer, "list") == 0 )
	{
		serial->write_p(PSTR("Trains:\n"));
		for(uint8_t i = 0; i < trains.count(); i++)
		{
			snprintf(buffer, SNPRINTF_BUFFER_SIZE, "NUMBER: %u ADDRESS: %u CURRENT PACKET: %x SPEED: %i FUNCTIONS: %s FUNCTIONMASK: %s\n", 
						i, trains[i].getAddress(),
						trains[i].getLastPacket(), trains[i].getSpeed(),
						bit_rep[trains[i].getFunctions() & 0x0F], bit_rep[trains[i].getFunctionMask() & 0x0F]);
			serial->write(buffer, SNPRINTF_BUFFER_SIZE);
		}
		serial->putChar('\n');
		return 0;
	}
	else if( strcmp(inBuffer, "probe") == 0 )
	{
		for(uint16_t j = 0; j < 255; j++)
		{
			snprintf(buffer, SNPRINTF_BUFFER_SIZE, "TRYING ADDR: %x\n", j);
			serial->write(buffer, strlen(buffer));
			cli();
			for(uint8_t k = 0; k < 10; k++)
			{
				Train::sendRawAddr(j, 60);
				_delay_ms(20);
			}
			for(uint8_t k = 0; k < 3; k++)
			{
				Train::sendRawAddr(j, 0);
				_delay_ms(20);
			}
			sei();
		}
		return 0;
	}
	else
	{
		uint8_t id = strtol(inBuffer, nullptr, 10);
		if(id < trains.count() )
		{
			setPower(true);
			char* token = strtok(NULL, " ");
			if( token != NULL &&  (strcmp(token, "speed") == 0 || strcmp(token, "s") == 0) )
			{
				token = strtok(NULL, " ");
				trains[id].setSpeed(atoi(token));
				serial->write_p(PSTR("Set Train speed\n"));
				return 0;
			}
			else if(token != NULL && strcmp(token, "function") == 0 ) 
			{
				token = strtok(NULL, " ");
				char* boolToken = strtok(NULL, " ");
				if(token != NULL && boolToken !=  NULL)
				{
					uint8_t functionId = atoi(token);
					bool on = (strcmp(boolToken, "on") == 0);
					trains[id].sendFunction(functionId, on );
					serial->write_p(PSTR("Set Train function "));
					serial->write(functionId);
					serial->write(on ? " on\n" : " off\n");
					return 0;
				}
			}
			else if(token != NULL && strcmp(token, "probe") == 0 ) 
			{
				serial->flush();
				for(uint16_t j = 0; j < 1024 && !serial->dataIsWaiting(); j++)
				{
					trains[id].sendRaw(j);
					snprintf(buffer, SNPRINTF_BUFFER_SIZE, "TRYING: %x\n", j);
					serial->write(buffer, strlen(buffer));
					sei();
					_delay_ms(250);
					cli();
				}
				sei();
				serial->flush();
				return 0;
			}
			else if(token != NULL && strcmp(token, "raw") == 0 ) 
			{
				token = strtok(NULL, " ");
				if(token != NULL)
				{
					cli();
					uint16_t i = strtol(token, nullptr, 16 );
					snprintf(buffer, SNPRINTF_BUFFER_SIZE, "SENDING: %x to %x\n", i, trains[id].getAddress());
					serial->write(buffer, strlen(buffer));
					for(uint8_t j = 0; j < 100; j++)
					{
						trains[id].sendRaw(i);
						_delay_ms(20);
					}
					serial->write_p(PSTR("Finished\n"));
					sei();
					return 0;
				}
			}
			else if( token != NULL && (strcmp(token, "reverse") == 0 || strcmp(token, "r") == 0) )
			{
				trains[id].reverse();
				serial->write_p(PSTR("Reversed Train\n"));
				return 0;
			}
			else if( token != NULL && (strcmp(token, "stop") == 0 || strcmp(token, "p") == 0))
			{
				trains[id].stop();
				return 0;
			}
			else if( token != NULL && strcmp(token, "delete") == 0)
			{
				trains.erase(id);
				serial->write_p(PSTR("Train: "));
				serial->write(id);
				serial->write_p(PSTR(" deleted\n"));
				save_state();
				return 0;
			}
			else 
			{
				serial->write_p(PSTR("Not a valid command\n"));
				return -1;
			}
		}
		else
		{
			serial->write_p(PSTR("Id out of range.\n"));
			return -2;
		}
	}
	return -3;
}

int powerDispatch(char* token, Serial* serial)
{
	if(strcmp(token, "off") == 0)
	{
		setPower(false);
		return 0;
	}
	else if( strcmp(token, "on") == 0)
	{
		for(uint16_t i = 0; i < trains.count(); i++)
		{
			trains[i].setSpeed(0);
		}
		setPower(true);
		return 0;
	}
	else if(strcmp(token, "auto") == 0)
	{
		token = strtok(NULL, " ");
		if(token != NULL && strcmp(token, "on") == 0)
		{
			autoff = true;
			serial->write_p(PSTR("auto power off turned on.\n"));
			save_state();
		}
		else if(token != NULL && strcmp(token, "off") == 0)
		{
			autoff = false;
			serial->write_p(PSTR("auto power off turned off.\n"));
			save_state();
		}
		else 
		{
			serial->write_p(PSTR("argument must be \"on\" or \"off\". This feature is currently "));
			autoff ? serial->write_p(PSTR("on.\n")) : serial->write_p(PSTR("off.\n"));
		}
		return 0;
	}
	return -1;
}


void serialDispatch(Serial* serial)
{
	if(serial->dataIsWaiting())
	{
		char buffer[COMMAND_BUFFER_SIZE];
		unsigned int length = serial->getString(buffer, COMMAND_BUFFER_SIZE);
		if(length > 0)
		{
			int ret = -1;
			char* token = strtok(buffer, " ");
			if(strcmp(token, "train") == 0 || strcmp(token, "t") == 0 )
			{
				token = strtok(NULL, " ");
				if(token != NULL)
					ret = trainDispatch(token, serial);
			}
			else if(strcmp(token, "turnout") == 0)
			{
				token = strtok(NULL, " ");
				if(token != NULL)
					ret = turnoutDispatch(token, serial);
			}
			else if(strncmp(token, "erase", 4) == 0)
			{
				for(uint16_t i = 0; i < EPPROM_SIZE; i++) EEPROM_write_char(i, 0);
				serial->write_p(PSTR("EEPROM erased\n"));
				trains.clear();
				ret = 0;
			}
			else if(strcmp(token, "dump") == 0)
			{
				for(uint16_t i = 0; i < EPPROM_SIZE; i++) 
				{
					if(i != 0) serial->putChar(',');
					serial->write((uint16_t)EEPROM_read_char(i));
				}
				serial->putChar('\n');
				ret = 0;
			}
			else if((strcmp(token, "stop") == 0 || strcmp(token, "s") == 0 ))
			{
				for(uint16_t i = 0; i < trains.count(); i++)
				{
					cli();
					trains[i].stop();
					sei();
				}
				ret = 0;
			}
			else if(strcmp(token, "power") == 0)
			{
				token = strtok(NULL, " ");
				if(token != NULL)
					ret = powerDispatch(token, serial);
			}
			else if(strcmp(token, "help") == 0)
			{
				printHelp(serial);
				ret = 0;
			}
			else
			{
				serial->putChar('\"');
				serial->write(buffer, length);
				serial->putChar('\"');
				serial->write_p(PSTR(" is not a valid command\n"));
				ret = -1;
			}
			
			if(ret < 0)
			{
				serial->write_p(PSTR("Command Failed\n"));
			}
			else
			{
				for(uint8_t i = 0; i < length; ++i)
				{
					if(buffer[i] == '\0' || buffer[i] == '\n')
						buffer[i] = ' ';
				}
				serial->write_p(PSTR("Sucess \""));
				serial->write(buffer, length);
				serial->write("\"\n");
			}
		}
	}
}

int main() 
{
	TCNT0 = 0;
	TCCR0B = (1<<CS00); // run timer0 with /1 scaler

	DDRD =  (1 << PD2) | (1 << PD3) | (1 << PD4) | (1 << PD5);

	restore_state();

	if(autoff) 
	{
		timer0InterruptEnable(false);
		Train::setOutput(Train::OFF);
	}
	else
	{
		timer0InterruptEnable(true);
		Train::setOutput(Train::LOW);
	}

	sei();

	Serial serial;

	serial.write_p(PSTR("TrainController v0.5 starting\n"));

	uint8_t itemToResend = 0;

	while(true)
	{        
		if(resendEvent && (trains.count() > 0 || turnouts.count())) 
		{
			timer0InterruptEnable(false);
			if(itemToResend < trains.count())
				trains[itemToResend].sendData();
			else if(itemToResend < trains.count() + turnouts.count())
				turnouts[itemToResend-trains.count()].sendData();
			itemToResend++;
			if(trains.count()+turnouts.count() <= itemToResend)
				itemToResend = 0;
			resendEvent = false;
			timer0InterruptEnable(true);
		}
		
		if(autoff) 
		{
			bool trainsRunning = false;
			for(uint16_t i = 0; i < trains.count(); i++)
				trainsRunning = trainsRunning || trains[i].isActive();
			if(!trainsRunning)
			{
				powerIsOn = false;
				timer0InterruptEnable(false);
				Train::setOutput(Train::OFF);
			}
		}
		serialDispatch(&serial);
	}
	return 0;
}