#include <stdint.h>
#include <stddef.h>

#include "bsp/soc.h"
#include "bsp/spi_flash.h"
#include "bsp/uart.h"

#ifndef APP_FLASH_OFFSET
#define APP_FLASH_OFFSET 0x06000000UL
#endif

#define BOOT_IMAGE_MAGIC   0x54425652u
#define BOOT_IMAGE_VERSION 1u
#define BOOT_HEADER_SIZE   64u

struct boot_image_header {
  uint32_t magic;
  uint32_t version;
  uint32_t header_size;
  uint32_t flags;
  uint32_t payload_flash_offset;
  uint32_t load_addr;
  uint32_t entry;
  uint32_t size;
  uint32_t crc32;
  uint32_t reserved[7];
};

static void uart_puts(const char *s)
{
  while (*s != '\0') {
    bsp_uart_putc(*s++);
  }
}

static void uart_puthex32(uint32_t value)
{
  static const char hex[] = "0123456789abcdef";
  for (int shift = 28; shift >= 0; shift -= 4) {
    bsp_uart_putc(hex[(value >> shift) & 0xfU]);
  }
}

static uint32_t crc32_update(uint32_t crc, const uint8_t *data, uint32_t len)
{
  crc = ~crc;
  for (uint32_t i = 0; i < len; ++i) {
    crc ^= data[i];
    for (uint32_t bit = 0; bit < 8; ++bit) {
      uint32_t mask = 0U - (crc & 1U);
      crc = (crc >> 1) ^ (0xedb88320U & mask);
    }
  }
  return ~crc;
}

static void die(const char *message, uint32_t value)
{
  uart_puts("BOOT ERROR: ");
  uart_puts(message);
  uart_puts(" 0x");
  uart_puthex32(value);
  uart_puts("\n");
  while (1) {
    asm volatile ("wfi");
  }
}

int main(void)
{
  struct boot_image_header header;

  uart_puts("\nRocket flash bootloader\n");
  uart_puts("Header offset: 0x");
  uart_puthex32((uint32_t)APP_FLASH_OFFSET);
  uart_puts("\n");

  bsp_spi_flash_init();
  bsp_spi_flash_reset();

  uint32_t flash_id = bsp_spi_flash_read_id();
  uart_puts("Flash JEDEC ID: 0x");
  uart_puthex32(flash_id);
  uart_puts("\n");

  bsp_spi_flash_read((uint32_t)APP_FLASH_OFFSET, &header, sizeof(header));

  if (header.magic != BOOT_IMAGE_MAGIC) {
    die("bad image magic", header.magic);
  }
  if (header.version != BOOT_IMAGE_VERSION || header.header_size != BOOT_HEADER_SIZE) {
    die("bad image version/header", header.version);
  }
  if (header.size == 0U || header.size > 0x000f0000U) {
    die("bad image size", header.size);
  }
  if (header.load_addr < BSP_DTIM_BASE || (header.load_addr + header.size) > (BSP_DTIM_BASE + BSP_DTIM_SIZE)) {
    die("image outside DTIM", header.load_addr);
  }

  uart_puts("Load: flash 0x");
  uart_puthex32(header.payload_flash_offset);
  uart_puts(" -> mem 0x");
  uart_puthex32(header.load_addr);
  uart_puts(", size 0x");
  uart_puthex32(header.size);
  uart_puts("\n");

  bsp_spi_flash_read(header.payload_flash_offset, (void *)(uintptr_t)header.load_addr, header.size);

  uint32_t crc = crc32_update(0U, (const uint8_t *)(uintptr_t)header.load_addr, header.size);
  if (crc != header.crc32) {
    die("crc mismatch", crc);
  }

  uart_puts("CRC ok, jumping 0x");
  uart_puthex32(header.entry);
  uart_puts("\n");

  asm volatile ("fence.i" ::: "memory");
  void (*entry)(void) = (void (*)(void))(uintptr_t)header.entry;
  entry();

  return 0;
}