#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 × 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 arbitrary bytes as 32-bit words ---- */
/*
 * Writes 'len' bytes from 'src' to flash at 'addr'.
 * len must be a multiple of 4.  Flash must already be unlocked.
 * Returns HAL_OK on success, or first failure status.
 */
static HAL_StatusTypeDef flash_write_words(uint32_t addr,
                                           const void *src,
                                           uint32_t len)
{
    const uint32_t *p = (const uint32_t *)src;
    HAL_StatusTypeDef rc = HAL_OK;
    for (uint32_t i = 0; i < len / 4u; i++) {
        rc = HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, addr, p[i]);
        if (rc != HAL_OK) return rc;
        addr += 4u;
    }
    return HAL_OK;
}

/* ---- 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;

    *out_n = p->num_bands;
    for (uint8_t i = 0; i < p->num_bands; i++) {
        out_entries[i] = p->bands[i];
    }
    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)
{
    if (num_bands == 0u || num_bands > PID_SCHED_MAX_BANDS) return false;

    HAL_StatusTypeDef rc;

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

    /* Single erase of sector 7 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 = 0;
    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;
    for (uint8_t i = 0; i < num_bands; i++) {
        srec.bands[i] = entries[i];
    }

    rc = flash_write_words(PID_SCHED_FLASH_ADDR, &srec, sizeof(srec));
    if (rc != HAL_OK) {
        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;

    rc = flash_write_words(PID_FLASH_STORE_ADDR, &prec, sizeof(prec));
    if (rc != HAL_OK) {
        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);
}