#include "clock.hpp"

#include "uart/uart.hpp"

#include <avr/boot.h>
#include <avr/interrupt.h>
#include <avr/pgmspace.h>

#include "command.hpp"

using uart_interface = uart::Hardware0<uart::Config<115200>, uart::Driven::BLOCKING>;
uart::Uart<uart_interface> serial;

struct Message {
	uint8_t start;
	uint8_t number;
	uint16_t size;
	uint8_t token;
	uint8_t body[275];
	uint8_t checksum;
};

static inline bool receiveByte(uint8_t &data, uint16_t timeout = 1000)
{
	for (uint16_t i = 0; i < timeout; ++i) {
		if (serial.rxByte(data))
			return true;
	}

	return false;
}

static inline uint8_t calcChecksum(const Message &msg)
{
	uint8_t checksum = msg.start;

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

	return checksum;
}

static inline bool receiveMessage(Message &msg)
{
	if (!receiveByte(msg.start) || msg.start != MESSAGE_START)
		return false;
	if (!receiveByte(msg.number))
		return false;
	if (!receiveByte(*(reinterpret_cast<uint8_t *>(&msg.size) + 1)))
		return false;
	if (!receiveByte(*reinterpret_cast<uint8_t *>(&msg.size)) || msg.size > sizeof(msg.body))
		return false;
	if (!receiveByte(msg.token) || msg.token != TOKEN)
		return false;
	for (uint16_t i = 0; i < msg.size; ++i) {
		if (!receiveByte(msg.body[i]))
			return false;
	}
	if (!receiveByte(msg.checksum) || msg.checksum != calcChecksum(msg))
		return false;

	return true;
}

static inline bool isSignOn(const Message &msg)
{
	if (msg.size == 1 && msg.body[0] == CMD_SIGN_ON)
		return true;
	return false;
}

static inline bool isGetParameter(const Message &msg)
{
	if (msg.size == 2 && msg.body[0] == CMD_GET_PARAMETER)
		return true;
	return false;
}

static inline bool isSetParameter(const Message &msg)
{
	if (msg.size == 3 && msg.body[0] == CMD_SET_PARAMETER)
		return true;
	return false;
}

static inline bool isEnterProgmodeIsp(const Message &msg)
{
	if (msg.size == 12 && msg.body[0] == CMD_ENTER_PROGMODE_ISP)
		return true;
	return false;
}

static inline bool isReadSignatureIsp(const Message &msg)
{
	if (msg.size == 6 && msg.body[0] == CMD_READ_SIGNATURE_ISP)
		return true;
	return false;
}

static inline bool isReadFuseIsp(const Message &msg)
{
	if (msg.size == 6 && msg.body[0] == CMD_READ_FUSE_ISP)
		return true;
	return false;
}

static inline bool isReadLockIsp(const Message &msg)
{
	if (msg.size == 6 && msg.body[0] == CMD_READ_LOCK_ISP)
		return true;
	return false;
}

static inline bool isLoadAddress(const Message &msg)
{
	if (msg.size == 5 && msg.body[0] == CMD_LOAD_ADDRESS)
		return true;
	return false;
}

static inline bool isReadFlashIsp(const Message &msg)
{
	if (msg.size == 4 && msg.body[0] == CMD_READ_FLASH_ISP)
		return true;
	return false;
}

static inline bool isLeaveProgmodeIsp(const Message &msg)
{
	if (msg.size == 3 && msg.body[0] == CMD_LEAVE_PROGMODE_ISP)
		return true;
	return false;
}

static inline void formatSignOnAnswer(Message &msg)
{
	msg.size = 3 + 8;
	msg.body[1] = STATUS_CMD_OK;
	msg.body[2] = 8;
	msg.body[3] = 'S';
	msg.body[4] = 'T';
	msg.body[5] = 'K';
	msg.body[6] = '5';
	msg.body[7] = '0';
	msg.body[8] = '0';
	msg.body[9] = '_';
	msg.body[10] = '2';
	msg.checksum = calcChecksum(msg);
}

static inline void formatGetParameterAnswer(Message &msg)
{
	msg.size = 3;

	if (msg.body[1] == PARAM_HW_VER) {
		msg.body[2] = 1;
	} else if (msg.body[1] == PARAM_SW_MAJOR) {
		msg.body[2] = 0x02;
	} else if (msg.body[1] == PARAM_SW_MINOR) {
		msg.body[2] = 0x0a;
	} else if (msg.body[1] == PARAM_SCK_DURATION) {
		msg.body[2] = 2;
	} else if (msg.body[1] == PARAM_VADJUST) {
		msg.body[2] = 25;
	} else if (msg.body[1] == PARAM_VTARGET) {
		msg.body[2] = 49;
	} else if (msg.body[1] == PARAM_OSC_PSCALE) {
		msg.body[2] = 2;
	} else if (msg.body[1] == PARAM_OSC_CMATCH) {
		msg.body[2] = 127;
	} else {
		msg.size = 2;
	}

	if (msg.size == 2) {
		msg.body[1] = STATUS_CMD_FAILED;
	} else {
		msg.body[1] = STATUS_CMD_OK;
	}

	msg.checksum = calcChecksum(msg);
}

static inline void formatSetParameterAnswer(Message &msg)
{
	msg.size = 2;
	msg.body[1] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatEnterProgmodeIspAnswer(Message &msg)
{
	msg.size = 2;
	msg.body[1] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatReadSignatureIspAnswer(Message &msg)
{
	msg.size = 4;
	msg.body[2] = boot_signature_byte_get(msg.body[4] * 2);
	msg.body[1] = STATUS_CMD_OK;
	msg.body[3] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatReadFuseIspAnswer(Message &msg)
{
	constexpr auto READ_LOW_FUSE_BITS = 0x0050;
	constexpr auto READ_HIGH_FUSE_BITS = 0x0858;
	constexpr auto READ_EXTENDED_FUSE_BITS = 0x0850;

	msg.size = 4;

	if (*reinterpret_cast<uint16_t *>(msg.body + 2) == READ_EXTENDED_FUSE_BITS) {
		msg.body[2] = boot_lock_fuse_bits_get(GET_EXTENDED_FUSE_BITS);
	}
	if (*reinterpret_cast<uint16_t *>(msg.body + 2) == READ_HIGH_FUSE_BITS) {
		msg.body[2] = boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS);
	}
	if (*reinterpret_cast<uint16_t *>(msg.body + 2) == READ_LOW_FUSE_BITS) {
		msg.body[2] = boot_lock_fuse_bits_get(GET_LOW_FUSE_BITS);
	}

	msg.body[1] = STATUS_CMD_OK;
	msg.body[3] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatReadLockIspAnswer(Message &msg)
{
	msg.size = 4;
	msg.body[2] = boot_lock_fuse_bits_get(GET_LOCK_BITS);
	msg.body[1] = STATUS_CMD_OK;
	msg.body[3] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatLoadAddressAnswer(Message &msg)
{
	msg.size = 2;
	msg.body[1] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatReadFlashIspAnswer(Message &msg, uint32_t &addr)
{
	const uint16_t numBytes = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
	msg.size = 3 + numBytes;
	msg.body[1] = STATUS_CMD_OK;
	for (uint16_t i = 0; i < numBytes; ++i) {
		msg.body[i + 2] = pgm_read_byte(static_cast<uint16_t>(addr + i));
	}
	addr += numBytes;
	msg.body[numBytes + 2] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatLeaveProgmodeIspAnswer(Message &msg)
{
	msg.size = 2;
	msg.body[1] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatErrorAnswer(Message &msg)
{
	msg.start = MESSAGE_START;
	msg.size = 1;
	msg.token = TOKEN;
	msg.body[0] = STATUS_CMD_UNKNOWN;
	msg.checksum = calcChecksum(msg);
}

static inline void handleMessage(Message &msg)
{
	static uint32_t s_address;

	if (isSignOn(msg))
		formatSignOnAnswer(msg);
	else if (isGetParameter(msg))
		formatGetParameterAnswer(msg);
	else if (isSetParameter(msg))
		formatSetParameterAnswer(msg);
	else if (isEnterProgmodeIsp(msg))
		formatEnterProgmodeIspAnswer(msg);
	else if (isReadSignatureIsp(msg))
		formatReadSignatureIspAnswer(msg);
	else if (isReadFuseIsp(msg))
		formatReadFuseIspAnswer(msg);
	else if (isReadLockIsp(msg))
		formatReadLockIspAnswer(msg);
	else if (isLoadAddress(msg)) {
		s_address = msg.body[1];
		s_address = (s_address << 8) | msg.body[2];
		s_address = (s_address << 8) | msg.body[3];
		s_address = (s_address << 8) | msg.body[4];
		formatLoadAddressAnswer(msg);
	} else if (isReadFlashIsp(msg))
		formatReadFlashIspAnswer(msg, s_address);
	else if (isLeaveProgmodeIsp(msg))
		formatLeaveProgmodeIspAnswer(msg);
	else
		formatErrorAnswer(msg);
}

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

int main()
{
	serial.init();

	static Message s_msg;

	while (true) {
		if (receiveMessage(s_msg)) {
			handleMessage(s_msg);
			transmitMessage(s_msg);
		}
	}

	return 0;
}

void startup() __attribute__((naked, section(".vectors")));
void startup()
{
	asm volatile("clr __zero_reg__");
	// SP = RAMEND;
	SREG = 0;
	asm volatile("jmp main");
}