#!/usr/bin/env python3
"""
Save audio debug dump from ESP32-S3 to a WAV file.

Usage:
    python audio_debug_save.py                      # save to debug_aec.wav on port 7200
    python audio_debug_save.py -o out.wav -p 7200   # custom output and port

The script listens for raw 16-bit mono PCM at 16 kHz streamed over UDP,
and writes it into a standard WAV file.  Press Ctrl+C to stop and finalize.
"""

import argparse
import socket
import struct
import sys

SAMPLE_RATE = 16000
CHANNELS = 1
SAMPLE_WIDTH = 2  # 16-bit


def main():
    parser = argparse.ArgumentParser(description="Save UDP audio debug dump to WAV")
    parser.add_argument("-p", "--port", type=int, default=7200, help="UDP port (default: 7200)")
    parser.add_argument("-o", "--output", default="debug_aec.wav", help="Output WAV file")
    args = parser.parse_args()

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(("0.0.0.0", args.port))
    print(f"Listening on UDP port {args.port}, writing to {args.output}")
    print("Press Ctrl+C to stop and finalize WAV.")

    # Write WAV header with placeholder data size, patch on exit
    f = open(args.output, "wb")
    data_bytes = 0
    byte_rate = SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH
    block_align = CHANNELS * SAMPLE_WIDTH

    def write_wav_header(total_data_bytes):
        f.seek(0)
        f.write(b"RIFF")
        f.write(struct.pack("<I", 36 + total_data_bytes))
        f.write(b"WAVE")
        f.write(b"fmt ")
        f.write(struct.pack("<I", 16))  # fmt chunk size
        f.write(struct.pack("<H", 1))   # PCM format
        f.write(struct.pack("<H", CHANNELS))
        f.write(struct.pack("<I", SAMPLE_RATE))
        f.write(struct.pack("<I", byte_rate))
        f.write(struct.pack("<H", block_align))
        f.write(struct.pack("<H", SAMPLE_WIDTH * 8))
        f.write(b"data")
        f.write(struct.pack("<I", total_data_bytes))

    write_wav_header(0)

    try:
        while True:
            data, _ = sock.recvfrom(2048)
            f.write(data)
            data_bytes += len(data)
    except KeyboardInterrupt:
        pass
    finally:
        write_wav_header(data_bytes)
        f.close()
        sock.close()
        duration = data_bytes / byte_rate if byte_rate else 0
        print(f"\nSaved {data_bytes} bytes ({duration:.1f}s) to {args.output}")


if __name__ == "__main__":
    main()