/*
 * Baby Mobile Audio Board v2 — nRF52840 Firmware
 * 
 * Platform: Seeed XIAO nRF52840 (or MDBT50Q module)
 * Framework: Adafruit nRF52 Arduino Core
 * 
 * Features:
 *   - USB Mass Storage: drag-and-drop .raw audio files
 *   - BLE: wireless audio upload from phone/Web Bluetooth
 *   - 8-bit 8kHz PWM audio playback from IS25LP128F SPI flash
 *   - DC motor PWM speed control
 *   - 8 button inputs with debounce
 *   - Deep sleep with button wake (~10µA)
 *   - Auto-shutoff timer
 * 
 * Audio format: unsigned 8-bit PCM, 8000 Hz, mono
 * Convert: ffmpeg -i song.mp3 -ar 8000 -ac 1 -f u8 -acodec pcm_u8 song.raw
 * 
 * Board setup in Arduino IDE:
 *   Board: "Seeed XIAO nRF52840"  (or "Adafruit Feather nRF52840")
 *   Install: Seeed nRF52 Boards via Board Manager
 *   
 * Wiring (XIAO pin names):
 *   D0  (P0.02) - Button 1
 *   D1  (P0.03) - Button 2
 *   D2  (P0.28) - Button 3
 *   D3  (P0.29) - Button 4
 *   D4  (P0.04) - Button 5
 *   D5  (P0.05) - Button 6
 *   D6  (P1.11) - Button 7
 *   D7  (P1.12) - Button 8
 *   D8  (P0.07) - SPI SCK  → Flash pin 6
 *   D9  (P0.06) - SPI MISO → Flash pin 2
 *   D10 (P0.05) - SPI MOSI → Flash pin 5
 *   A0  (P0.02) - Audio PWM output → R+C LPF → PAM8302A
 *   A1  (P0.03) - Amp ~SD (HIGH=on)
 *   A2  (P0.28) - Motor PWM → MOSFET gate
 *   A3  (P0.29) - LED output
 *   A4  (P0.04) - Flash ~CS
 *   A5  (P0.05) - (spare)
 * 
 * NOTE: Pin assignments above are illustrative. Adjust to your actual 
 * XIAO pinout and PCB routing. The XIAO has 11 usable GPIOs on the 
 * edge pins plus additional pads on the bottom.
 */

#include <SPI.h>
#include <Adafruit_TinyUSB.h>
#include <bluefruit.h>

// ============================================================
// POC MODE: stream audio from internal LittleFS instead of SPI flash.
// Upload tracks via BLE (CMD 0x02 [track_num] + data packets).
// WAV files are accepted — 44-byte PCM header is stripped on receipt.
// Disable this define to revert to full SPI flash mode.
// ============================================================
#define POC_INTERNAL_FLASH

#ifdef POC_INTERNAL_FLASH
#include <Adafruit_LittleFS.h>
#include <InternalFileSystem.h>
using namespace Adafruit_LittleFS_Namespace;
#endif

// ============================================================
// PIN DEFINITIONS — adjust for your PCB
// ============================================================

// Using raw nRF52840 GPIO numbers for flexibility.
// Map these to your actual PCB connections.

// Audio
#define PIN_AUDIO_PWM    8    // Any PWM-capable pin
#define PIN_AMP_SD       14   // PAM8302A shutdown (HIGH=enabled)

// Motor
#define PIN_MOTOR_PWM    15   // PWM to MOSFET gate

// SPI Flash
#define PIN_FLASH_CS     2    // Flash chip select
// SPI MOSI/MISO/SCK use default SPI pins

// Buttons (directly to GPIO, active LOW with internal pull-up)
#define PIN_BTN1         0
#define PIN_BTN2         1
#define PIN_BTN3         5
#define PIN_BTN4         3
/*
#define PIN_BTN5         4
#define PIN_BTN6         5
#define PIN_BTN7         3
#define PIN_BTN8         4
*/

// PIN_LED is 11 (LED_R)
// LED_B is 13/12
// LED_G is 12/13

const uint8_t BTN_PINS[] = {
    PIN_BTN1, PIN_BTN2, PIN_BTN3, PIN_BTN4,
    /*PIN_BTN5, PIN_BTN6, PIN_BTN7, PIN_BTN8*/
};
#define NUM_BUTTONS 4

// ============================================================
// FLASH CONSTANTS
// ============================================================

#define FLASH_SIZE        (16UL * 1024UL * 1024UL)  // 16MB
#define FLASH_SECTOR      4096
#define FLASH_PAGE        256
#define FLASH_BLOCK_64K   (64UL * 1024UL)

// IS25LP128F / W25Q128 compatible commands
#define CMD_READ          0x03
#define CMD_FAST_READ     0x0B
#define CMD_WRITE_EN      0x06
#define CMD_PAGE_PROG     0x02
#define CMD_SECT_ERASE    0x20
#define CMD_BLOCK_ERASE   0xD8
#define CMD_CHIP_ERASE    0xC7
#define CMD_READ_SR1      0x05
#define CMD_JEDEC_ID      0x9F
#define CMD_POWER_DOWN    0xB9
#define CMD_WAKE          0xAB

// ============================================================
// TRACK TABLE FORMAT
// ============================================================
// Sector 0 (first 4KB) holds the track index:
//   [0]     uint8_t  magic = 0xBB (indicates valid table)
//   [1]     uint8_t  num_tracks
//   [2..3]  reserved
//   [4..]   Track entries, 12 bytes each:
//             uint32_t start_addr (big-endian)
//             uint32_t byte_length (big-endian)
//             uint8_t  sample_rate_khz (8 = 8kHz)
//             uint8_t  bits (8 = 8-bit unsigned)
//             uint16_t reserved
// Audio data starts at AUDIO_START (sector 1 = 0x1000)

#define TRACK_TABLE_ADDR  0x000000
#define AUDIO_START_ADDR  0x001000
#define TABLE_MAGIC       0xBB
#define MAX_TRACKS        32
#define TRACK_ENTRY_SIZE  12

// ============================================================
// AUDIO CONFIG
// ============================================================

#define SAMPLE_RATE       8000
#define AUDIO_BUF_SIZE    512   // larger buffer = fewer SPI transactions
#define SILENCE           128   // 0x80 = center for unsigned 8-bit

// ============================================================
// USB MASS STORAGE
// ============================================================

Adafruit_USBD_MSC usb_msc;

// Simple FAT12 filesystem on the SPI flash
// We present the entire 16MB flash as a USB drive.
// The host can write raw files; firmware scans for audio on boot.

// MSC callbacks
int32_t msc_read_cb(uint32_t lba, void* buffer, uint32_t bufsize);
int32_t msc_write_cb(uint32_t lba, uint8_t* buffer, uint32_t bufsize);
void    msc_flush_cb(void);
bool    msc_start_stop_cb(uint8_t power_condition, bool start, bool load_eject);

#define MSC_BLOCK_SIZE  512
#define MSC_BLOCK_COUNT (FLASH_SIZE / MSC_BLOCK_SIZE)

// ============================================================
// BLE SERVICE
// ============================================================

BLEService        audioSvc  = BLEService("12340001-0000-1000-8000-00805f9b34fb");
BLECharacteristic audioCmd  = BLECharacteristic("12340002-0000-1000-8000-00805f9b34fb");
BLECharacteristic audioData = BLECharacteristic("12340003-0000-1000-8000-00805f9b34fb");
BLECharacteristic audioStat = BLECharacteristic("12340004-0000-1000-8000-00805f9b34fb");

// BLE upload state
volatile bool     g_bleUploading = false;
volatile uint32_t g_bleWriteAddr = 0;
volatile uint32_t g_bleWriteLen = 0;

#ifdef POC_INTERNAL_FLASH
File     g_pocFile(InternalFS);   // open file handle (read or write)
uint8_t  g_pocWriteTrack = 0;     // track slot being written via BLE
bool     g_pocSkipHeader = false; // strip WAV header from first data packet
#endif

// ============================================================
// GLOBAL STATE
// ============================================================

volatile bool     g_playing = false;
volatile bool     g_motorOn = false;
uint8_t           g_motorSpeed = 160;

uint8_t  g_numTracks = 0;
uint32_t g_trackStart[MAX_TRACKS];
uint32_t g_trackLen[MAX_TRACKS];
uint8_t  g_currentTrack = 0;

// Audio double-buffer
uint8_t  g_audioBuf[2][AUDIO_BUF_SIZE];
volatile uint8_t  g_activeBuf = 0;
volatile uint16_t g_bufPos = 0;
volatile bool     g_bufReady[2] = {false, false};
uint32_t g_nextReadAddr = 0;
uint32_t g_playEnd = 0;

// Timers
unsigned long g_lastActivity = 0;
#define AUTO_OFF_MS  (30UL * 60UL * 1000UL)  // 30 min auto-shutoff
#define IDLE_SLEEP_MS (60UL * 1000UL)         // 1 min idle → sleep

// USB connected flag
volatile bool g_usbConnected = false;

#ifdef POC_INTERNAL_FLASH
static void pocFilename(uint8_t n, char *buf) {  // buf must be >=16 bytes
    snprintf(buf, 16, "/track%d.raw", n);
}
#endif

// ============================================================
// SPI FLASH DRIVER
// ============================================================

SPISettings flashSPI(8000000, MSBFIRST, SPI_MODE0);  // 8MHz

void flashSelect()   { digitalWrite(PIN_FLASH_CS, LOW); }
void flashDeselect() { digitalWrite(PIN_FLASH_CS, HIGH); }

void flashWake() {
    flashSelect();
    SPI.transfer(CMD_WAKE);
    flashDeselect();
    delayMicroseconds(5);
}

void flashSleep() {
    flashSelect();
    SPI.transfer(CMD_POWER_DOWN);
    flashDeselect();
}

uint32_t flashReadJEDEC() {
    SPI.beginTransaction(flashSPI);
    flashSelect();
    SPI.transfer(CMD_JEDEC_ID);
    uint32_t id = (uint32_t)SPI.transfer(0) << 16;
    id |= (uint32_t)SPI.transfer(0) << 8;
    id |= SPI.transfer(0);
    flashDeselect();
    SPI.endTransaction();
    return id;
}

void flashReadBytes(uint32_t addr, uint8_t *buf, uint32_t len) {
    SPI.beginTransaction(flashSPI);
    flashSelect();
    SPI.transfer(CMD_READ);
    SPI.transfer((addr >> 16) & 0xFF);
    SPI.transfer((addr >> 8) & 0xFF);
    SPI.transfer(addr & 0xFF);
    for (uint32_t i = 0; i < len; i++) {
        buf[i] = SPI.transfer(0);
    }
    flashDeselect();
    SPI.endTransaction();
}

void flashWaitBusy() {
    SPI.beginTransaction(flashSPI);
    flashSelect();
    SPI.transfer(CMD_READ_SR1);
    while (SPI.transfer(0) & 0x01) { /* spin */ }
    flashDeselect();
    SPI.endTransaction();
}

void flashWriteEnable() {
    SPI.beginTransaction(flashSPI);
    flashSelect();
    SPI.transfer(CMD_WRITE_EN);
    flashDeselect();
    SPI.endTransaction();
}

void flashEraseSector(uint32_t addr) {
    flashWriteEnable();
    SPI.beginTransaction(flashSPI);
    flashSelect();
    SPI.transfer(CMD_SECT_ERASE);
    SPI.transfer((addr >> 16) & 0xFF);
    SPI.transfer((addr >> 8) & 0xFF);
    SPI.transfer(addr & 0xFF);
    flashDeselect();
    SPI.endTransaction();
    flashWaitBusy();
}

void flashEraseBlock64K(uint32_t addr) {
    flashWriteEnable();
    SPI.beginTransaction(flashSPI);
    flashSelect();
    SPI.transfer(CMD_BLOCK_ERASE);
    SPI.transfer((addr >> 16) & 0xFF);
    SPI.transfer((addr >> 8) & 0xFF);
    SPI.transfer(addr & 0xFF);
    flashDeselect();
    SPI.endTransaction();
    flashWaitBusy();
}

void flashPageProgram(uint32_t addr, const uint8_t *data, uint16_t len) {
    flashWriteEnable();
    SPI.beginTransaction(flashSPI);
    flashSelect();
    SPI.transfer(CMD_PAGE_PROG);
    SPI.transfer((addr >> 16) & 0xFF);
    SPI.transfer((addr >> 8) & 0xFF);
    SPI.transfer(addr & 0xFF);
    for (uint16_t i = 0; i < len; i++) {
        SPI.transfer(data[i]);
    }
    flashDeselect();
    SPI.endTransaction();
    flashWaitBusy();
}


// ============================================================
// TRACK TABLE
// ============================================================

void loadTrackTable() {
#ifdef POC_INTERNAL_FLASH
    g_numTracks = 0;
    char fname[16];
    for (uint8_t i = 0; i < MAX_TRACKS; i++) {
        pocFilename(i, fname);
        File f(InternalFS);
        if (!f.open(fname, FILE_O_READ)) break;
        g_trackLen[i] = f.size();
        f.close();
        g_numTracks = i + 1;
    }
    Serial.print("Found ");
    Serial.print(g_numTracks);
    Serial.println(" tracks in internal flash");
#else
    uint8_t header[4];
    flashReadBytes(TRACK_TABLE_ADDR, header, 4);

    if (header[0] != TABLE_MAGIC) {
        g_numTracks = 0;
        Serial.println("No track table found (flash may be empty)");
        return;
    }

    g_numTracks = header[1];
    if (g_numTracks > MAX_TRACKS) g_numTracks = MAX_TRACKS;

    uint8_t entry[TRACK_ENTRY_SIZE];
    for (uint8_t i = 0; i < g_numTracks; i++) {
        flashReadBytes(TRACK_TABLE_ADDR + 4 + i * TRACK_ENTRY_SIZE, entry, TRACK_ENTRY_SIZE);
        g_trackStart[i] = ((uint32_t)entry[0] << 24) | ((uint32_t)entry[1] << 16) |
                          ((uint32_t)entry[2] << 8)  | entry[3];
        g_trackLen[i]   = ((uint32_t)entry[4] << 24) | ((uint32_t)entry[5] << 16) |
                          ((uint32_t)entry[6] << 8)  | entry[7];
    }

    Serial.print("Loaded ");
    Serial.print(g_numTracks);
    Serial.println(" tracks from flash");
    for (uint8_t i = 0; i < g_numTracks; i++) {
        Serial.print("  Track ");
        Serial.print(i + 1);
        Serial.print(": addr=0x");
        Serial.print(g_trackStart[i], HEX);
        Serial.print(", ");
        Serial.print(g_trackLen[i] / SAMPLE_RATE);
        Serial.println("s");
    }
#endif
}

void writeTrackTable() {
#ifdef POC_INTERNAL_FLASH
    return;  // no table needed; LittleFS files are the store
#else
    // Erase sector 0
    flashEraseSector(0);

    // Build table in RAM
    uint8_t table[4 + MAX_TRACKS * TRACK_ENTRY_SIZE];
    memset(table, 0xFF, sizeof(table));
    table[0] = TABLE_MAGIC;
    table[1] = g_numTracks;
    table[2] = 0;
    table[3] = 0;

    for (uint8_t i = 0; i < g_numTracks; i++) {
        uint8_t *e = &table[4 + i * TRACK_ENTRY_SIZE];
        e[0] = (g_trackStart[i] >> 24) & 0xFF;
        e[1] = (g_trackStart[i] >> 16) & 0xFF;
        e[2] = (g_trackStart[i] >> 8) & 0xFF;
        e[3] = g_trackStart[i] & 0xFF;
        e[4] = (g_trackLen[i] >> 24) & 0xFF;
        e[5] = (g_trackLen[i] >> 16) & 0xFF;
        e[6] = (g_trackLen[i] >> 8) & 0xFF;
        e[7] = g_trackLen[i] & 0xFF;
        e[8] = 8;   // sample rate kHz
        e[9] = 8;   // bits
        e[10] = 0;
        e[11] = 0;
    }

    // Write in pages
    uint16_t total = 4 + g_numTracks * TRACK_ENTRY_SIZE;
    for (uint16_t offset = 0; offset < total; offset += FLASH_PAGE) {
        uint16_t chunk = min((uint16_t)FLASH_PAGE, (uint16_t)(total - offset));
        flashPageProgram(offset, &table[offset], chunk);
    }
#endif
}


// ============================================================
// USB MASS STORAGE CALLBACKS
// ============================================================

int32_t msc_read_cb(uint32_t lba, void* buffer, uint32_t bufsize) {
    uint32_t addr = lba * MSC_BLOCK_SIZE;
    flashReadBytes(addr, (uint8_t*)buffer, bufsize);
    return bufsize;
}

int32_t msc_write_cb(uint32_t lba, uint8_t* buffer, uint32_t bufsize) {
    uint32_t addr = lba * MSC_BLOCK_SIZE;

    // Erase sector if we're at a sector boundary
    if ((addr % FLASH_SECTOR) == 0) {
        flashEraseSector(addr);
    }

    // Write in page-sized chunks
    uint32_t written = 0;
    while (written < bufsize) {
        uint16_t pageOffset = (addr + written) % FLASH_PAGE;
        uint16_t chunk = min((uint16_t)(FLASH_PAGE - pageOffset), (uint16_t)(bufsize - written));
        flashPageProgram(addr + written, buffer + written, chunk);
        written += chunk;
    }

    return bufsize;
}

void msc_flush_cb(void) {
    // After USB write completes, reload the track table
    // (host may have written a new FAT filesystem)
    // We scan for audio data on eject instead
}

bool msc_start_stop_cb(uint8_t power_condition, bool start, bool load_eject) {
    if (!start && load_eject) {
        // Host ejected the drive — rescan for tracks
        Serial.println("USB ejected, rescanning tracks...");
        loadTrackTable();
    }
    return true;
}


// ============================================================
// BLE SETUP
// ============================================================

void ble_connect_cb(uint16_t conn_handle) {
    Serial.println("BLE connected");
    g_lastActivity = millis();
}

void ble_disconnect_cb(uint16_t conn_handle, uint8_t reason) {
    Serial.println("BLE disconnected");
    g_bleUploading = false;
}

// BLE command characteristic: receives commands
//   CMD 0x01 [num_tracks] [track_entries...] = write track table
//   CMD 0x02 [addr_3bytes] = start writing audio at address
//   CMD 0x03 = finish upload, reload tracks
//   CMD 0x04 [track_num] = play track
//   CMD 0x05 = stop playback
void audioCmd_write_cb(uint16_t conn_handle, BLECharacteristic* chr,
                       uint8_t* data, uint16_t len) {
    if (len < 1) return;
    g_lastActivity = millis();

    switch (data[0]) {
        case 0x01: // Write track table
            if (len >= 2) {
                g_numTracks = data[1];
                if (g_numTracks > MAX_TRACKS) g_numTracks = MAX_TRACKS;
                // Parse track entries from data[2..]
                for (uint8_t i = 0; i < g_numTracks && (2 + i * 8 + 7) < len; i++) {
                    uint8_t *e = &data[2 + i * 8];
                    g_trackStart[i] = ((uint32_t)e[0] << 24) | ((uint32_t)e[1] << 16) |
                                      ((uint32_t)e[2] << 8) | e[3];
                    g_trackLen[i] = ((uint32_t)e[4] << 24) | ((uint32_t)e[5] << 16) |
                                    ((uint32_t)e[6] << 8) | e[7];
                }
                writeTrackTable();
                Serial.println("BLE: Track table updated");
            }
            break;

        case 0x02: // Start audio write
#ifdef POC_INTERNAL_FLASH
            if (len >= 2) {
                g_pocWriteTrack = data[1];
                if (g_pocFile) g_pocFile.close();
                char fname[16];
                pocFilename(g_pocWriteTrack, fname);
                if (!g_pocFile.open(fname, FILE_O_WRITE)) {
                    Serial.print("BLE: cannot open ");
                    Serial.println(fname);
                    break;
                }
                g_bleUploading = true;
                g_pocSkipHeader = true;
                g_bleWriteLen = 0;
                Serial.print("BLE: Start write track ");
                Serial.println(g_pocWriteTrack);
            }
#else
            if (len >= 4) {
                g_bleWriteAddr = ((uint32_t)data[1] << 16) | ((uint32_t)data[2] << 8) | data[3];
                g_bleUploading = true;
                // Erase 64K block at target
                flashEraseBlock64K(g_bleWriteAddr & ~(FLASH_BLOCK_64K - 1));
                Serial.print("BLE: Start write at 0x");
                Serial.println(g_bleWriteAddr, HEX);
            }
#endif
            break;

        case 0x03: // Finish upload
#ifdef POC_INTERNAL_FLASH
            if (g_pocFile) g_pocFile.close();
            g_pocSkipHeader = false;
#endif
            g_bleUploading = false;
            loadTrackTable();
            Serial.println("BLE: Upload complete");
            break;

        case 0x04: // Play track
            if (len >= 2 && data[1] < g_numTracks) {
                audioStart(data[1]);
            }
            break;

        case 0x05: // Stop
            audioStop();
            motorStop();
            break;
    }
}

// BLE data characteristic: receives raw audio bytes for upload
void audioData_write_cb(uint16_t conn_handle, BLECharacteristic* chr,
                        uint8_t* data, uint16_t len) {
    if (!g_bleUploading) return;
    g_lastActivity = millis();

#ifdef POC_INTERNAL_FLASH
    uint8_t *src = data;
    uint16_t srcLen = len;

    // Strip 44-byte WAV header from first packet if present
    if (g_pocSkipHeader) {
        g_pocSkipHeader = false;
        if (srcLen >= 4 && src[0]=='R' && src[1]=='I' && src[2]=='F' && src[3]=='F') {
            if (srcLen > 44) {
                src    += 44;
                srcLen -= 44;
            } else {
                // Header spans packets — drop whole packet (document: send raw PCM instead)
                return;
            }
        }
    }

    g_pocFile.write(src, srcLen);
    g_bleWriteLen += srcLen;
#else
    // Erase new sectors as we cross boundaries
    uint32_t endAddr = g_bleWriteAddr + len;
    uint32_t currentSector = g_bleWriteAddr / FLASH_SECTOR;
    uint32_t endSector = (endAddr - 1) / FLASH_SECTOR;
    for (uint32_t s = currentSector + 1; s <= endSector; s++) {
        flashEraseSector(s * FLASH_SECTOR);
    }

    // Write data page by page
    uint16_t written = 0;
    while (written < len) {
        uint16_t pageOff = (g_bleWriteAddr + written) % FLASH_PAGE;
        uint16_t chunk = min((uint16_t)(FLASH_PAGE - pageOff), (uint16_t)(len - written));
        flashPageProgram(g_bleWriteAddr + written, data + written, chunk);
        written += chunk;
    }

    g_bleWriteAddr += len;
    g_bleWriteLen += len;
#endif

    // Update status characteristic with bytes written
    uint8_t stat[4];
    stat[0] = (g_bleWriteLen >> 24) & 0xFF;
    stat[1] = (g_bleWriteLen >> 16) & 0xFF;
    stat[2] = (g_bleWriteLen >> 8) & 0xFF;
    stat[3] = g_bleWriteLen & 0xFF;
    audioStat.write(stat, 4);
    audioStat.notify(stat, 4);
}

void setupBLE() {
    Bluefruit.begin();
    Bluefruit.setName("BabyMobile");
    Bluefruit.setTxPower(0);  // 0 dBm — save power, short range is fine
    Bluefruit.Periph.setConnectCallback(ble_connect_cb);
    Bluefruit.Periph.setDisconnectCallback(ble_disconnect_cb);

    // Audio service
    audioSvc.begin();

    // Command characteristic (write)
    audioCmd.setProperties(CHR_PROPS_WRITE);
    audioCmd.setPermission(SECMODE_OPEN, SECMODE_OPEN);
    audioCmd.setMaxLen(240);
    audioCmd.setWriteCallback(audioCmd_write_cb);
    audioCmd.begin();

    // Data characteristic (write without response for speed)
    audioData.setProperties(CHR_PROPS_WRITE_WO_RESP);
    audioData.setPermission(SECMODE_OPEN, SECMODE_OPEN);
    audioData.setMaxLen(240);
    audioData.setWriteCallback(audioData_write_cb);
    audioData.begin();

    // Status characteristic (read + notify)
    audioStat.setProperties(CHR_PROPS_READ | CHR_PROPS_NOTIFY);
    audioStat.setPermission(SECMODE_OPEN, SECMODE_NO_ACCESS);
    audioStat.setMaxLen(4);
    audioStat.begin();

    // Start advertising
    Bluefruit.Advertising.addFlags(BLE_GAP_ADV_FLAGS_LE_ONLY_GENERAL_DISC_MODE);
    Bluefruit.Advertising.addTxPower();
    Bluefruit.Advertising.addService(audioSvc);
    Bluefruit.Advertising.addName();
    Bluefruit.Advertising.restartOnDisconnect(true);
    Bluefruit.Advertising.setInterval(160, 320); // 100-200ms
    Bluefruit.Advertising.start(0);  // advertise forever

    Serial.println("BLE advertising as 'BabyMobile'");
}


// ============================================================
// AUDIO PLAYBACK
// ============================================================

// nRF52840 has hardware PWM (up to 4 instances, 16MHz base clock)
// We use the nrf_pwm peripheral directly for audio output
// and a TIMER for the sample rate interrupt.

// For Arduino compatibility, we use analogWrite for the PWM
// and a software timer for sample feeding.

// micros() timestamp for next audio sample (loop-driven at 8kHz)
uint32_t g_nextSampleUs = 0;

// Software gain: 1=unity, 2=2x, etc. (clips at 0/255). Adjustable via 'u'/'d' serial.
static uint8_t g_audioGain = 3;

// Called from loop() every 125µs while g_playing
void audioTick(void) {
    // Output sample with software gain (stretches away from center 128)
    int16_t s = (int16_t)g_audioBuf[g_activeBuf][g_bufPos] - 128;
    s *= g_audioGain;
    if (s >  127) s =  127;
    if (s < -128) s = -128;
    analogWrite(PIN_AUDIO_PWM, (uint8_t)(s + 128));
    g_bufPos++;

    if (g_bufPos >= AUDIO_BUF_SIZE) {
        g_bufReady[g_activeBuf] = false;
        g_activeBuf ^= 1;
        g_bufPos = 0;

        if (!g_bufReady[g_activeBuf]) {
            // Buffer underrun — stop cleanly
            g_playing = false;
            analogWrite(PIN_AUDIO_PWM, SILENCE);
        }
    }
}

void audioInit() {
    pinMode(PIN_AUDIO_PWM, OUTPUT);
    analogWrite(PIN_AUDIO_PWM, SILENCE);
    analogWriteResolution(8);  // 8-bit PWM
}

void audioStart(uint8_t trackNum) {
    Serial.print("Playing track");
    Serial.print(trackNum);
    Serial.print(" of ");
    Serial.println(g_numTracks);
    if (trackNum >= g_numTracks) return;

    g_currentTrack = trackNum;
    g_playEnd = g_trackLen[trackNum];
    g_nextReadAddr = 0;

#ifdef POC_INTERNAL_FLASH
    if (g_pocFile) g_pocFile.close();
    char fname[16];
    pocFilename(trackNum, fname);
    if (!g_pocFile.open(fname, FILE_O_READ)) {
        Serial.print("audioStart: cannot open ");
        Serial.println(fname);
        return;
    }

    // Pre-fill buf 0
    uint32_t toRead = min((uint32_t)AUDIO_BUF_SIZE, g_playEnd - g_nextReadAddr);
    g_pocFile.read(g_audioBuf[0], toRead);
    g_nextReadAddr += toRead;
    g_bufReady[0] = true;

    // Pre-fill buf 1
    if (g_nextReadAddr < g_playEnd) {
        toRead = min((uint32_t)AUDIO_BUF_SIZE, g_playEnd - g_nextReadAddr);
        g_pocFile.read(g_audioBuf[1], toRead);
        g_nextReadAddr += toRead;
        g_bufReady[1] = true;
    } else {
        memset(g_audioBuf[1], SILENCE, AUDIO_BUF_SIZE);
        g_bufReady[1] = true;
    }
#else
    uint32_t start = g_trackStart[trackNum];
    g_playEnd = start + g_trackLen[trackNum];
    g_nextReadAddr = start;

    // Pre-fill both buffers
    uint32_t toRead = min((uint32_t)AUDIO_BUF_SIZE, g_playEnd - g_nextReadAddr);
    flashReadBytes(g_nextReadAddr, g_audioBuf[0], toRead);
    g_nextReadAddr += toRead;
    g_bufReady[0] = true;

    if (g_nextReadAddr < g_playEnd) {
        toRead = min((uint32_t)AUDIO_BUF_SIZE, g_playEnd - g_nextReadAddr);
        flashReadBytes(g_nextReadAddr, g_audioBuf[1], toRead);
        g_nextReadAddr += toRead;
        g_bufReady[1] = true;
    } else {
        memset(g_audioBuf[1], SILENCE, AUDIO_BUF_SIZE);
        g_bufReady[1] = true;
    }
#endif

    g_activeBuf = 0;
    g_bufPos = 0;

    // Enable amplifier
    digitalWrite(PIN_AMP_SD, HIGH);
    delay(10);

    g_playing = true;
    g_nextSampleUs = micros();

    Serial.print("Playing track ");
    Serial.println(trackNum + 1);
}

void audioStop() {
    g_playing = false;
    analogWrite(PIN_AUDIO_PWM, SILENCE);

#ifdef POC_INTERNAL_FLASH
    if (g_pocFile) g_pocFile.close();
#endif

    // Disable amp
    digitalWrite(PIN_AMP_SD, LOW);

    Serial.println("Playback stopped");
}


// ============================================================
// MOTOR
// ============================================================

void motorInit() {
    pinMode(PIN_MOTOR_PWM, OUTPUT);
    analogWrite(PIN_MOTOR_PWM, 0);
}

void motorStart(uint8_t speed) {
    g_motorSpeed = speed;
    g_motorOn = true;
    analogWrite(PIN_MOTOR_PWM, speed);
}

void motorStop() {
    g_motorOn = false;
    analogWrite(PIN_MOTOR_PWM, 0);
}


// ============================================================
// BUTTONS
// ============================================================

void buttonsInit() {
    for (uint8_t i = 0; i < NUM_BUTTONS; i++) {
        pinMode(BTN_PINS[i], INPUT_PULLUP);
    }
}

// Returns 1-8, or 0 if none pressed
uint8_t buttonRead() {
    for (uint8_t i = 0; i < NUM_BUTTONS; i++) {
        // Buttons pull to ground
        if (analogRead(BTN_PINS[i]) < 128) {
            delay(20);  // debounce
            if (analogRead(BTN_PINS[i]) < 128) {
                return i + 1;
            }
        }
    }
    return 0;
}

void waitButtonRelease(uint8_t btn) {
    unsigned long deadline = millis() + 300;
    while (buttonRead() == btn && millis() < deadline) {
        // Keep audio running while waiting for button release
        if (g_playing) {
            uint32_t now = micros();
            if ((int32_t)(now - g_nextSampleUs) >= 0) {
                g_nextSampleUs += 125;
                audioTick();
            }
        }
    }
}


// ============================================================
// SLEEP / WAKE
// ============================================================

void enterDeepSleep() {
    // Don't sleep if any button is currently LOW — would wake instantly.
    // Also catches floating pins that the internal pull-up isn't winning against.
    for (uint8_t i = 0; i < NUM_BUTTONS; i++) {
        if (analogRead(BTN_PINS[i]) < 128) {
            g_lastActivity = millis();  // postpone
            return;
        }
    }

    Serial.println("Entering deep sleep...");
    audioStop();
    motorStop();
#ifndef POC_INTERNAL_FLASH
    flashSleep();
#endif
    digitalWrite(PIN_LED, LOW);

    // Stop BLE advertising to save power
    Bluefruit.Advertising.stop();

    // Configure buttons as wake sources.
    // nrf_gpio_cfg_sense_input takes nRF GPIO numbers, not Arduino pin numbers —
    // use g_ADigitalPinMap[] to convert.
    for (uint8_t i = 0; i < NUM_BUTTONS; i++) {
        nrf_gpio_cfg_sense_input(g_ADigitalPinMap[BTN_PINS[i]],
                                 NRF_GPIO_PIN_PULLUP,
                                 NRF_GPIO_PIN_SENSE_LOW);
    }

    // System OFF — lowest power, wakes via GPIO sense or reset
    sd_power_system_off();

    // Execution stops here. Device resets on wake.
    // (This line is never reached)
}


// ============================================================
// SETUP
// ============================================================

void setup() {
    // LED early — visual feedback
    pinMode(PIN_LED, OUTPUT);
    digitalWrite(PIN_LED, HIGH);

    Serial.begin(115200);
    Serial.println("...");
    
    int c=0;
    while(c<10) {
      delay(250);
      digitalWrite(PIN_LED, LOW);
      // Serial.println("0");
      delay(250);
      digitalWrite(PIN_LED, HIGH);
      // Serial.println("1");
      c++;
    }

    Serial.println("=== Baby Mobile v2 ===");
    Serial.println("nRF52840 + IS25LP128F + PAM8302A");

    // Pin setup
    pinMode(PIN_AMP_SD, OUTPUT);
    digitalWrite(PIN_AMP_SD, LOW);

#ifdef POC_INTERNAL_FLASH
    InternalFS.begin();
    Serial.println("InternalFS mounted");
#else
    pinMode(PIN_FLASH_CS, OUTPUT);
    digitalWrite(PIN_FLASH_CS, HIGH);

    // SPI
    SPI.begin();

    // Flash
    flashWake();
    delay(1);
    uint32_t jedec = flashReadJEDEC();
    Serial.print("Flash JEDEC: 0x");
    Serial.println(jedec, HEX);

    if (jedec == 0x9D6018 || jedec == 0x9D6017 ||  // IS25LP128F / IS25LP064
        jedec == 0xEF4018 || jedec == 0xEF4017) {  // W25Q128 / W25Q064
        Serial.println("Flash detected OK");
    } else if (jedec == 0x000000 || jedec == 0xFFFFFF) {
        Serial.println("WARNING: No flash detected! Check SPI wiring.");
    }
#endif

    // Load tracks
    loadTrackTable();

    // Buttons
    buttonsInit();

    // Motor
    motorInit();

    // Audio
    audioInit();

#ifdef POC_INTERNAL_FLASH
    // Boot test tone: 440Hz square wave for 1 second
    // If you hear a beep, the amp/PWM chain is working.
    Serial.println("Boot tone...");
    digitalWrite(PIN_AMP_SD, HIGH);
    delay(20);
    for (int i = 0; i < 8000; i++) {
        // 440Hz at 8kHz sample rate: 8000/440 ≈ 18 samples/cycle
        analogWrite(PIN_AUDIO_PWM, (i % 18) < 9 ? 255 : 0);
        delayMicroseconds(125);
    }
    analogWrite(PIN_AUDIO_PWM, SILENCE);
    delay(20);
    digitalWrite(PIN_AMP_SD, LOW);
    Serial.println("Boot tone done");

    // Hex dump track 0 + blocking playback diagnostic
    // (after audioInit so amp/PWM are ready)
    if (g_numTracks > 0) {
        char diagName[16]; pocFilename(0, diagName);
        File diagF(InternalFS);
        if (diagF.open(diagName, FILE_O_READ)) {
            uint32_t fsz = diagF.size();
            Serial.print("Track0 size: "); Serial.println(fsz);
            uint8_t hbuf[64];
            int nr = diagF.read(hbuf, sizeof(hbuf));
            diagF.close();
            Serial.print("First "); Serial.print(nr); Serial.println(" bytes (hex):");
            for (int i = 0; i < nr; i++) {
                if (hbuf[i] < 0x10) Serial.print("0");
                Serial.print(hbuf[i], HEX);
                Serial.print(i % 16 == 15 ? "\n" : " ");
            }
            Serial.println();
        } else {
            Serial.println("Cannot open track0 for read!");
        }

        Serial.println("Diag: blocking play track0 for 3s...");
        audioStart(0);
        uint32_t diagEnd = millis() + 3000;
        while (millis() < diagEnd && g_playing) {
            uint32_t now = micros();
            if ((int32_t)(now - g_nextSampleUs) >= 0) {
                g_nextSampleUs += 125;
                audioTick();
            }
            for (uint8_t b = 0; b < 2; b++) {
                if (!g_bufReady[b] && g_nextReadAddr < g_playEnd) {
                    uint32_t toRead = min((uint32_t)AUDIO_BUF_SIZE, g_playEnd - g_nextReadAddr);
                    g_pocFile.read(g_audioBuf[b], toRead);
                    for (uint32_t i = toRead; i < AUDIO_BUF_SIZE; i++) g_audioBuf[b][i] = SILENCE;
                    g_nextReadAddr += toRead;
                    g_bufReady[b] = true;
                }
            }
        }
        audioStop();
        Serial.println("Diag done");
    }
#endif

    // USB Mass Storage
#ifndef POC_INTERNAL_FLASH
    usb_msc.setID("BabyMobile", "Audio Drive", "2.0");
    usb_msc.setReadWriteCallback(msc_read_cb, msc_write_cb, msc_flush_cb);
    //usb_msc.setStartStopCallback(msc_start_stop_cb); // unused in nrf?
    usb_msc.setCapacity(MSC_BLOCK_COUNT, MSC_BLOCK_SIZE);
    //usb_msc.setReadOnly(false);
    usb_msc.setUnitReady(true);
    usb_msc.begin();
#endif

    // BLE
    setupBLE();

    g_lastActivity = millis();

    digitalWrite(PIN_LED, HIGH);
    Serial.println("Ready! Plug in USB to upload audio, or press a button.");
}


// ============================================================
// MAIN LOOP
// ============================================================

// ============================================================
// SERIAL UPLOAD STATE MACHINE (POC_INTERNAL_FLASH only)
// ============================================================
#ifdef POC_INTERNAL_FLASH
enum SerUploadState { SER_IDLE, SER_RECEIVING };
static SerUploadState  g_serState       = SER_IDLE;
static char            g_serLineBuf[64] = {0};
static uint8_t         g_serLineLen     = 0;
static uint8_t         g_serTrack       = 0;
static uint32_t        g_serBytesExpected = 0;
static uint32_t        g_serBytesReceived = 0;
static bool            g_serSkipHeader  = false;
static File            g_serFile(InternalFS);

static void serUploadTick() {
    if (g_serState == SER_IDLE) {
        // Accumulate characters until newline
        while (Serial.available()) {
            // Serial.print("r");
            char c = (char)Serial.read();
            if (c == '\n' || c == '\r') {
                g_serLineBuf[g_serLineLen] = '\0';
                if (g_serLineLen == 0) { g_serLineLen = 0; break; }
                // Parse: UPLOAD <track> <len>
                unsigned int utrk = 0, ulen = 0;
                if (sscanf(g_serLineBuf, "UPLOAD %u %u", &utrk, &ulen) == 2) {
                    g_serTrack         = (uint8_t)utrk;
                    g_serBytesExpected = (uint32_t)ulen;
                    g_serBytesReceived = 0;
                    g_serSkipHeader    = true;

                    char fname[16];
                    pocFilename(g_serTrack, fname);
                    if (g_serFile) g_serFile.close();
                    // Remove first — FILE_O_WRITE has no truncate flag
                    InternalFS.remove(fname);
                    if (!g_serFile.open(fname, FILE_O_WRITE)) {
                        Serial.print("ERR cannot open ");
                        Serial.println(fname);
                    } else {
                        g_serState = SER_RECEIVING;
                        g_usbConnected = true;
                        Serial.println("READY");
                    }
                } else if (sscanf(g_serLineBuf, "DUMP %u", &utrk) == 1) {
                    // Hex-dump first 256 bytes of a track file
                    char fname[16];
                    pocFilename((uint8_t)utrk, fname);
                    File df(InternalFS);
                    if (df.open(fname, FILE_O_READ)) {
                        uint32_t fsz = df.size();
                        Serial.print("SIZE "); Serial.println(fsz);
                        uint8_t dbuf[16];
                        uint32_t limit = min(fsz, (uint32_t)256);
                        uint32_t off = 0;
                        while (off < limit) {
                            int n = df.read(dbuf, min((uint32_t)sizeof(dbuf), limit - off));
                            if (n <= 0) break;
                            for (int j = 0; j < n; j++) {
                                if (dbuf[j] < 0x10) Serial.print("0");
                                Serial.print(dbuf[j], HEX);
                                Serial.print(j % 16 == 15 || (off + j + 1) == limit ? "\n" : " ");
                            }
                            off += n;
                        }
                        df.close();
                        Serial.println("END");
                    } else {
                        Serial.println("NO FILE");
                    }
                } else if (g_serLineLen == 1 && (g_serLineBuf[0] == 'u' || g_serLineBuf[0] == 'd')) {
                    if (g_serLineBuf[0] == 'u') g_audioGain++;
                    else if (g_audioGain > 1) g_audioGain--;
                    Serial.print("GAIN "); Serial.println(g_audioGain);
                } else {
                    Serial.print("ERR bad cmd: ");
                    Serial.println(g_serLineBuf);
                }
                g_serLineLen = 0;
                // Serial.print("0");
            } else {
                if (g_serLineLen < (sizeof(g_serLineBuf) - 1)) {
                    g_serLineBuf[g_serLineLen++] = c;
                }
                // Serial.print(",");
            }
        }
    } else { // SER_RECEIVING
        // Serial.print("e");
        uint8_t chunk[64];
        while (Serial.available() && g_serBytesReceived < g_serBytesExpected) {
            // Serial.print(";");
            int n = Serial.readBytes(chunk, min((int)sizeof(chunk),
                                    (int)(g_serBytesExpected - g_serBytesReceived)));
            if (n <= 0) break;

            uint8_t *src    = chunk;
            uint16_t srcLen = (uint16_t)n;

            // Strip 44-byte WAV header from first chunk if present
            if (g_serSkipHeader) {
                g_serSkipHeader = false;
                if (srcLen >= 4 && src[0]=='R' && src[1]=='I' && src[2]=='F' && src[3]=='F') {
                    if (srcLen > 44) { src += 44; srcLen -= 44; }
                    else { g_serBytesReceived += n; continue; }
                }
            }

            int32_t wr = g_serFile.write(src, srcLen);
            if (wr != (int32_t)srcLen) {
                // Write failed (LittleFS full or error) — abort immediately.
                // Do NOT loop printing errors; that fills USB CDC TX and hangs.
                g_serFile.close();
                g_usbConnected = false;
                g_serState = SER_IDLE;
                Serial.print("ERR WRITE_FAIL at=");
                Serial.println(g_serBytesReceived);
                break;
            }
            g_serBytesReceived += n;
        }

        if (g_serState != SER_RECEIVING) return;  // aborted in write-fail handler above

        if (g_serBytesReceived >= g_serBytesExpected) {
            g_serFile.close();
            // Reopen to read the actual flushed size from LittleFS
            char fname2[16]; pocFilename(g_serTrack, fname2);
            File tmp(InternalFS);
            uint32_t fsz = 0;
            if (tmp.open(fname2, FILE_O_READ)) { fsz = tmp.size(); tmp.close(); }
            g_usbConnected = false;
            g_serState = SER_IDLE;
            loadTrackTable();
            Serial.print("OK ");
            Serial.print(fsz);
            Serial.print("/");
            Serial.println(g_serBytesReceived);
        }
    }
}
#endif

void loop() {
    // ---- Audio sample output (8kHz, loop-driven) ----
    if (g_playing) {
        g_lastActivity = millis();  // stay awake while audio is playing
        uint32_t now = micros();
        if ((int32_t)(now - g_nextSampleUs) >= 0) {
            g_nextSampleUs += 125;  // 1/8000s = 125µs
            audioTick();
        }
    }

    // ---- Serial upload (POC mode) ----
#ifdef POC_INTERNAL_FLASH
    serUploadTick();
#endif

    // ---- Refill audio buffers ----
    if (g_playing) {
        for (uint8_t b = 0; b < 2; b++) {
            if (!g_bufReady[b] && g_nextReadAddr < g_playEnd) {
                uint32_t remaining = g_playEnd - g_nextReadAddr;
                uint32_t toRead = min((uint32_t)AUDIO_BUF_SIZE, remaining);

#ifdef POC_INTERNAL_FLASH
                g_pocFile.read(g_audioBuf[b], toRead);
#else
                flashReadBytes(g_nextReadAddr, g_audioBuf[b], toRead);
#endif
                // Pad with silence
                for (uint32_t i = toRead; i < AUDIO_BUF_SIZE; i++) {
                    g_audioBuf[b][i] = SILENCE;
                }
                g_nextReadAddr += toRead;
                g_bufReady[b] = true;
            }
        }

        // Track finished? Auto-advance and loop
        if (!g_bufReady[0] && !g_bufReady[1] && g_nextReadAddr >= g_playEnd) {
            audioStop();
            if (g_numTracks > 0) {
                g_currentTrack = (g_currentTrack + 1) % g_numTracks;
                audioStart(g_currentTrack);
            }
        }
    }

    // ---- Handle buttons ----
    // Per-button cooldown prevents phantom/stuck pins from re-firing
    static unsigned long btnLastMs[NUM_BUTTONS + 1] = {0};
    uint8_t btn = buttonRead();
    if (btn > 0 && millis() - btnLastMs[btn] > 800) {
        btnLastMs[btn] = millis();
        Serial.print("BTN"); Serial.println(btn);
        g_lastActivity = millis();

        switch (btn) {
            case 1:  // Play / Pause
            case 2:
            case 3:
                if (g_playing) {
                    audioStop();
                } else if (g_numTracks > 0) {
                    audioStart(g_currentTrack);
                }
                break;

            case 4:  // Next track
                if (g_numTracks > 0) {
                    g_currentTrack = (g_currentTrack + 1) % g_numTracks;
                    if (g_playing) { audioStop(); audioStart(g_currentTrack); }
                }
                break;

            case 5:  // Previous track
                if (g_numTracks > 0) {
                    g_currentTrack = (g_currentTrack == 0) ? g_numTracks - 1 : g_currentTrack - 1;
                    if (g_playing) { audioStop(); audioStart(g_currentTrack); }
                }
                break;

            case 6:  // Motor toggle
                if (g_motorOn) motorStop(); else motorStart(g_motorSpeed);
                break;

            case 7:  // Motor faster
                if (g_motorSpeed < 240) g_motorSpeed += 20;
                if (g_motorOn) analogWrite(PIN_MOTOR_PWM, g_motorSpeed);
                break;

            case 8:  // Motor slower
                if (g_motorSpeed > 60) g_motorSpeed -= 20;
                if (g_motorOn) analogWrite(PIN_MOTOR_PWM, g_motorSpeed);
                break;

            case 9:  // Play all (music + motor)
                motorStart(g_motorSpeed);
                if (!g_playing && g_numTracks > 0) audioStart(0);
                break;

            case 10:  // Stop all
                audioStop();
                motorStop();
                break;
        }

        waitButtonRelease(btn);
    }

    // ---- Auto shutoff / idle sleep ----
    if (millis() - g_lastActivity > AUTO_OFF_MS) {
        enterDeepSleep();
    }

    if (!g_playing && !g_motorOn && !g_usbConnected && !g_bleUploading) {
        if (millis() - g_lastActivity > IDLE_SLEEP_MS) {
            enterDeepSleep();
        }
    }

}
