import numpy as np
import pytest

from fasr.data.audio import (
    Audio,
    AudioChannel,
    AudioChannelList,
    AudioChunk,
    AudioChunkList,
    AudioSpan,
    AudioSpanList,
    AudioToken,
    AudioTokenList,
)
from fasr.data.waveform import Waveform


@pytest.fixture
def sample_rate():
    return 16000


@pytest.fixture
def sine_wave(sample_rate):
    """1 秒 440Hz 正弦波"""
    t = np.linspace(0, 1, sample_rate, dtype=np.float32)
    return np.sin(2 * np.pi * 440 * t)


@pytest.fixture
def short_wave(sample_rate):
    """0.5 秒正弦波"""
    t = np.linspace(0, 0.5, sample_rate // 2, dtype=np.float32)
    return np.sin(2 * np.pi * 440 * t)


@pytest.fixture
def waveform(sine_wave, sample_rate):
    return Waveform(data=sine_wave, sample_rate=sample_rate)


@pytest.fixture
def short_waveform(short_wave, sample_rate):
    return Waveform(data=short_wave, sample_rate=sample_rate)


@pytest.fixture
def audio_token():
    return AudioToken(text="你", start_ms=0, end_ms=100, follow="")


@pytest.fixture
def audio_token_list():
    return AudioTokenList(
        [
            AudioToken(text="你", start_ms=0, end_ms=100, follow=""),
            AudioToken(text="好", start_ms=100, end_ms=200, follow=""),
            AudioToken(text="世", start_ms=200, end_ms=300, follow=""),
            AudioToken(text="界", start_ms=300, end_ms=400, follow=" "),
        ]
    )


@pytest.fixture
def audio_span(audio_token_list, waveform):
    return AudioSpan(
        start_ms=0,
        end_ms=1000,
        waveform=waveform,
        tokens=audio_token_list,
    )


@pytest.fixture
def audio_channel(waveform, audio_token_list):
    segments = AudioSpanList(
        [
            AudioSpan(
                start_ms=0,
                end_ms=500,
                waveform=waveform,
                tokens=audio_token_list,
            )
        ]
    )
    return AudioChannel(waveform=waveform, segments=segments)


@pytest.fixture
def audio_with_channel(waveform, sample_rate):
    channel = AudioChannel(waveform=waveform)
    audio = Audio(sample_rate=sample_rate)
    audio.channels = AudioChannelList([channel])
    audio.duration = waveform.duration
    return audio


@pytest.fixture
def audio_chunk(waveform):
    return AudioChunk(stream_id="test-stream-id", waveform=waveform)
