#pragma once

#include <math.h>
#include <stddef.h>
#include <stdint.h>

namespace stk500v2 {

enum class Msg : uint8_t {
	START = 0x1B, // ASCII ESC
	TOKEN = 0x0E,
};

inline constexpr bool operator==(const Msg &msg, const uint8_t &num)
{
	return static_cast<uint8_t>(msg) == num;
}

enum class Cmd : uint8_t {
	SIGN_ON = 0x01,
	SET_PARAMETER = 0x02,
	GET_PARAMETER = 0x03,
	SET_DEVICE_PARAMETERS = 0x04,
	OSCCAL = 0x05,
	LOAD_ADDRESS = 0x06,
	FIRMWARE_UPGRADE = 0x07,
};

inline constexpr bool operator==(const Cmd &cmd, const uint8_t &num)
{
	return static_cast<uint8_t>(cmd) == num;
}

enum class IspCmd : uint8_t {
	ENTER_PROGMODE = 0x10,
	LEAVE_PROGMODE = 0x11,
	CHIP_ERASE = 0x12,
	PROGRAM_FLASH = 0x13,
	READ_FLASH = 0x14,
	PROGRAM_EEPROM = 0x15,
	READ_EEPROM = 0x16,
	PROGRAM_FUSE = 0x17,
	READ_FUSE = 0x18,
	PROGRAM_LOCK = 0x19,
	READ_LOCK = 0x1A,
	READ_SIGNATURE = 0x1B,
	READ_OSCCAL = 0x1C,

	SPI_MULTI = 0x1D,
};

enum class PpCmd : uint8_t {
	ENTER_PROGMODE = 0x20,
	LEAVE_PROGMODE = 0x21,
	CHIP_ERASE = 0x22,
	PROGRAM_FLASH = 0x23,
	READ_FLASH = 0x24,
	PROGRAM_EEPROM = 0x25,
	READ_EEPROM = 0x26,
	PROGRAM_FUSE = 0x27,
	READ_FUSE = 0x28,
	PROGRAM_LOCK = 0x29,
	READ_LOCK = 0x2A,
	READ_SIGNATURE = 0x2B,
	READ_OSCCAL = 0x2C,

	SET_CONTROL_STACK = 0x2D,
};

enum class HvspCmd : uint8_t {
	ENTER_PROGMODE = 0x30,
	LEAVE_PROGMODE = 0x31,
	CHIP_ERASE = 0x32,
	PROGRAM_FLASH = 0x33,
	READ_FLASH = 0x34,
	PROGRAM_EEPROM = 0x35,
	READ_EEPROM = 0x36,
	PROGRAM_FUSE = 0x37,
	READ_FUSE = 0x38,
	PROGRAM_LOCK = 0x39,
	READ_LOCK = 0x3A,
	READ_SIGNATURE = 0x3B,
	READ_OSCCAL = 0x3C,
};

enum class Status : uint8_t {
	// Success
	CMD_OK = 0x00,

	// Warnings
	CMD_TOUT = 0x80,
	RDY_BSY_TOUT = 0x81,
	SET_PARAM_MISSING = 0x82,

	// Errors
	CMD_FAILED = 0xC0,
	CKSUM_ERROR = 0xC1,
	CMD_UNKNOWN = 0xC9,
};

enum class Param : uint8_t {
	BUILD_NUMBER_LOW = 0x80,
	BUILD_NUMBER_HIGH = 0x81,
	HW_VER = 0x90,
	SW_MAJOR = 0x91,
	SW_MINOR = 0x92,
	VTARGET = 0x94,
	VADJUST = 0x95,
	OSC_PSCALE = 0x96,
	OSC_CMATCH = 0x97,
	SCK_DURATION = 0x98,
	TOPCARD_DETECT = 0x9A,
	STATUS = 0x9C,
	DATA = 0x9D,
	RESET_POLARITY = 0x9E,
	CONTROLLER_INIT = 0x9F,
};

enum class Answer : uint8_t {
	CKSUM_ERROR = 0xB0,
};

template <size_t Size>
struct Message {
	uint8_t start;
	uint8_t number;
	uint16_t size;
	uint8_t token;
	uint8_t body[Size];
	uint8_t checksum;
};

template <auto SignOnFn, auto SetParamFn, auto GetParamFn>
struct Callbacks {
	static constexpr auto onSignOn = SignOnFn;
	static constexpr auto onSetParam = SetParamFn;
	static constexpr auto onGetParam = GetParamFn;
};

template <class Uart, uint16_t Timeout, auto CallbackFns>
class Stk500v2 {
  public:
	inline void init()
	{
		m_serial.init();
	}

	inline bool callback()
	{
		if (receiveMessage()) {
			handleMessage();
		}

		if (m_timeout)
			return true;

		return false;
	}

  private:
	static constexpr auto MSG_BUFFER_SIZE = 275;
	using msg_t = Message<MSG_BUFFER_SIZE>;

	msg_t m_msg;
	Uart m_serial;
	uint16_t m_timeout = Timeout;

	inline bool receiveByte(uint8_t &data)
	{
		constexpr auto MICROSECOND = 1000.0 * 1000;
		constexpr auto SYMBOL_SIZE = 9;
		constexpr auto BYTE_DELAY_US = (SYMBOL_SIZE * MICROSECOND) / Uart::BAUD_RATE;
		constexpr auto NUM_MS_DELAY_STEPS = static_cast<uint16_t>(round(1000 / BYTE_DELAY_US));
		uint16_t msDelay = NUM_MS_DELAY_STEPS;

		while (m_timeout) {
			if (m_serial.rxByte(data)) {
				m_timeout = Timeout;
				return true;
			}

			_delay_us(BYTE_DELAY_US);

			if (--msDelay == 0) {
				msDelay = NUM_MS_DELAY_STEPS;
				--m_timeout;
			}
		}

		return false;
	}

	inline uint8_t calcChecksum() const
	{
		uint8_t checksum = static_cast<uint8_t>(m_msg.start);

		for (uint16_t i = 1; i < 5 + m_msg.size; ++i) {
			checksum ^= *(reinterpret_cast<const uint8_t *>(&m_msg) + i);
		}

		return checksum;
	}

	inline bool receiveMessage()
	{
		constexpr auto CHECK_MESSAGE_FORMAT = true;
		constexpr auto CHECK_MESSAGE_SIZE = true;
		constexpr auto CHECK_MESSAGE_CHECKSUM = true;

		if (!receiveByte(m_msg.start)) {
			return false;
		}
		if constexpr (CHECK_MESSAGE_FORMAT) {
			if (m_msg.start != Msg::START)
				return false;
		}

		if (!receiveByte(m_msg.number))
			return false;

		if (!receiveByte(*(reinterpret_cast<uint8_t *>(&m_msg.size) + 1)))
			return false;
		if (!receiveByte(*reinterpret_cast<uint8_t *>(&m_msg.size)))
			return false;

		if constexpr (CHECK_MESSAGE_SIZE) {
			if (m_msg.size > sizeof(m_msg.body))
				return false;
		}

		if (!receiveByte(m_msg.token))
			return false;

		if constexpr (CHECK_MESSAGE_FORMAT) {
			if (m_msg.token != Msg::TOKEN)
				return false;
		}

		for (uint16_t i = 0; i < m_msg.size; ++i) {
			if (!receiveByte(m_msg.body[i]))
				return false;
		}

		if (!receiveByte(m_msg.checksum))
			return false;

		if constexpr (CHECK_MESSAGE_CHECKSUM) {
			if (m_msg.checksum != calcChecksum())
				return false;
		}

		return true;
	}

	inline void transmitMessage() const
	{
		m_serial.txByte(m_msg.start);
		m_serial.txByte(m_msg.number);
		m_serial.txByte(m_msg.size >> 8);
		m_serial.txByte(m_msg.size & 0xFF);
		m_serial.txByte(m_msg.token);
		for (uint16_t i = 0; i < m_msg.size; ++i)
			m_serial.txByte(m_msg.body[i]);
		m_serial.txByte(m_msg.checksum);
	}

	inline void handleMessage()
	{
		constexpr auto CHECK_MESSAGE_LENGTH = true;

		switch (m_msg.body[0]) {
		case static_cast<uint8_t>(Cmd::SIGN_ON): {
			if constexpr (CHECK_MESSAGE_LENGTH) {
				if (m_msg.size != 1)
					break;
			}
			CallbackFns.onSignOn();
			sendReply<static_cast<uint8_t>(Cmd::SIGN_ON)>();
			break;
		}
		case static_cast<uint8_t>(Cmd::SET_PARAMETER): {
			if constexpr (CHECK_MESSAGE_LENGTH) {
				if (m_msg.size != 3)
					break;
			}
			CallbackFns.onSetParam();
			sendReply<static_cast<uint8_t>(Cmd::SET_PARAMETER)>();
			break;
		}
		case static_cast<uint8_t>(Cmd::GET_PARAMETER): {
			if constexpr (CHECK_MESSAGE_LENGTH) {
				if (m_msg.size != 2)
					break;
			}
			const auto errorOccurred = CallbackFns.onGetParam(static_cast<Param>(m_msg.body[1]), m_msg.body[2]);
			sendReply<static_cast<uint8_t>(Cmd::GET_PARAMETER)>(errorOccurred);
			break;
		}
		case static_cast<uint8_t>(Cmd::SET_DEVICE_PARAMETERS):
			break;
		case static_cast<uint8_t>(Cmd::OSCCAL):
			break;
		case static_cast<uint8_t>(Cmd::LOAD_ADDRESS):
			break;
		case static_cast<uint8_t>(Cmd::FIRMWARE_UPGRADE):
			break;
		}
	}

	template <uint8_t Command>
	inline void sendReply(bool errorOccurred = false)
	{
		if constexpr (Command == Cmd::SIGN_ON) {
			m_msg.size = 3 + 8;
			m_msg.body[1] = static_cast<uint8_t>(Status::CMD_OK);
			m_msg.body[2] = 8;
			m_msg.body[3] = 'S';
			m_msg.body[4] = 'T';
			m_msg.body[5] = 'K';
			m_msg.body[6] = '5';
			m_msg.body[7] = '0';
			m_msg.body[8] = '0';
			m_msg.body[9] = '_';
			m_msg.body[10] = '2';
		} else if (Command == Cmd::SET_PARAMETER) {
			m_msg.size = 2;
			m_msg.body[1] = static_cast<uint8_t>(Status::CMD_OK);
		} else if (Command == Cmd::GET_PARAMETER) {
			if (!errorOccurred) {
				m_msg.size = 3;
				m_msg.body[1] = static_cast<uint8_t>(Status::CMD_OK);
			} else {
				m_msg.size = 2;
				m_msg.body[1] = static_cast<uint8_t>(Status::CMD_FAILED);
			}
		}

		m_msg.checksum = calcChecksum();

		transmitMessage();
	}
};

} // namespace stk500v2