#include "pid_flash.h"
#include "stm32f7xx_hal.h"
#include <string.h>

bool pid_flash_load(float *kp, float *ki, float *kd)
{
    const pid_flash_t *p = (const pid_flash_t *)PID_FLASH_STORE_ADDR;

    if (p->magic != PID_FLASH_MAGIC) return false;

    /* Basic sanity bounds -- same as JLINK_CMD_PID_SET handler */
    if (p->kp < 0.0f || p->kp > 500.0f) return false;
    if (p->ki < 0.0f || p->ki > 50.0f)  return false;
    if (p->kd < 0.0f || p->kd > 50.0f)  return false;

    *kp = p->kp;
    *ki = p->ki;
    *kd = p->kd;
    return true;
}

bool pid_flash_save(float kp, float ki, float kd)
{
    HAL_StatusTypeDef rc;

    /* Unlock flash */
    rc = HAL_FLASH_Unlock();
    if (rc != HAL_OK) return false;

    /* Erase sector 7 */
    FLASH_EraseInitTypeDef erase = {
        .TypeErase    = FLASH_TYPEERASE_SECTORS,
        .Sector       = PID_FLASH_SECTOR,
        .NbSectors    = 1,
        .VoltageRange = PID_FLASH_SECTOR_VOLTAGE,
    };
    uint32_t sector_error = 0;
    rc = HAL_FLASHEx_Erase(&erase, &sector_error);
    if (rc != HAL_OK || sector_error != 0xFFFFFFFFUL) {
        HAL_FLASH_Lock();
        return false;
    }

    /* Build record */
    pid_flash_t rec;
    memset(&rec, 0xFF, sizeof(rec));
    rec.magic = PID_FLASH_MAGIC;
    rec.kp    = kp;
    rec.ki    = ki;
    rec.kd    = kd;

    /* Write 64 bytes as 16 x 32-bit words */
    const uint32_t *src  = (const uint32_t *)&rec;
    uint32_t        addr = PID_FLASH_STORE_ADDR;
    for (uint8_t i = 0; i < sizeof(rec) / 4u; i++) {
        rc = HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, addr, src[i]);
        if (rc != HAL_OK) {
            HAL_FLASH_Lock();
            return false;
        }
        addr += 4u;
    }

    HAL_FLASH_Lock();

    /* Verify readback */
    const pid_flash_t *stored = (const pid_flash_t *)PID_FLASH_STORE_ADDR;
    return (stored->magic == PID_FLASH_MAGIC &&
            stored->kp == kp &&
            stored->ki == ki &&
            stored->kd == kd);
}

/* ---- Helper: write N bytes to flash as 32-bit words ---- */
static bool flash_write_words(uint32_t addr, const void *src_buf, uint32_t byte_len)
{
    const uint32_t *w = (const uint32_t *)src_buf;
    for (uint32_t i = 0; i < byte_len / 4u; i++) {
        if (HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, addr, w[i]) != HAL_OK)
            return false;
        addr += 4u;
    }
    return true;
}

/* ---- pid_flash_load_schedule() ---- */
bool pid_flash_load_schedule(pid_sched_entry_t *out_entries, uint8_t *out_n)
{
    const pid_sched_flash_t *p = (const pid_sched_flash_t *)PID_SCHED_FLASH_ADDR;

    if (p->magic != PID_SCHED_MAGIC) return false;
    if (p->num_bands == 0u || p->num_bands > PID_SCHED_MAX_BANDS) return false;

    memcpy(out_entries, p->bands, p->num_bands * sizeof(pid_sched_entry_t));
    *out_n = p->num_bands;
    return true;
}

/* ---- pid_flash_save_all() ---- */
bool pid_flash_save_all(float kp_single, float ki_single, float kd_single,
                        const pid_sched_entry_t *entries, uint8_t num_bands)
{
    HAL_StatusTypeDef rc;

    if (num_bands == 0u || num_bands > PID_SCHED_MAX_BANDS) return false;

    rc = HAL_FLASH_Unlock();
    if (rc != HAL_OK) return false;

    /* Erase sector 7 -- one erase covers both records */
    FLASH_EraseInitTypeDef erase = {
        .TypeErase    = FLASH_TYPEERASE_SECTORS,
        .Sector       = PID_FLASH_SECTOR,
        .NbSectors    = 1,
        .VoltageRange = PID_FLASH_SECTOR_VOLTAGE,
    };
    uint32_t sector_error = 0u;
    rc = HAL_FLASHEx_Erase(&erase, &sector_error);
    if (rc != HAL_OK || sector_error != 0xFFFFFFFFUL) {
        HAL_FLASH_Lock();
        return false;
    }

    /* Build and write schedule record at PID_SCHED_FLASH_ADDR */
    pid_sched_flash_t srec;
    memset(&srec, 0xFF, sizeof(srec));
    srec.magic     = PID_SCHED_MAGIC;
    srec.num_bands = num_bands;
    srec.flags     = 0u;
    srec._pad0[0]  = 0u;
    srec._pad0[1]  = 0u;
    memcpy(srec.bands, entries, num_bands * sizeof(pid_sched_entry_t));

    if (!flash_write_words(PID_SCHED_FLASH_ADDR, &srec, sizeof(srec))) {
        HAL_FLASH_Lock();
        return false;
    }

    /* Build and write single-PID record at PID_FLASH_STORE_ADDR */
    pid_flash_t prec;
    memset(&prec, 0xFF, sizeof(prec));
    prec.magic = PID_FLASH_MAGIC;
    prec.kp    = kp_single;
    prec.ki    = ki_single;
    prec.kd    = kd_single;

    if (!flash_write_words(PID_FLASH_STORE_ADDR, &prec, sizeof(prec))) {
        HAL_FLASH_Lock();
        return false;
    }

    HAL_FLASH_Lock();

    /* Verify both records */
    const pid_sched_flash_t *sv = (const pid_sched_flash_t *)PID_SCHED_FLASH_ADDR;
    const pid_flash_t       *pv = (const pid_flash_t *)PID_FLASH_STORE_ADDR;

    return (sv->magic == PID_SCHED_MAGIC && sv->num_bands == num_bands &&
            pv->magic == PID_FLASH_MAGIC && pv->kp == kp_single &&
            pv->ki == ki_single && pv->kd == kd_single);
}