#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include "firmware.h"
#include "control.h"
#include "io.h"

extern I2C_HandleTypeDef hi2c1;
extern UART_HandleTypeDef huart2;

typedef enum {
	TARGET_NONE = 0,
	TARGET_ROM,
	TARGET_I2C
} UpdateTarget;

extern uint8_t ack[1];
extern uint8_t nack[1];
void firmware_update() {
	UpdateTarget target;

	HAL_UART_Receive(&huart2, &target, 1, 1000);
	HAL_UART_Transmit(&huart2, ack, sizeof(ack), HAL_MAX_DELAY);

	uint8_t length_buffer[2];
	HAL_UART_Receive(&huart2, length_buffer, sizeof(length_buffer), HAL_MAX_DELAY);
	HAL_UART_Transmit(&huart2, length_buffer, sizeof(length_buffer), HAL_MAX_DELAY);
	uint16_t length = length_buffer[0] + (length_buffer[1] << 8);

	uint8_t* data = malloc(length);
	HAL_UART_Receive(&huart2, data, length, 1000);
	HAL_UART_Transmit(&huart2, ack, sizeof(ack), HAL_MAX_DELAY);

	switch (target) {
		case TARGET_ROM:
			control_program_eeprom(data, length);
			break;

		case TARGET_I2C: {
			uint8_t address = 0;
			HAL_UART_Receive(&huart2, &address, 1, 1000);
			HAL_UART_Transmit(&huart2, &address, sizeof(ack), HAL_MAX_DELAY);

			if (HAL_I2C_IsDeviceReady(&hi2c1, 0x29 << 1, 10, 1000) != HAL_OK) {
				HAL_UART_Transmit(&huart2, nack, sizeof(nack), HAL_MAX_DELAY);
				return;
			}
			HAL_UART_Transmit(&huart2, ack, sizeof(ack), HAL_MAX_DELAY);

			// @todo This assumes that the target it not already in bootloader mode, otherwise this will mess things up
			uint8_t reboot_command[] = {0xFF};
			HAL_I2C_Master_Transmit(&hi2c1, address << 1, reboot_command, sizeof(reboot_command), 1000);

			while (HAL_I2C_IsDeviceReady(&hi2c1, 0x29 << 1, 10, 1000) != HAL_OK);

			HAL_UART_Transmit(&huart2, ack, sizeof(ack), HAL_MAX_DELAY);

			uint8_t abort_command[] = {0x00};
			HAL_I2C_Master_Transmit(&hi2c1, address << 1, abort_command, sizeof(abort_command), 1000);

			uint8_t version_command[] = {0x01};
			HAL_I2C_Master_Transmit(&hi2c1, address << 1, version_command, sizeof(version_command), 1000);
			uint8_t version_string[16] = {0};
			HAL_I2C_Master_Receive(&hi2c1, address << 1, version_string, sizeof(version_string), 1000);
			HAL_UART_Transmit(&huart2, version_string, sizeof(version_string), HAL_MAX_DELAY);

			uint8_t upload_command[4 + 0x80];
			upload_command[0] = 0x02;
			upload_command[1] = 0x01;

			uint16_t remaining = length;
			for (uint8_t offset = 0; remaining > 0; ++offset) {
				upload_command[2] = ((offset*0x80) >> 8) & 0xFF;
				upload_command[3] = (offset*0x80) & 0xFF;

				uint8_t amount = remaining > 0x80 ? 0x80 : remaining;
				remaining -= amount;

				// We need to handle the last section properly
				memcpy(&upload_command[4], &data[offset*0x80], amount);

				HAL_I2C_Master_Transmit(&hi2c1, address << 1, upload_command, sizeof(upload_command), 1000);

				while (HAL_I2C_IsDeviceReady(&hi2c1, 0x29 << 1, 10, 1000) != HAL_OK);

				uint8_t progress[] = {(length-remaining) & 0xFF, (length-remaining) >> 8};
				HAL_UART_Transmit(&huart2, progress, sizeof(progress), HAL_MAX_DELAY);
			}

			// Start application
			uint8_t start_command[] = {0x01, 0x80};
			HAL_I2C_Master_Transmit(&hi2c1, address << 1, start_command, sizeof(start_command), 1000);

			HAL_UART_Transmit(&huart2, ack, sizeof(ack), HAL_MAX_DELAY);

			break;
		}

		default:
			printf("Target not implemented!\n\r");
			break;
	}

	printf("Complete!\n\r");
}