import pyaudiowpatch as pyaudio
import numpy as np
import serial

# =========================
# SETTINGS
# =========================

COM_PORT = "COM4"
BAUD = 2000000

CHUNK = 480

# =========================
# FIND WASAPI LOOPBACK
# =========================

p = pyaudio.PyAudio()

print("WASAPI devices:")
print()

for i in range(p.get_device_count()):
    info = p.get_device_info_by_index(i)

    if info["isLoopbackDevice"]:
        print(
            i,
            info["name"],
            "channels:", info["maxInputChannels"],
            "rate:", info["defaultSampleRate"]
        )

# =========================
# GET DEFAULT LOOPBACK
# =========================

wasapi_info = p.get_host_api_info_by_type(
    pyaudio.paWASAPI
)

default_output = p.get_device_info_by_index(
    wasapi_info["defaultOutputDevice"]
)

print("\nDefault Windows output:")
print(default_output["name"])

# Find its loopback device
loopback_device = None

for i in range(p.get_device_count()):

    info = p.get_device_info_by_index(i)

    if (
        info["isLoopbackDevice"]
        and default_output["name"] in info["name"]
    ):
        loopback_device = info
        break

if loopback_device is None:
    raise RuntimeError("Could not find WASAPI loopback device")

print("\nUsing loopback:")
print(loopback_device["name"])

# =========================
# SERIAL TO XIAO
# =========================

ser = serial.Serial(COM_PORT, BAUD)

# =========================
# AUDIO STREAM
# =========================

stream = p.open(
    format=pyaudio.paFloat32,
    channels=loopback_device["maxInputChannels"],
    rate=int(loopback_device["defaultSampleRate"]),
    input=True,
    input_device_index=loopback_device["index"],
    frames_per_buffer=CHUNK
)

print("\nStreaming computer audio to XIAO...")
print("Press Ctrl+C to stop.")

try:

    while True:

        data = stream.read(
            CHUNK,
            exception_on_overflow=False
        )

        audio = np.frombuffer(
            data,
            dtype=np.float32
        )

        # Stereo -> mono
        audio = audio.reshape(
            -1,
            loopback_device["maxInputChannels"]
        )

        audio = np.mean(audio, axis=1)
        # audio = audio[::3]

        # [-1, +1] -> signed 16-bit PCM
        samples = np.clip(audio, -1.0, 1.0)
        samples = (samples * 32767).astype(np.int16)

        ser.write(samples.tobytes())

except KeyboardInterrupt:
    print("\nStopping...")

finally:
    stream.stop_stream()
    stream.close()
    ser.close()
    p.terminate()