#!/usr/bin/env python3
"""
flash_upload.py — Upload audio to Baby Mobile v2 via serial.

This is a fallback for uploading audio if USB mass storage isn't working.
Normally you'd just drag files onto the USB drive.

Usage:
  python3 flash_upload.py /dev/ttyACM0 song1.mp3 song2.wav song3.raw

Audio files will be auto-converted to 8kHz/8-bit/unsigned PCM using ffmpeg.

Track table format (v2):
  [0]     magic = 0xBB
  [1]     num_tracks
  [2..3]  reserved
  [4..]   12-byte entries: start(4) + length(4) + rate_khz(1) + bits(1) + pad(2)
"""

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

BAUD = 115200
PAGE_SIZE = 256
SECTOR_SIZE = 4096
AUDIO_START = 0x1000
MAX_TRACKS = 32
FLASH_SIZE = 16 * 1024 * 1024
TABLE_MAGIC = 0xBB
ENTRY_SIZE = 12


def open_serial(port):
    ser = serial.Serial(port, BAUD, timeout=2)
    time.sleep(2)
    ser.read(ser.in_waiting)
    return ser


def convert_to_raw(input_path):
    """Convert any audio file to 8kHz 8-bit unsigned PCM."""
    if input_path.endswith('.raw'):
        return input_path

    raw_path = input_path + ".raw"
    if not shutil.which('ffmpeg'):
        print(f"  ERROR: ffmpeg not found. 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)}...")
    result = subprocess.run([
        'ffmpeg', '-y', '-i', input_path,
        '-ar', '8000', '-ac', '1', '-f', 'u8', '-acodec', 'pcm_u8',
        raw_path
    ], capture_output=True)

    if result.returncode != 0:
        print(f"  ERROR: ffmpeg failed: {result.stderr.decode()[-200:]}")
        sys.exit(1)

    return raw_path


def build_track_table(tracks):
    """Build v2 track table."""
    table = bytearray()
    table.append(TABLE_MAGIC)
    table.append(len(tracks))
    table.extend(b'\x00\x00')  # reserved

    for start, length in tracks:
        table.extend(struct.pack('>I', start))
        table.extend(struct.pack('>I', length))
        table.append(8)   # sample rate kHz
        table.append(8)   # bits per sample
        table.extend(b'\x00\x00')

    # Pad to page boundary
    while len(table) % PAGE_SIZE != 0:
        table.append(0xFF)

    return bytes(table)


def serial_cmd(ser, cmd):
    """Send a single-char command and read response line."""
    ser.write(cmd.encode())
    return ser.readline().decode().strip()


def serial_write_page(ser, addr, data):
    """Write up to 256 bytes via the serial 'W' command."""
    assert len(data) <= 256
    ser.write(b'W')
    ser.write(struct.pack('>I', addr)[1:])   # 3-byte addr
    ser.write(struct.pack('>H', len(data)))   # 2-byte length
    ser.write(data)
    return ser.readline().decode().strip()


def serial_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]
        if len(chunk) < PAGE_SIZE:
            chunk = chunk + b'\x80' * (PAGE_SIZE - len(chunk))

        serial_write_page(ser, addr, chunk)
        written += PAGE_SIZE
        addr += PAGE_SIZE

        pct = min(100, written * 100 // total)
        print(f"\r  Writing {label}: {pct}%", end="", flush=True)

    print()


def main():
    if len(sys.argv) < 3:
        print("Usage: python3 flash_upload.py <port> <audio_file> [...]")
        print("  Supported: .raw .mp3 .wav .ogg .flac .aac")
        sys.exit(1)

    port = sys.argv[1]
    audio_files = sys.argv[2:]

    # Convert
    print("=== Preparing audio ===")
    raw_files = []
    for f in audio_files:
        raw = convert_to_raw(f)
        size = os.path.getsize(raw)
        print(f"  {os.path.basename(f)}: {size/1024:.0f} KB ({size/8000:.1f}s)")
        raw_files.append(raw)

    total = sum(os.path.getsize(f) for f in raw_files)
    print(f"  Total: {total/1024:.0f} KB ({total/8000/60:.1f} min)")

    # Connect
    print(f"\n=== Connecting to {port} ===")
    ser = open_serial(port)
    print(f"  {serial_cmd(ser, 'I')}")

    # Erase
    print("\n=== Erasing flash ===")
    ser.write(b'E')
    while True:
        line = ser.readline().decode().strip()
        if line: print(f"  {line}")
        if "Done" in line: break

    # Write tracks
    print("\n=== Writing audio ===")
    tracks = []
    addr = AUDIO_START

    for i, raw in enumerate(raw_files):
        with open(raw, 'rb') as f:
            data = f.read()

        tracks.append((addr, len(data)))
        serial_write_data(ser, addr, data, label=f"Track {i+1}")

        addr += len(data)
        # Align to sector
        if addr % SECTOR_SIZE:
            addr = ((addr // SECTOR_SIZE) + 1) * SECTOR_SIZE

    # Write track table
    print("\n=== Writing track table ===")
    table = build_track_table(tracks)
    serial_write_data(ser, 0, table, label="Track table")

    # Summary
    print(f"\n=== Done! {len(tracks)} tracks uploaded ===")
    for i, (s, l) in enumerate(tracks):
        print(f"  Track {i+1}: 0x{s:06X}, {l/8000:.1f}s")

    # Cleanup
    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()