#include "flash-writer.h"
#include <stm32/stm32f4xx.h>
#include <stdbool.h>

static bool flash_op_busy(void)
{
	return !!(FLASH->SR & FLASH_SR_BSY);
}

static void lock_flash_cr(void)
{
	FLASH->CR |= FLASH_CR_LOCK;
}

void flash_writer_enable_access(void)
{
	const uint32_t key1 = 0x45670123UL;
	const uint32_t key2 = 0xCDEF89ABUL;

	FLASH->KEYR = key1;
	FLASH->KEYR = key2;
}

void flash_writer_perform_mass_erase(void)
{
	flash_writer_enable_access();
	while (flash_op_busy());

	FLASH->CR = DMA_SxCR_PSIZE_1;
	FLASH->CR |= FLASH_CR_MER;
	FLASH->CR |= FLASH_CR_STRT;

	while(flash_op_busy());
	lock_flash_cr();
}

uint32_t flash_writer_get_flash_size(void)
{
	uint32_t flash_size;
	const uint16_t *flash_size_ptr = (const uint16_t *)0x1FFF7A22UL;

	flash_size = (uint32_t)*flash_size_ptr;
	flash_size *= 1024;

	return flash_size;
}

int flash_writer_write_to_memory(void *dest, const void *src, uint32_t size)
{
	uint32_t full_word_cnt;
	uint32_t byte_cnt;
	uint32_t idx;
	const uint32_t *word_src_ptr;
	uint32_t *word_dest_ptr;
	const char *char_src_ptr;
	char *char_dest_ptr;
	uint32_t pre_byte_count;

	flash_writer_enable_access();
	while (flash_op_busy());

	/* Number of full words to program */
	full_word_cnt = size / 4u;
	byte_cnt = size % 4;

	word_dest_ptr = dest;
	word_src_ptr = src;

	/* Do the first bytes, in case the destination is not word aligned */
	pre_byte_count = (4 - ((uint32_t)dest % 4u)) % 4;
	if (pre_byte_count) {
		FLASH->CR = 0u;
		FLASH->CR |= FLASH_CR_PG;

		char_src_ptr = src;
		char_dest_ptr = dest;

		/* Write bytes to memory until we hit the next word aligned address */
		for (idx = 0; idx < pre_byte_count; idx++) {
			*(char_dest_ptr++) = *(char_src_ptr++);
		}

		/* Correct the word addresses set above */
		full_word_cnt = (size - pre_byte_count) / 4u;
		byte_cnt = (size - pre_byte_count) % 4;
		word_dest_ptr = (uint32_t *)char_dest_ptr;
		word_src_ptr = (uint32_t *)char_src_ptr;

		while(flash_op_busy());
		FLASH->CR = 0u;
	}

	/* Do the full word flash write */
	if (full_word_cnt) {
		FLASH->CR = FLASH_CR_PSIZE_1;
		FLASH->CR |= FLASH_CR_PG;

		for (idx = 0; idx < full_word_cnt; idx++) {
			*word_dest_ptr = *word_src_ptr;
			word_dest_ptr++;
			word_src_ptr++;
		}

		while (flash_op_busy());
		FLASH->CR = 0u;
	}

	/* write remaining bytes */
	if (byte_cnt) {
		char_src_ptr = (char *)word_src_ptr;
		char_dest_ptr = (char *)word_dest_ptr;

		FLASH->CR = 0u;
		FLASH->CR |= FLASH_CR_PG;

		for (idx = 0; idx < byte_cnt; idx++) {
			*char_dest_ptr = *char_src_ptr;
			char_dest_ptr++;
			char_src_ptr++;
		}

		while (flash_op_busy());
		FLASH->CR = 0u;
	}

	lock_flash_cr();

	return 0;
}

uint32_t flash_writer_get_base_address(void)
{
	return FLASH_BASE;
}