#!/usr/bin/env python3
"""
upload_track.py — serial audio uploader for baby_mobile_v2 POC_INTERNAL_FLASH mode

Usage:
    python upload_track.py <port> <track_num> <file>

Supported formats: .wav, .mp3, .ogg, .flac, .aac, .m4a, .raw (anything ffmpeg handles)

Output format: 16-bit signed PCM WAV, 16 kHz, mono (WAV header included).
  - Already-conformant WAV files (PCM, mono, 8 or 16-bit, 8/16/32 kHz) are sent as-is.
  - All other formats are converted via ffmpeg to 16-bit 16 kHz mono WAV.
  - .raw files are treated as legacy 8-bit unsigned 16 kHz and wrapped in a WAV header.

Requires: pyserial  (pip install pyserial)
          ffmpeg in PATH for non-WAV/RAW files or non-conformant WAV files
"""

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

CHUNK_SIZE = 512
BAUD_RATE  = 115200
TIMEOUT_S  = 10.0

# Sample rates the firmware PWM can reproduce (32kHz carrier / (REFRESH+1))
SUPPORTED_RATES = {8000, 16000, 32000}


def _make_wav_header(num_samples: int, sample_rate: int, bits: int) -> bytes:
    """Build a minimal 44-byte PCM WAV header."""
    num_channels  = 1
    byte_rate     = sample_rate * num_channels * bits // 8
    block_align   = num_channels * bits // 8
    data_size     = num_samples * block_align
    chunk_size    = 36 + data_size
    return struct.pack('<4sI4s4sIHHIIHH4sI',
        b'RIFF', chunk_size, b'WAVE',
        b'fmt ', 16, 1, num_channels, sample_rate,
        byte_rate, block_align, bits,
        b'data', data_size)


def load_wav(path: str) -> bytes:
    """Return full WAV bytes (header + PCM data) ready to upload."""
    ext = os.path.splitext(path)[1].lower()

    # ── Legacy .raw: 8-bit unsigned 16 kHz mono ─────────────────────────────
    if ext == '.raw':
        with open(path, 'rb') as f:
            pcm = f.read()
        print(f"RAW file: wrapping {len(pcm):,} bytes as 8-bit 16 kHz mono WAV")
        header = _make_wav_header(len(pcm), 16000, 8)
        return header + pcm

    # ── WAV: check if already conformant ────────────────────────────────────
    if ext == '.wav':
        with open(path, 'rb') as f:
            data = f.read()
        if data[:4] == b'RIFF' and len(data) >= 44:
            try:
                audio_format  = struct.unpack_from('<H', data, 20)[0]
                num_channels  = struct.unpack_from('<H', data, 22)[0]
                sample_rate   = struct.unpack_from('<I', data, 24)[0]
                bits_per_samp = struct.unpack_from('<H', data, 34)[0]
                conformant = (
                    audio_format == 1 and          # PCM
                    num_channels == 1 and           # mono
                    bits_per_samp in (8, 16) and
                    sample_rate in SUPPORTED_RATES
                )
            except Exception:
                conformant = False

            if conformant:
                duration = (len(data) - 44) / (bits_per_samp // 8) / sample_rate
                print(f"WAV already conformant: {bits_per_samp}-bit {sample_rate // 1000}kHz "
                      f"mono — {duration:.1f}s")
                return data  # send as-is, header included
            print("WAV needs conversion — running ffmpeg")
        # fall through to ffmpeg

    # ── All other formats (and non-conformant WAV): convert via ffmpeg ───────
    if not shutil.which('ffmpeg'):
        sys.exit(
            "ERROR: ffmpeg not found. Install it or convert manually:\n"
            f"  ffmpeg -i \"{path}\" -ar 16000 -ac 1 -acodec pcm_s16le out.wav"
        )

    print(f"Converting {os.path.basename(path)} via ffmpeg -> 16-bit 16 kHz mono WAV...")
    result = subprocess.run(
        ['ffmpeg', '-y', '-i', path,
         '-ar', '16000', '-ac', '1', '-f', 'wav', '-acodec', 'pcm_s16le', 'pipe:1'],
        capture_output=True
    )
    if result.returncode != 0:
        sys.exit(f"ERROR: ffmpeg failed:\n{result.stderr.decode(errors='replace')[-400:]}")

    wav = result.stdout
    if len(wav) >= 44:
        bits   = struct.unpack_from('<H', wav, 34)[0]
        rate   = struct.unpack_from('<I', wav, 24)[0]
        audio_bytes = len(wav) - 44
        duration = audio_bytes / (bits // 8) / rate
        print(f"  Converted: {audio_bytes:,} audio bytes ({duration:.1f}s), "
              f"{bits}-bit {rate // 1000}kHz")
    return wav


def upload(port: str, track_num: int, wav: bytes) -> None:
    total = len(wav)
    print(f"Uploading {total} bytes as track {track_num} via {port} ...")

    with serial.Serial(port, BAUD_RATE, timeout=2.0,
                       dsrdtr=False, rtscts=False) as ser:
        time.sleep(0.5)
        ser.reset_input_buffer()
        ser.write(b"\n")
        ser.flush()
        time.sleep(0.1)
        ser.reset_input_buffer()

        # Send UPLOAD command
        cmd = f"UPLOAD {track_num} {total}\n"
        ser.write(cmd.encode())
        ser.flush()

        # Wait for READY
        deadline = time.time() + TIMEOUT_S
        while True:
            if time.time() > deadline:
                sys.exit("Timed out waiting for READY")
            line = ser.readline().decode(errors="replace").strip()
            if line == "READY":
                break
            if line:
                print(f"  firmware: {line}")

        # Stream WAV data in chunks; poll for ERR
        sent = 0
        err_line = None
        ser.timeout = 0.05
        while sent < total:
            chunk = wav[sent:sent + CHUNK_SIZE]
            ser.write(chunk)
            sent += len(chunk)
            pct = sent * 100 // total
            print(f"\r  {sent}/{total} bytes ({pct}%)   ", end="", flush=True)
            line = ser.readline().decode(errors="replace").strip()
            if line.startswith("ERR"):
                err_line = line
                break
        print()
        ser.timeout = 2.0

        if err_line:
            print(f"ERROR from firmware: {err_line}")
            if "at=" in err_line:
                accepted = int(err_line.split("at=")[1].split()[0])
                safe = int(accepted * 0.9)
                # estimate audio bytes (subtract header)
                audio_accepted = max(0, safe - 44)
                bits = struct.unpack_from('<H', wav, 34)[0] if len(wav) >= 44 else 16
                rate = struct.unpack_from('<I', wav, 24)[0] if len(wav) >= 44 else 16000
                max_s = audio_accepted / (bits // 8) / rate
                print(f"  LittleFS accepted {accepted} bytes before full.")
                print(f"  Safe clip length: ~{max_s:.1f}s")
                print(f"  Trim with:  ffmpeg -i input.mp3 -t {max_s:.0f} -ar 16000 -ac 1 "
                      f"-acodec pcm_s16le out.wav")
            sys.exit(1)

        # Wait for OK
        deadline = time.time() + TIMEOUT_S
        while True:
            if time.time() > deadline:
                sys.exit("Timed out waiting for OK")
            line = ser.readline().decode(errors="replace").strip()
            if line.startswith("OK"):
                print(f"Upload done: {line}")
                break
            if line.startswith("ERR"):
                print(f"ERROR from firmware: {line}")
                sys.exit(1)
            if line:
                print(f"  firmware: {line}")

        # Verify: hex-dump first 256 bytes and check WAV header is intact
        print("Verifying...")
        ser.reset_input_buffer()
        ser.write(f"DUMP {track_num}\n".encode())
        ser.flush()

        deadline = time.time() + TIMEOUT_S
        fw_size = None
        fw_bytes = bytearray()
        while True:
            if time.time() > deadline:
                print("WARNING: timed out waiting for DUMP response")
                break
            line = ser.readline().decode(errors="replace").strip()
            if not line:
                continue
            if line.startswith("SIZE"):
                fw_size = int(line.split()[1])
            elif line == "END":
                break
            elif line == "NO FILE":
                print("ERROR: firmware says file does not exist!")
                break
            else:
                for tok in line.split():
                    try:
                        fw_bytes.append(int(tok, 16))
                    except ValueError:
                        pass

        if fw_size is not None:
            print(f"  Firmware file size: {fw_size} bytes")
            if fw_size == 0:
                print("  ERROR: file is empty — write failed (LittleFS full?)")
            elif fw_size < total:
                print(f"  WARNING: only {fw_size}/{total} bytes written (LittleFS full?)")
            else:
                print(f"  Size OK: {fw_size}/{total}")

        # Compare first 256 bytes (includes WAV header — confirms format metadata made it)
        compare_len = min(len(fw_bytes), len(wav), 256)
        if compare_len > 0:
            mismatches = sum(1 for i in range(compare_len) if fw_bytes[i] != wav[i])
            if mismatches == 0:
                print(f"  Data OK: first {compare_len} bytes match (WAV header intact)")
            else:
                print(f"  ERROR: {mismatches}/{compare_len} byte mismatches in first {compare_len} bytes")
                print(f"  Expected: {' '.join(f'{b:02X}' for b in wav[:16])}")
                print(f"  Got:      {' '.join(f'{b:02X}' for b in fw_bytes[:16])}")


def main():
    if len(sys.argv) != 4:
        print(__doc__)
        sys.exit(1)

    port      = sys.argv[1]
    track_num = int(sys.argv[2])
    path      = sys.argv[3]

    if not 0 <= track_num < 32:
        sys.exit("track_num must be 0–31")

    wav = load_wav(path)
    if not wav:
        sys.exit("File is empty after conversion")

    upload(port, track_num, wav)


if __name__ == "__main__":
    main()