import serial
import wave
import struct

# === CONFIGURATION ===
SERIAL_PORT = 'COM7'         # Update this to match your system
BAUD_RATE = 1000000
SAMPLE_RATE = 8000
DURATION_SECONDS = 5
OUTPUT_FILE = 'esp32_audio.wav'

# === SETUP SERIAL ===
ser = serial.Serial(SERIAL_PORT, BAUD_RATE, timeout=0.01)

# === PREPARE TO CAPTURE ===
print(f"Recording {DURATION_SECONDS} seconds of audio...")

samples = []
max_samples = SAMPLE_RATE * DURATION_SECONDS
chunk_size = 4096  # bytes

# === CAPTURE LOOP ===
while len(samples) < max_samples:
    data = ser.read(chunk_size)
    for i in range(0, len(data) - 1, 2):
        sample = int.from_bytes(data[i:i+2], byteorder='little', signed=True)
        samples.append(sample)
    if len(samples) % 1000 == 0:
        print(f"{len(samples)} samples captured")

ser.close()

# === NORMALIZE AUDIO ===
peak = max(abs(s) for s in samples)
scale = 32767 / peak if peak > 0 else 1.0
normalized = [int(max(-32768, min(32767, s * scale))) for s in samples]

# === WRITE TO WAV FILE ===
with wave.open(OUTPUT_FILE, 'w') as wav:
    wav.setnchannels(1)
    wav.setsampwidth(2)  # 16-bit
    wav.setframerate(SAMPLE_RATE)
    for sample in normalized:
        wav.writeframes(struct.pack('<h', sample))

print(f"Saved normalized audio to {OUTPUT_FILE}")