#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;

void firmware_update() {
	printf("FIRMWARE UPDATE\n\r");

	UpdateTarget target;
	HAL_UART_Receive(&huart2, &target, 1, 1000);
	printf("T: %i\n\r", target);

	uint8_t address = 0;
	if (target == TARGET_I2C) {
		HAL_UART_Receive(&huart2, &address, 1, 1000);
		printf("A: %i\n\r", address);
	}

	uint8_t length_buffer[2];
	HAL_UART_Receive(&huart2, length_buffer, 2, 1000);
	uint16_t length = length_buffer[0] + (length_buffer[1] << 8);
	printf("L: %i\n\r", length);

	uint8_t* data = malloc(length);
	HAL_UART_Receive(&huart2, data, length, 1000);

	printf("Firmware received!\n\r");

	switch (target) {
		case TARGET_ROM:
			control_program_eeprom(data, length);
			break;
		case TARGET_I2C: {
			// Make sure tell the device to enter bootloader mode first
			// Upload the application
			printf("Uploading to target...\n\r");
			uint8_t upload_command[4 + length];
			upload_command[0] = 0x02;
			upload_command[1] = 0x01;

			printf("0/%i\r", length);

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

				printf("%i/%i\r", length-remaining, length);
			}

			printf("\n");

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

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

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