
#include <string>
#include <iostream>
#include <map>
#include <array>
#include <thread>
#include <chrono>

#include "intelhex.h"
#include "serial/serial.h"


#define ACK			0x79
#define NACK		0x1F
#define CMD_ERASE	0x43
#define CMD_GETID	0x02
#define CMD_WRITE	0x2b


using std::cout, std::endl, std::cerr;


typedef uint8_t u8;
typedef uint16_t u16;
typedef uint32_t u32;
typedef uint64_t u64;


#define APPNAME "C:/Users/gritt/Documents/workspace/STM32/Custom-Bootloader/App1/build/App1.hex"
#define COMPORT	"COM3"

std::map <uint16_t, std::string> deviceIDs = {
	{0x410, "STM32F103RB"},
	{0x415, "STM32L152RG"},
	{0x417, "STM32L053R8"},
	{0x421,  "STM32F446RE"},
	{0x431, "STM32F411RE"},
	{0x433, "STM32F401RE"},
	{0x437, "STM32L152RE"},
	{0x439, "STM32F302R8"},
	{0x438, "STM32F334R8"},
	{0x440, "STM32F030R8"},
	{0x442, "STM32F091RC"},
	{0x446, "STM32F303RE"},
	{0x447, "STM32L073RZ"},
	{0x448, "STM32F070RB/STM32F072RB"},
	{0x458, "STM32F410RB"},
};


u32 computeStm32Crc(u32* buffer, u32 len) {
	u32 crc = 0xffffffff;

	for (u32 i = 0; i < len; ++i) {
		crc ^= buffer[i];

		for (u8 j = 0; j < 32; ++j) {

			if (crc & 0x80000000) {
				crc = (crc << 1) ^ 0x04C11DB7;  // Polynomial used in STM32
			}
			else {
				crc = (crc << 1);
			}
		}
	}
	return (crc & 0xffffffff);
}

u32 computeStm32Crc_bytes(u8* buffer, u32 len) {
	u32 crc = 0xffffffff;

	for (u32 i = 0; i < len; ++i) {
		u32 data = buffer[i];
		crc ^= data;

		for (u8 j = 0; j < 32; ++j) {

			if (crc & 0x80000000) {
				crc = (crc << 1) ^ 0x04C11DB7;  // Polynomial used in STM32
			}
			else {
				crc = (crc << 1);
			}
		}
	}
	return (crc & 0xffffffff);
}



class Flasher {

public:

	Flasher(std::string com) :
	
		timeout (0, 1000, 1, 1000, 1)
	{

		stm32Serial.setPort(com);
		stm32Serial.setBaudrate(115200);
		stm32Serial.setTimeout(timeout);
		try {
			stm32Serial.open();

			u32 len = stm32Serial.available();
			if (len != 0) {
				stm32Serial.read(len);
			}
		}

		catch (...) {
			cerr << "Unable to open serial port" << endl;
		}

	}

	~Flasher() {
		stm32Serial.close();
	}

	u32 getDeviceID() {
		if (!stm32Serial.isOpen()) {
			cerr << "Unable to open serial " << COMPORT << endl;
			return 0;
		}

		std::array <u8, 20> payload;
		std::fill(std::begin(payload), std::end(payload), 0);

		payload[0] = CMD_GETID;
		u32 crc = computeStm32Crc((u32*)payload.data(), 1);
		memcpy(payload.data() + 1, &crc, sizeof(u32));

		stm32Serial.write(payload.data(), 20);

		// read 1 byte (ACK or NACK)
		u32 len = stm32Serial.read(payload.data(), 1);
		if (len != 1) {
			cerr << "Error reading device ID\n" << endl;
			return 0;
		}
		if (payload[0] == ACK) {
			len = stm32Serial.read(payload.data(), 2);
			if (len == 2) {

				u32 deviceID = 0;
				memcpy(&deviceID, payload.data(), 2);
				return deviceID;
			}
		}
		else {
			cerr << "Received NACK on device ID request" << endl;
			return 0;
		}
	}

	u8 eraseFlash () {

		if (!stm32Serial.isOpen())
			return 1;

		std::array <u8, 20> payload;
		payload[0] = CMD_ERASE;
		payload[1] = 0xff;
		u32 crc = computeStm32Crc_bytes(payload.data(), 2);
		memcpy(payload.data()+2, &crc, sizeof(u32));



		stm32Serial.write(payload.data(), 20);

		// read 1 byte (ACK or NACK)
		// since erase operation can take several seconds, increase read timeout to 50s
		timeout.read_timeout_constant = 50000;
		stm32Serial.setTimeout(timeout);

		u32 len = stm32Serial.read(payload.data(), 1);

		if (len == 0) {
			cerr << "read timed out" << endl;
			return 2;
		}
		else {
			if (payload[0] == ACK) {
				cout << "received ACK" << endl;
				return 0;
			}
			else {
				cerr << "received NACK" << endl;
				return 3;
			}
		}
	}

	void flashImage (std::string image) {

		intelhex::hex_data program = intelhex::hex_data(APPNAME);
		std::array <u8, 20> requestPayload;
		std::array <u8, 20> dataPayload;
		std::array <u8, 20> readBuffer;

		requestPayload[0] = CMD_WRITE;
		
		u32 line = 1;
		for (auto& record : program) {

			cout << std::dec << (int)line++ << ":Write at " << std::hex << (int) record.first << endl;

			// send write request and address
			memcpy(requestPayload.data()+1, &record.first,sizeof (u32));
			u32 crc = computeStm32Crc_bytes(requestPayload.data(), 5);
			memcpy(requestPayload.data()+5, &crc, sizeof(u32));

			stm32Serial.write(requestPayload.data(), 20);

			u8 len = stm32Serial.read(readBuffer.data(), 1);
			if (len == 0 || len==1 && readBuffer[0] == NACK)
				return;

			// send data
			for (u8 i = 0; i < 16; i++) {
				if (i < record.second.size())
					dataPayload[i] = record.second[i];
				else
					dataPayload[i] = 0xff;
			}

			crc = computeStm32Crc((u32*)dataPayload.data(), 4);
			memcpy(dataPayload.data()+16, &crc, sizeof(u32));

			stm32Serial.write(dataPayload.data(), 20);

			len = stm32Serial.read(readBuffer.data(), 1);
			if (len == 0) {
				cerr << "TIMEOUT" << endl;
				return;
			}
			if (len == 1 && readBuffer[0] == NACK) {
				cerr << "NACK" << endl;
				return;
			}

			cout << "Done" << endl;
		}
	}

private:

	// open serial port
	serial::Timeout timeout;
	serial::Serial stm32Serial;
	

};


int main() {


	Flasher flasher(COMPORT);
	u32 devID = flasher.getDeviceID();
	cout << "Devide ID: " << devID << endl;

	if (devID == 0)
		return 1;



	std::string device = deviceIDs[devID];
	cout << "found " << device << endl;
	
	cout << "Erase flash" << endl;
	if (flasher.eraseFlash() != 0) {
		cerr << "Error erasing flash" << endl;
		return 2;
	}

	cout << "Done" << endl;

	
	flasher.flashImage(APPNAME);
	cout << "IMAGE WRITING DONE" << endl;

	return 0;
}