// SPDX-License-Identifier: BSD-3-Clause
// Copyright (c) 2025, Alex Taradov <alex@taradov.com>. All rights reserved.

/*- Includes ----------------------------------------------------------------*/
#include <unistd.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include "target.h"
#include "edbg.h"
#include "dap.h"

/*- Definitions -------------------------------------------------------------*/
#define FLASH_ADDR             0x08000000

#define FLASH_PAGE_SIZE_1      4096
#define FLASH_ROW_SIZE_1       512

#define FLASH_PAGE_SIZE_2      2048
#define FLASH_ROW_SIZE_2       256

#define DHCSR                  0xe000edf0
#define DHCSR_DEBUGEN          (1 << 0)
#define DHCSR_HALT             (1 << 1)
#define DHCSR_DBGKEY           (0xa05f << 16)

#define DEMCR                  0xe000edfc
#define DEMCR_VC_CORERESET     (1 << 0)

#define AIRCR                  0xe000ed0c
#define AIRCR_VECTKEY          (0x05fa << 16)
#define AIRCR_SYSRESETREQ      (1 << 2)

#define FLASH_ACR              0x40022000
#define FLASH_PDKEYR           0x40022004
#define FLASH_KEYR             0x40022008
#define FLASH_OPTKEYR          0x4002200c
#define FLASH_SR               0x40022010
#define FLASH_CR               0x40022014
#define FLASH_ECCR             0x40022018
#define FLASH_OPTR             0x40022020
#define FLASH_PCROP1SR         0x40022024
#define FLASH_PCROP1ER         0x40022028
#define FLASH_WRP1AR           0x4002202c
#define FLASH_WRP1BR           0x40022030
#define FLASH_PCROP2SR         0x40022044
#define FLASH_PCROP2ER         0x40022048
#define FLASH_WRP2AR           0x4002204c
#define FLASH_WRP2BR           0x40022050
#define FLASH_SEC1R            0x40022070
#define FLASH_SEC2R            0x40022074

#define DBGMCU_IDCODE          0xe0042000
#define FLASH_SIZE_REG         0x1fff75e0
#define PACKAGE_REG            0x1fff7500

#define FLASH_KEYR_KEY1        0x45670123
#define FLASH_KEYR_KEY2        0xcdef89ab

#define FLASH_OPTKEYR_KEY1     0x08192a3b
#define FLASH_OPTKEYR_KEY2     0x4c5d6e7f

#define FLASH_SR_EOP           (1 << 0)
#define FLASH_SR_OPERR         (1 << 1)
#define FLASH_SR_PROGERR       (1 << 3)
#define FLASH_SR_WRPERR        (1 << 4)
#define FLASH_SR_PGAERR        (1 << 5)
#define FLASH_SR_SIZERR        (1 << 6)
#define FLASH_SR_PGSERR        (1 << 7)
#define FLASH_SR_MISSERR       (1 << 8)
#define FLASH_SR_FASTERR       (1 << 9)
#define FLASH_SR_RDERR         (1 << 14)
#define FLASH_SR_OPTVERR       (1 << 15)
#define FLASH_SR_BSY           (1 << 16)

#define FLASH_SR_ALL_ERRORS    (FLASH_SR_OPERR | FLASH_SR_PROGERR | FLASH_SR_WRPERR | \
    FLASH_SR_PGAERR | FLASH_SR_SIZERR | FLASH_SR_PGSERR | FLASH_SR_MISSERR | \
    FLASH_SR_FASTERR | FLASH_SR_RDERR | FLASH_SR_OPTVERR)

#define FLASH_CR_PG            (1 << 0)
#define FLASH_CR_PER           (1 << 1)
#define FLASH_CR_MER1          (1 << 2)
#define FLASH_CR_PNB(x)        ((x) << 3)
#define FLASH_CR_BKER          (1 << 11)
#define FLASH_CR_MER2          (1 << 15)
#define FLASH_CR_STRT          (1 << 16)
#define FLASH_CR_OPTSTRT       (1 << 17)
#define FLASH_CR_FSTPG         (1 << 18)
#define FLASH_CR_OBL_LAUNCH    (1 << 27)
#define FLASH_CR_SEC_PROT1     (1 << 28)
#define FLASH_CR_SEC_PROT2     (1 << 29)
#define FLASH_CR_OPT_LOCK      (1 << 30)
#define FLASH_CR_LOCK          (1 << 31)

#define FLASH_SIZE_REG_MASK    0xffff
#define FLASH_SIZE_REG_MULT    1024

#define FLASH_OPTR_RDP_MASK    0x000000ff
#define FLASH_OPTR_DBANK       (1 << 22)

#define OPTIONS_SIZE           4

#define DEVICE_ID_MASK         0x0000ffff

/*- Types -------------------------------------------------------------------*/
typedef struct
{
  uint32_t  idcode;
  char      *family;
  char      *name;
} device_t;

/*- Variables ---------------------------------------------------------------*/
static device_t devices[] =
{
  { 0x6468, "stm32g4", "STM32G431/41" },
  { 0x6469, "stm32g4", "STM32G471/73/74/83/84" },
  { 0x6479, "stm32g4", "STM32G491/A1" },
};

static device_t target_device;
static target_options_t target_options;
static bool target_dbank;
static int target_page_size;
static int target_row_size;

/*- Implementations ---------------------------------------------------------*/

//-----------------------------------------------------------------------------
static void flash_wait_done(void)
{
  while (dap_read_word(FLASH_SR) & FLASH_SR_BSY);

  uint32_t sr = dap_read_word(FLASH_SR);

  if (sr & FLASH_SR_ALL_ERRORS)
    error_exit("flash operation failed. FLASH_SR = 0x%08x", sr);
}

//-----------------------------------------------------------------------------
static void target_select(target_options_t *options)
{
  dap_disconnect();
  dap_connect(DAP_INTERFACE_SWD);
  dap_reset_pin(0);
  dap_reset_link();

  // Stop the core
  dap_write_word(DHCSR, DHCSR_DBGKEY | DHCSR_DEBUGEN | DHCSR_HALT);
  dap_write_word(DEMCR, DEMCR_VC_CORERESET);
  dap_write_word(AIRCR, AIRCR_VECTKEY | AIRCR_SYSRESETREQ);

  dap_reset_pin(1);
  sleep_ms(10);

  uint32_t idcode = dap_read_word(DBGMCU_IDCODE);

  for (int i = 0; i < ARRAY_SIZE(devices); i++)
  {
    if (devices[i].idcode != (idcode & DEVICE_ID_MASK))
      continue;

    verbose("Target: %s\n", devices[i].name);

    target_device = devices[i];
    target_options = *options;

    uint32_t flash_size = (dap_read_word(FLASH_SIZE_REG) & FLASH_SIZE_REG_MASK) * FLASH_SIZE_REG_MULT;

    target_dbank     = (dap_read_word(FLASH_OPTR) & FLASH_OPTR_DBANK);
    target_page_size = target_dbank ? FLASH_PAGE_SIZE_2 : FLASH_PAGE_SIZE_1;
    target_row_size  = target_dbank ? FLASH_ROW_SIZE_2 : FLASH_ROW_SIZE_1;

    verbose("Flash size: %d bytes (%s bank mode)\n", flash_size, target_dbank ? "dual" : "single");

    target_check_options(&target_options, flash_size, target_page_size);

    bool locked = (0xaa != (dap_read_word(FLASH_OPTR) & FLASH_OPTR_RDP_MASK));

    if (locked && !options->unlock)
      error_exit("target is locked, unlock is necessary");

    dap_write_word(FLASH_KEYR, FLASH_KEYR_KEY1);
    dap_write_word(FLASH_KEYR, FLASH_KEYR_KEY2);
    dap_write_word(FLASH_OPTKEYR, FLASH_OPTKEYR_KEY1);
    dap_write_word(FLASH_OPTKEYR, FLASH_OPTKEYR_KEY2);
    dap_write_word(FLASH_CR, 0);

    check(0 == dap_read_word(FLASH_CR), "Failed to unlock the flash for write operation. Try to power cycle the target.");

    return;
  }

  error_exit("unknown target device (DBGMCU_IDCODE = 0x%08x)", idcode);
}

//-----------------------------------------------------------------------------
static void target_deselect(void)
{
  dap_write_word(DEMCR, 0);
  dap_write_word(AIRCR, AIRCR_VECTKEY | AIRCR_SYSRESETREQ);

  target_free_options(&target_options);
}

//-----------------------------------------------------------------------------
static void target_erase(void)
{
  dap_write_word(FLASH_CR, FLASH_CR_MER1 | FLASH_CR_MER2);
  dap_write_word(FLASH_CR, FLASH_CR_MER1 | FLASH_CR_MER2 | FLASH_CR_STRT);
  flash_wait_done();
  dap_write_word(FLASH_CR, 0);
}

//-----------------------------------------------------------------------------
static void target_lock(void)
{
  error_exit("target_lock() is not implemented yet");
}

//-----------------------------------------------------------------------------
static void target_unlock(void)
{
  error_exit("target_unlock() is not implemented yet");
}

//-----------------------------------------------------------------------------
static void target_program(void)
{
  uint32_t addr = FLASH_ADDR + target_options.offset;
  uint32_t offs = 0;
  uint32_t start_page, number_of_pages;
  uint8_t *buf = target_options.file_data;
  uint32_t size = target_options.file_size;

  start_page = target_options.offset / target_page_size;
  number_of_pages = (size + target_page_size - 1) / target_page_size;

  for (uint32_t page = 0; page < number_of_pages; page++)
  {
    // Erase Page
    dap_write_word(FLASH_CR, FLASH_CR_PER | FLASH_CR_PNB(start_page + page));
    dap_write_word(FLASH_CR, FLASH_CR_PER | FLASH_CR_PNB(start_page + page) | FLASH_CR_STRT);
    flash_wait_done();

    dap_write_word(FLASH_CR, FLASH_CR_PG);

    // Program Page
    for (int i = 0; i < (target_page_size / target_row_size); i++)
    {
      dap_write_block(addr, &buf[offs], target_row_size);
      addr += target_row_size;
      offs += target_row_size;
    }

    flash_wait_done();

    verbose(".");
  }

  dap_write_word(FLASH_CR, 0);
}

//-----------------------------------------------------------------------------
static void target_verify(void)
{
  uint32_t addr = FLASH_ADDR + target_options.offset;
  int block_size;
  uint32_t offs = 0;
  uint8_t *bufb;
  uint8_t *bufa = target_options.file_data;
  int size = target_options.file_size;

  bufb = buf_alloc(target_row_size);

  while (size)
  {
    dap_read_block(addr, bufb, target_row_size);

    block_size = (size > target_row_size) ? target_row_size : size;

    for (int i = 0; i < block_size; i++)
    {
      if (bufa[offs + i] != bufb[i])
      {
        verbose("\nat address 0x%x expected 0x%02x, read 0x%02x\n",
            addr + i, bufa[offs + i], bufb[i]);
        buf_free(bufb);
        error_exit("verification failed");
      }
    }

    addr += target_row_size;
    offs += target_row_size;
    size -= block_size;

    if (0 == (offs % target_page_size))
      verbose(".");
  }

  buf_free(bufb);
}

//-----------------------------------------------------------------------------
static void target_read(void)
{
  uint32_t addr = FLASH_ADDR + target_options.offset;
  uint32_t offs = 0;
  uint8_t *buf = target_options.file_data;
  uint32_t size = target_options.size;

  while (size)
  {
    dap_read_block(addr, &buf[offs], target_row_size);

    addr += target_row_size;
    offs += target_row_size;
    size -= target_row_size;

    if (0 == (offs % target_page_size))
      verbose(".");
  }

  save_file(target_options.name, buf, target_options.size);
}

//-----------------------------------------------------------------------------
static uint32_t section_to_reg(int section)
{
  uint32_t regs[] = { FLASH_OPTR, FLASH_PCROP1SR, FLASH_PCROP1ER, FLASH_WRP1AR, FLASH_WRP1BR,
      FLASH_SEC1R, FLASH_PCROP2SR, FLASH_PCROP2ER, FLASH_WRP2AR, FLASH_WRP2BR, FLASH_SEC2R };

  if (section < ARRAY_SIZE(regs))
    return regs[section];
  else
    return 0;
}

//-----------------------------------------------------------------------------
static int target_fuse_read(int section, uint8_t *data)
{
  uint32_t reg = section_to_reg(section);

  if (0 == reg)
    return 0;

  uint32_t value = dap_read_word(reg);

  data[0] = value;
  data[1] = value >> 8;
  data[2] = value >> 16;
  data[3] = value >> 24;

  return OPTIONS_SIZE;
}

//-----------------------------------------------------------------------------
static void target_fuse_write(int section, uint8_t *data)
{
  uint32_t reg = section_to_reg(section);

  if (0 == reg)
    error_exit("internal: incorrect section index in target_fuse_write()");

  uint32_t value = (data[3] << 24) | (data[2] << 16) | (data[1] << 8) | data[0];

  dap_write_word_req(reg, value);
  dap_write_word_req(FLASH_CR, FLASH_CR_OPTSTRT);
  dap_transfer();

  flash_wait_done();
}

//-----------------------------------------------------------------------------
static char *target_enumerate(int i)
{
  if (i < ARRAY_SIZE(devices))
    return devices[i].family;

  return NULL;
}

//-----------------------------------------------------------------------------
static char target_help[] =
  "Fuses:\n"
  "  The option bytes are represented by the following sections (32-bits each):\n"
  "     0 - OPTR     (Option register)\n"
  "     1 - PCROP1SR (Bank 1 PCROP start address register)\n"
  "     2 - PCROP1ER (Bank 1 PCROP end address register)\n"
  "     3 - WRP1AR   (Bank 1 WRP area A address register)\n"
  "     4 - WRP1BR   (Bank 1 WRP area B address register)\n"
  "     5 - SEC1R    (Bank 1 securable memory area)\n"
  "     6 - PCROP2SR (Bank 2 PCROP start address register)\n"
  "     7 - PCROP2ER (Bank 2 PCROP end address register)\n"
  "     8 - WRP2AR   (Bank 2 WRP area A address register)\n"
  "     9 - WRP2BR   (Bank 2 WRP area B address register)\n"
  "    10 - SEC2R    (Bank 2 securable memory area)\n";

//-----------------------------------------------------------------------------
target_ops_t target_st_stm32g4_ops =
{
  .select    = target_select,
  .deselect  = target_deselect,
  .erase     = target_erase,
  .lock      = target_lock,
  .unlock    = target_unlock,
  .program   = target_program,
  .verify    = target_verify,
  .read      = target_read,
  .fread     = target_fuse_read,
  .fwrite    = target_fuse_write,
  .enumerate = target_enumerate,
  .help      = target_help,
};


