/* SPDX-License-Identifier: GPL-2.0 */

#include "rtl8710_flasher.h"

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

#define printf DiagPrintf

static void print_flash_info(const uint8_t *flash_id) {
  printf("[FLASH] Device ID: %02X %02X %02X\n", flash_id[0], flash_id[1],
         flash_id[2]);
  printf("[FLASH] Vendor: 0x%08X\n", flash_init_para.vendor_id);
  printf("[FLASH] Config: Mode=%d, Clock=%d, ReadCmd=0x%02X\n",
         flash_init_para.current_mode, flash_init_para.clock_div,
         flash_init_para.current_read_cmd);
  printf("[FLASH] Timing: SampleDelay=%d, DummyCycles=[%d,%d,%d]\n",
         flash_init_para.read_sample_delay,
         flash_init_para.read_dummy_cycles[0],
         flash_init_para.read_dummy_cycles[1],
         flash_init_para.read_dummy_cycles[2]);
}

int __attribute__((section(".vectors"))) main(void) {
  __asm__ volatile("cpsid if");

  uint8_t flash_id[4] = {0};

  // sometimes we boot up and the ROM already inits flash?
  if (PERI_ON_SOC_FUNC_EN & FLASH_PERIPH_BIT) {
    printf("[FLASHER] Flash peripheral already enabled...?\n");
  }

  RCC_PeriphClockCmd(FLASH_PERIPH_BIT, FLASH_CLOCK_BIT, 0);
  FLASH_ClockDiv(0);
  RCC_PeriphClockCmd(FLASH_PERIPH_BIT, FLASH_CLOCK_BIT, 1);

  PINMUX_Ctrl(10, 0, 1);
  FLASH_StructInit_GD(&flash_init_para);
  FLASH_Init(2);

  FLASH_RxCmd(flash_init_para.cmd_read_id, 3, flash_id);
  print_flash_info(flash_id);

  /* clear block protection */
  FLASH_SetStatusBits(0x1c, 0);
  printf("[FLASH] Block protection cleared\n");
  Cache_Enable(0);

  while (1) {
    FLASH_CONTROL->start = 0;
    while (!FLASH_CONTROL->start) __asm__("nop");

    switch (FLASH_CONTROL->cmd) {
      case CMD_READ_ID:
        FLASH_RxCmd(flash_init_para.cmd_read_id, 4, flash_id);
        FLASH_CONTROL->param = *(uint32_t *)flash_id;
        break;

      case CMD_MASS_ERASE:
        FLASH_Erase(0, 0);
        printf("[FLASHER] Full chip erase completed\n");
        break;

      case CMD_SECTOR_ERASE:
        FLASH_Erase(2, FLASH_CONTROL->offset);
        break;

      case CMD_READ:
        for (uint32_t i = 0; i < FLASH_CONTROL->len; i += 4) {
          volatile uint32_t *flash_addr =
              (volatile uint32_t *)(SPI_FLASH_BASE + FLASH_CONTROL->offset + i);
          uint32_t *dest = (uint32_t *)&FLASH_CONTROL->data[i];
          *dest = *flash_addr;
        }
        break;

      case CMD_WRITE:
        printf("[FLASHER] Starting write of %d bytes at offset 0x%08X\n",
               FLASH_CONTROL->len, FLASH_CONTROL->offset);

        for (uint32_t offset = 0; offset < FLASH_CONTROL->len; offset += 4) {
          uint32_t write_addr = FLASH_CONTROL->offset + offset;
          FLASH_TxData12B(write_addr, 4,
                          (uint8_t *)&FLASH_CONTROL->data[offset]);
        }

        printf("[FLASHER] Write completed successfully\n");
        break;

      case CMD_VERIFY: {
        uint8_t tbuf[FLASH_SECTOR_SIZE];
        uint32_t start_sector = FLASH_CONTROL->offset / FLASH_SECTOR_SIZE;
        uint32_t end_sector = (FLASH_CONTROL->offset + FLASH_CONTROL->len - 1) /
                              FLASH_SECTOR_SIZE;

        for (uint32_t i = 0; i < FLASH_CONTROL->len; i += 4) {
          volatile uint32_t *flash_addr =
              (volatile uint32_t *)(SPI_FLASH_BASE + FLASH_CONTROL->offset + i);
          *((uint32_t *)tbuf) = *flash_addr;

          size_t l = MIN(4, FLASH_CONTROL->len - i);
          size_t k;
          uint32_t current_sector =
              (FLASH_CONTROL->offset + i) / FLASH_SECTOR_SIZE;
          if (current_sector > end_sector) {
            break;
          }

          for (k = 0; k < l; k++) {
            if (tbuf[k] != FLASH_CONTROL->data[i + k]) {
              printf("[VERIFY] mismatch at offset %x: flash=%02x buf=%02x\n",
                     i + k, tbuf[k], FLASH_CONTROL->data[i + k]);
              break;
            }
          }
          if (k < l) {
            FLASH_CONTROL->status = 1;
            FLASH_CONTROL->param = i;
            break;
          }
        }
        break;
      }

      default:
        FLASH_CONTROL->status = 1;
        break;
    }
  }
}