# type: ignore

import numpy as np
import pytest

from livekit.rtc import AudioMixer
from livekit.rtc.utils import sine_wave_generator

SAMPLE_RATE = 48000
BLOCKSIZE = SAMPLE_RATE // 10


@pytest.mark.asyncio
async def test_mixer_two_sine_waves():
    """
    Test that mixing two sine waves (440Hz and 880Hz) produces an output
    containing both frequency components.
    """
    duration = 1.0
    mixer = AudioMixer(
        sample_rate=SAMPLE_RATE,
        num_channels=1,
        blocksize=BLOCKSIZE,
        stream_timeout_ms=100,
        capacity=100,
    )
    stream1 = sine_wave_generator(440, duration, SAMPLE_RATE)
    stream2 = sine_wave_generator(880, duration, SAMPLE_RATE)
    mixer.add_stream(stream1)
    mixer.add_stream(stream2)
    mixer.end_input()

    mixed_signals = []
    async for frame in mixer:
        data = np.frombuffer(frame.data.tobytes(), dtype=np.int16)
        mixed_signals.append(data)

    await mixer.aclose()

    if not mixed_signals:
        pytest.fail("No frames were produced by the mixer.")

    mixed_signal = np.concatenate(mixed_signals)

    # Use FFT to analyze frequency components.
    fft = np.fft.rfft(mixed_signal)
    freqs = np.fft.rfftfreq(len(mixed_signal), 1 / SAMPLE_RATE)
    magnitude = np.abs(fft)

    # Identify peak frequencies. We'll pick the 5 highest peaks.
    peak_indices = np.argsort(magnitude)[-5:]
    peak_freqs = freqs[peak_indices]

    print("Peak frequencies:", peak_freqs)

    # Assert that the peaks include 440Hz and 880Hz (with a tolerance of ±5 Hz)
    assert any(np.isclose(peak_freqs, 440, atol=5)), f"Expected 440 Hz in peaks, got: {peak_freqs}"
    assert any(np.isclose(peak_freqs, 880, atol=5)), f"Expected 880 Hz in peaks, got: {peak_freqs}"