#!/usr/bin/env python3
"""
直接播放测试音频到 ESP32-S3，绕过 Flask 服务器。
用法: python3 play_test.py <文件名或编号> [串口]

可用测试文件:
  1  sweep_100_8000.wav  - 对数扫频 100Hz→8kHz, 3s (测频响)
  2  square_1k.wav       - 1kHz 方波, 2s (测失真)
  3  vowel_a.wav         - 合成元音 "a", 1.5s
  4  vowel_i.wav         - 合成元音 "i", 1.5s (高频更多)
  5  five_tones.wav      - 5个单频 500/1k/2k/4k/6kHz
  6  noise_burst.wav     - 白噪声脉冲, 1s (测宽带)
  7  synth_counting.wav  - 合成数字 1-2-3-4-5 (模拟语音节奏)
  8  gap_test.wav        - 扫频+0.5s静音+扫频 (测PLL保持)
  9  am_1k_4hz.wav       - 1kHz载波, 4Hz包络 (模拟语音幅度)
  10 click_train.wav     - 10个脉冲 (测瞬态响应)
"""

import os
import sys
import struct
import time
import serial
import serial.tools.list_ports

TEST_DIR = os.path.join(os.path.dirname(__file__), "test_audio")

TEST_FILES = {
    "1": "sweep_100_8000.wav",
    "2": "square_1k.wav",
    "3": "vowel_a.wav",
    "4": "vowel_i.wav",
    "5": "five_tones.wav",
    "6": "noise_burst.wav",
    "7": "synth_counting.wav",
    "8": "gap_test.wav",
    "9": "am_1k_4hz.wav",
    "10": "click_train.wav",
}

SERIAL_BAUD = 921600
CHUNK_SIZE = 2048
SEND_DELAY = 0.025  # 25ms ≈ 82KB/s


def parse_wav(filepath):
    with open(filepath, "rb") as f:
        data = f.read()

    if len(data) < 44 or data[:4] != b"RIFF" or data[8:12] != b"WAVE":
        print("[Error] Not a valid WAV file")
        return None

    offset = 12
    fmt_parsed = False
    sample_rate = 0
    bits_per_sample = 0
    channels = 0

    while offset < len(data) - 8:
        chunk_id = data[offset:offset+4]
        chunk_size = struct.unpack_from("<I", data, offset + 4)[0]

        if chunk_id == b"fmt ":
            _, channels, sample_rate = struct.unpack_from("<HHH", data, offset + 8)
            bits_per_sample = struct.unpack_from("<H", data, offset + 8 + 14)[0]
            fmt_parsed = True
        elif chunk_id == b"data":
            if not fmt_parsed:
                print("[Error] data chunk before fmt")
                return None
            pcm_data = data[offset + 8 : offset + 8 + chunk_size]
            return pcm_data, sample_rate, bits_per_sample, channels

        offset += 8 + chunk_size
        if chunk_size % 2 == 1:
            offset += 1

    print("[Error] No data chunk found")
    return None


def send_pcm(ser, pcm_data):
    total = len(pcm_data)
    sent = 0
    print(f"[Serial] Sending {total} bytes PCM to ESP32...")

    ser.reset_input_buffer()
    ser.reset_output_buffer()

    while sent < total:
        end = min(sent + CHUNK_SIZE, total)
        chunk = pcm_data[sent:end]
        written = ser.write(chunk)
        ser.flush()
        sent += written
        time.sleep(SEND_DELAY)
        if sent % (CHUNK_SIZE * 5) == 0 or sent == total:
            print(f"[Serial] Progress: {sent}/{total} bytes")

    print(f"[Serial] Send complete ({sent} bytes)")


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        print("Usage: python3 play_test.py <1-6 或文件名> [串口]")
        print(f"\n可用串口: {[p.device for p in serial.tools.list_ports.comports()]}")
        sys.exit(1)

    arg = sys.argv[1]
    if arg in TEST_FILES:
        filename = TEST_FILES[arg]
    else:
        filename = arg

    filepath = os.path.join(TEST_DIR, filename)
    if not os.path.exists(filepath):
        print(f"[Error] File not found: {filepath}")
        print(f"Available: {list(TEST_FILES.values())}")
        sys.exit(1)

    # Serial port
    if len(sys.argv) >= 3:
        port = sys.argv[2]
    else:
        ports = [p.device for p in serial.tools.list_ports.comports()]
        if not ports:
            print("[Error] No serial port found")
            sys.exit(1)
        port = ports[0]
        if len(ports) > 1:
            print(f"[Info] Multiple ports found, using first: {port}")
            print(f"       Available: {ports}")

    print(f"[WAV] Loading: {filename}")
    result = parse_wav(filepath)
    if result is None:
        sys.exit(1)

    pcm_data, sample_rate, bits, channels = result
    print(f"[WAV] {sample_rate}Hz, {bits}bit, {channels}ch, {len(pcm_data)} bytes PCM")

    if sample_rate != 16000:
        print(f"[Warning] Sample rate {sample_rate} != 16000, ESP32 will play at wrong speed!")
    if bits != 16:
        print(f"[Warning] Bits per sample {bits} != 16!")

    print(f"[Serial] Opening {port} @ {SERIAL_BAUD}")
    try:
        ser = serial.Serial(port, SERIAL_BAUD, timeout=1)
        ser.dtr = False
        ser.rts = False
    except Exception as e:
        print(f"[Error] Failed to open serial: {e}")
        sys.exit(1)

    try:
        send_pcm(ser, pcm_data)
    finally:
        ser.close()
        print("[Serial] Port closed")


if __name__ == "__main__":
    main()
