baby-mobile/flash_upload.py

#!/usr/bin/env python3
"""
flash_upload.py — Upload audio files to the Baby Mobile's W25Q128 flash.

Usage:
  python3 flash_upload.py /dev/ttyUSB0 song1.raw song2.raw song3.raw ...

Audio files must be raw unsigned 8-bit PCM at 8000 Hz mono.
Convert with ffmpeg:
  ffmpeg -i song.mp3 -ar 8000 -ac 1 -f u8 -acodec pcm_u8 song.raw

The script writes a track table to sector 0, then audio data starting
at sector 1 (address 0x1000).
"""

import serial
import struct
import sys
import time
import os
import subprocess
import shutil

BAUD = 9600
PAGE_SIZE = 256
SECTOR_SIZE = 4096
AUDIO_START = 0x1000
MAX_TRACKS = 32
FLASH_SIZE = 16 * 1024 * 1024  # 16MB


def open_serial(port):
    ser = serial.Serial(port, BAUD, timeout=2)
    time.sleep(2)  # wait for Arduino reset
    # Flush
    ser.read(ser.in_waiting)
    return ser


def flash_identify(ser):
    ser.write(b'I')
    line = ser.readline().decode().strip()
    print(f"  {line}")
    return line


def flash_erase(ser):
    print("  Erasing entire flash (takes ~40 seconds)...")
    ser.write(b'E')
    line = ser.readline().decode().strip()
    while "Done" not in line:
        line = ser.readline().decode().strip()
        if line:
            print(f"  {line}")
    print("  Erase complete.")


def flash_write_page(ser, addr, data):
    """Write up to 256 bytes to flash."""
    assert len(data) <= 256
    cmd = b'W'
    cmd += struct.pack('>I', addr)[1:]  # 3 bytes address
    cmd += struct.pack('>H', len(data))
    ser.write(cmd)
    ser.write(data)
    resp = ser.readline().decode().strip()
    return resp


def flash_write_data(ser, start_addr, data, label=""):
    """Write arbitrary length data, page by page."""
    total = len(data)
    written = 0
    addr = start_addr
    
    while written < total:
        chunk = data[written:written + PAGE_SIZE]
        # Pad last page
        if len(chunk) < PAGE_SIZE:
            chunk = chunk + b'\x80' * (PAGE_SIZE - len(chunk))  # 0x80 = silence
        
        flash_write_page(ser, addr, chunk)
        written += PAGE_SIZE
        addr += PAGE_SIZE
        
        pct = min(100, written * 100 // total)
        print(f"\r  Writing {label}: {pct}% ({written}/{total} bytes)", end="", flush=True)
    
    print()
    return addr


def convert_to_raw(input_path):
    """Convert any audio file to 8kHz 8-bit unsigned PCM using ffmpeg."""
    raw_path = input_path + ".raw"
    
    if input_path.endswith('.raw'):
        return input_path
    
    if not shutil.which('ffmpeg'):
        print(f"ERROR: ffmpeg not found. Please convert manually:")
        print(f"  ffmpeg -i {input_path} -ar 8000 -ac 1 -f u8 -acodec pcm_u8 {raw_path}")
        sys.exit(1)
    
    print(f"  Converting {os.path.basename(input_path)} to 8kHz/8bit PCM...")
    subprocess.run([
        'ffmpeg', '-y', '-i', input_path,
        '-ar', '8000', '-ac', '1', '-f', 'u8', '-acodec', 'pcm_u8',
        raw_path
    ], capture_output=True)
    
    return raw_path


def build_track_table(tracks):
    """Build the track index (stored in sector 0 of flash).
    
    Format:
      byte 0:    num_tracks (uint8)
      bytes 1-3: padding
      bytes 4+:  track_start[MAX_TRACKS] as uint32 big-endian
      then:      track_length[MAX_TRACKS] as uint32 big-endian
    """
    num = len(tracks)
    table = struct.pack('B', num) + b'\x00' * 3
    
    # Start addresses
    for start, length in tracks:
        table += struct.pack('>I', start)
    for _ in range(MAX_TRACKS - num):
        table += struct.pack('>I', 0)
    
    # Lengths
    for start, length in tracks:
        table += struct.pack('>I', length)
    for _ in range(MAX_TRACKS - num):
        table += struct.pack('>I', 0)
    
    # Pad to sector size
    table += b'\xFF' * (SECTOR_SIZE - len(table))
    return table


def main():
    if len(sys.argv) < 3:
        print("Usage: python3 flash_upload.py <serial_port> <audio_file> [audio_file ...]")
        print()
        print("Supported input formats: .raw (8kHz/8bit/unsigned), .mp3, .wav, .ogg, .flac")
        print("Non-.raw files will be converted automatically using ffmpeg.")
        print()
        print("Example:")
        print("  python3 flash_upload.py /dev/ttyUSB0 lullaby.mp3 twinkle.wav brahms.mp3")
        sys.exit(1)
    
    port = sys.argv[1]
    audio_files = sys.argv[2:]
    
    if len(audio_files) > MAX_TRACKS:
        print(f"ERROR: Maximum {MAX_TRACKS} tracks supported.")
        sys.exit(1)
    
    # Convert files to raw PCM
    print("=== Preparing audio files ===")
    raw_files = []
    for f in audio_files:
        raw = convert_to_raw(f)
        size = os.path.getsize(raw)
        duration = size / 8000
        print(f"  {os.path.basename(f)}: {size} bytes ({duration:.1f} seconds)")
        raw_files.append(raw)
    
    # Check total size
    total_audio = sum(os.path.getsize(f) for f in raw_files)
    available = FLASH_SIZE - AUDIO_START
    if total_audio > available:
        print(f"ERROR: Audio ({total_audio} bytes) exceeds flash capacity ({available} bytes)")
        sys.exit(1)
    
    total_duration = total_audio / 8000
    print(f"\n  Total: {total_audio} bytes ({total_duration:.1f} seconds / {total_duration/60:.1f} minutes)")
    print(f"  Flash usage: {(total_audio + AUDIO_START) * 100 // FLASH_SIZE}%")
    
    # Connect
    print(f"\n=== Connecting to {port} ===")
    ser = open_serial(port)
    flash_identify(ser)
    
    # Erase
    print("\n=== Erasing flash ===")
    flash_erase(ser)
    
    # Write audio data and build track list
    print("\n=== Writing audio ===")
    tracks = []
    addr = AUDIO_START
    
    for i, raw_path in enumerate(raw_files):
        with open(raw_path, 'rb') as f:
            data = f.read()
        
        start = addr
        length = len(data)
        tracks.append((start, length))
        
        name = os.path.basename(audio_files[i])
        addr = flash_write_data(ser, start, data, label=f"Track {i+1} ({name})")
        
        # Align to sector boundary for clean layout
        if addr % SECTOR_SIZE != 0:
            addr = ((addr // SECTOR_SIZE) + 1) * SECTOR_SIZE
    
    # Write track table to sector 0
    print("\n=== Writing track table ===")
    table = build_track_table(tracks)
    flash_write_data(ser, 0, table, label="Track table")
    
    # Summary
    print("\n=== Upload complete! ===")
    print(f"  {len(tracks)} tracks written")
    for i, (start, length) in enumerate(tracks):
        print(f"  Track {i+1}: addr=0x{start:06X}, {length} bytes ({length/8000:.1f}s)")
    print(f"\n  Reset the board to start playback.")
    
    # Cleanup temp files
    for raw, orig in zip(raw_files, audio_files):
        if raw != orig and os.path.exists(raw):
            os.remove(raw)
    
    ser.close()


if __name__ == "__main__":
    main()