#include "clock.hpp"

#include "uart/uart.hpp"

#include <math.h>

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

#include "command.hpp"

static constexpr auto TIMEOUT = 5000;
static constexpr auto BAUD_RATE = 115200;

using uart_interface = uart::Hardware0<uart::Config<BAUD_RATE>, 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)
{
	constexpr auto MICROSECOND = 1000.0 * 1000;
	constexpr auto SYMBOL_SIZE = 9;
	constexpr auto BYTE_DELAY_US = (SYMBOL_SIZE * MICROSECOND) / 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 (timeout) {
		if (serial.rxByte(data)) {
			timeout = TIMEOUT;
			return true;
		}

		_delay_us(BYTE_DELAY_US);

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

	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, uint16_t &timeout)
{
	if (!receiveByte(msg.start, timeout) || msg.start != MESSAGE_START)
		return false;
	if (!receiveByte(msg.number, timeout))
		return false;
	if (!receiveByte(*(reinterpret_cast<uint8_t *>(&msg.size) + 1), timeout))
		return false;
	if (!receiveByte(*reinterpret_cast<uint8_t *>(&msg.size), timeout) || msg.size > sizeof(msg.body))
		return false;
	if (!receiveByte(msg.token, timeout) || msg.token != TOKEN)
		return false;
	for (uint16_t i = 0; i < msg.size; ++i) {
		if (!receiveByte(msg.body[i], timeout))
			return false;
	}
	if (!receiveByte(msg.checksum, timeout) || msg.checksum != calcChecksum(msg))
		return false;

	return true;
}

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);
}

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 isReadEepromIsp(const Message &msg)
{
	if (msg.size == 4 && msg.body[0] == CMD_READ_EEPROM_ISP)
		return true;
	return false;
}

static inline bool isChipEraseIsp(const Message &msg)
{
	if (msg.size == 7 && msg.body[0] == CMD_CHIP_ERASE_ISP)
		return true;
	return false;
}

static inline bool isProgramFlashIsp(const Message &msg)
{
	if (msg.body[0] == CMD_PROGRAM_FLASH_ISP) {
		const auto dataSize = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
		if (msg.size == (dataSize + 10) && dataSize == SPM_PAGESIZE)
			return true;
	}
	return false;
}

static inline bool isProgramEepromIsp(const Message &msg)
{
	if (msg.body[0] == CMD_PROGRAM_EEPROM_ISP) {
		if (msg.size == (static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2]) + 10)
			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 if (msg.body[1] == PARAM_TOPCARD_DETECT) {
		msg.body[2] = 0xFF;
	} 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 byteAddress = 2 * 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>(byteAddress + i));
	}
	const auto numWords = numBytes / 2;
	addr += numWords;
	msg.body[numBytes + 2] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

namespace {

bool isEepromReady()
{
	return (EECR & (1 << EEPE)) ? false : true;
}

void waitEepromReady()
{
	while (!isEepromReady())
		;
}

uint8_t readEepromByte(const uint8_t *addr)
{
	EEAR = reinterpret_cast<uint16_t>(addr);
	EECR |= (1 << EERE);
	return EEDR;
}

void writeEepromByte(uint8_t *addr, uint8_t value)
{
	EECR = 0;
	EEAR = reinterpret_cast<uint16_t>(addr);
	EEDR = value;
	EECR |= (1 << EEMPE);
	EECR |= (1 << EEPE);
}

//////////////////////////////////////////////////////////////////////////

void writeFlashPage(uint32_t pageAddress, const uint8_t *data)
{
	boot_page_erase(pageAddress);
	boot_spm_busy_wait();

	for (uint16_t i = 0; i < SPM_PAGESIZE; i += 2) {
		uint16_t dataWord = *data++;
		dataWord |= (*data++) << 8;
		boot_page_fill(pageAddress + i, dataWord);
	}

	boot_page_write(pageAddress);
	boot_spm_busy_wait();
	boot_rww_enable();
}

} // namespace

static inline uint16_t getBootloaderSize()
{
	const auto highFuse = boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS);
	constexpr auto BOOTSZ0 = 1;
	constexpr auto BOOTSZ1 = 2;

	if (highFuse & (1 << BOOTSZ1) && highFuse & (1 << BOOTSZ0))
		return 256 * 2;
	else if (highFuse & (1 << BOOTSZ1))
		return 512 * 2;
	else if (highFuse & (1 << BOOTSZ0))
		return 1024 * 2;
	return 2048 * 2;
}

static inline uint32_t getFlashSize()
{
	const auto bootloaderSize = getBootloaderSize();
	return (FLASHEND - bootloaderSize + 1);
}

static inline void performChipErase(uint16_t flashStartAddress = 0x0000)
{
	constexpr auto getEepromEraseFuseBit = []() -> bool {
		constexpr auto EESAVE = 3;
		return boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS) & (1 << EESAVE);
	};

	constexpr auto eraseFlash = [](const uint16_t &flashStartAddress) {
		const auto flashSize = getFlashSize();
		const auto byteAddress = 2 * flashStartAddress;
		for (uint16_t i = byteAddress; i < flashSize; i += SPM_PAGESIZE) {
			boot_page_erase(i);
			boot_spm_busy_wait();
		}
		boot_rww_enable();
	};

	constexpr auto eraseEeprom = [getEepromEraseFuseBit]() {
		const auto eraseEeprom = getEepromEraseFuseBit();
		if (eraseEeprom) {
			constexpr auto EEPROM_SIZE = E2END + 1;
			for (uint16_t i = 0; i < EEPROM_SIZE; ++i) {
				writeEepromByte(reinterpret_cast<uint8_t *>(i), 0xFF);
				waitEepromReady();
			}
		}
	};

	eraseFlash(flashStartAddress);
	eraseEeprom();
}

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

	msg.checksum = calcChecksum(msg);
}

static inline void formatReadEepromIspAnswer(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] = readEepromByte(reinterpret_cast<const uint8_t *>(addr + i));
	}
	addr += numBytes;
	msg.body[numBytes + 2] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatProgramFlashIspAnswer(Message &msg, uint32_t &addr)
{
	const auto byteAddress = 2 * addr;
	if (byteAddress < getFlashSize())
		writeFlashPage(byteAddress, msg.body + 10);
	const uint16_t numBytes = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
	const auto numWords = numBytes / 2;
	addr += numWords;

	msg.size = 2;
	msg.body[1] = STATUS_CMD_OK;

	msg.checksum = calcChecksum(msg);
}

static inline void formatProgramEepromIspAnswer(Message &msg, uint32_t &addr)
{
	const uint16_t numBytes = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
	for (uint16_t i = 0; i < numBytes; ++i) {
		writeEepromByte(reinterpret_cast<uint8_t *>(addr + i), msg.body[10 + i]);
		waitEepromReady();
	}
	addr += numBytes;
	msg.size = 2;
	msg.body[1] = 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);
}

enum class ChipEraseState {
	NONE = 0,
	REQUEST = (1 << 1),
	RESPONSE = (1 << 2),
	PERFORM = REQUEST | RESPONSE,
	PROGRAM = (1 << 3),
	FINISH = REQUEST | RESPONSE | PROGRAM,
};

constexpr ChipEraseState operator|(const ChipEraseState &self, const ChipEraseState &other)
{
	return static_cast<ChipEraseState>(static_cast<uint8_t>(self) | static_cast<uint8_t>(other));
}

constexpr ChipEraseState &operator|=(ChipEraseState &self, const ChipEraseState &other)
{
	self = self | other;
	return self;
}

static inline void handleMessage(Message &msg, uint32_t &addr, uint16_t &finishEraseAddress,
                                 ChipEraseState &chipEraseFlag)
{
	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)) {
		addr = msg.body[1];
		addr = (addr << 8) | msg.body[2];
		addr = (addr << 8) | msg.body[3];
		addr = (addr << 8) | msg.body[4];
		formatLoadAddressAnswer(msg);
	} else if (isReadFlashIsp(msg))
		formatReadFlashIspAnswer(msg, addr);
	else if (isReadEepromIsp(msg))
		formatReadEepromIspAnswer(msg, addr);
	else if (isChipEraseIsp(msg)) {
		chipEraseFlag |= ChipEraseState::REQUEST;
		formatChipEraseIspAnswer(msg);
	} else if (isProgramFlashIsp(msg)) {
		chipEraseFlag |= ChipEraseState::PROGRAM;
		formatProgramFlashIspAnswer(msg, addr);
		finishEraseAddress = addr;
	} else if (isProgramEepromIsp(msg))
		formatProgramEepromIspAnswer(msg, addr);
	else if (isLeaveProgmodeIsp(msg)) {
		chipEraseFlag |= ChipEraseState::RESPONSE;
		formatLeaveProgmodeIspAnswer(msg);
	} else
		formatErrorAnswer(msg);

	transmitMessage(msg);
}

int main()
{
	serial.init();

	Message msg;
	uint32_t addr = 0x0000;
	uint16_t finishEraseAddress = 0x0000;
	ChipEraseState chipEraseFlag = ChipEraseState::NONE;
	uint16_t timeout = TIMEOUT;

	while (true) {
		if (receiveMessage(msg, timeout)) {
			handleMessage(msg, addr, finishEraseAddress, chipEraseFlag);
		}

		if (timeout == 0) {
			if (chipEraseFlag == ChipEraseState::PERFORM) {
				performChipErase();
				chipEraseFlag = ChipEraseState::NONE;
			} else if (chipEraseFlag == ChipEraseState::FINISH) {
				performChipErase(finishEraseAddress);
				chipEraseFlag = ChipEraseState::NONE;
			}

			asm volatile("jmp 0x0000");
		}
	}

	return 0;
}

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