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

# ==================== AudioToken ====================


class TestAudioToken:
    def test_create(self):
        token = AudioToken(text="你", start_ms=0, end_ms=100)
        assert token.text == "你"
        assert token.start_ms == 0
        assert token.end_ms == 100
        assert token.follow == " "

    def test_duration_ms(self, audio_token):
        assert audio_token.duration_ms == 100

    def test_is_punctuation(self):
        assert AudioToken(text="，", start_ms=0, end_ms=50).is_punctuation()
        assert AudioToken(text="。", start_ms=0, end_ms=50).is_punctuation()
        assert not AudioToken(text="你", start_ms=0, end_ms=50).is_punctuation()
        assert not AudioToken(text="a", start_ms=0, end_ms=50).is_punctuation()

    def test_default_follow(self):
        token = AudioToken(text="x", start_ms=0, end_ms=10)
        assert token.follow == " "

    def test_custom_follow(self):
        token = AudioToken(text="x", start_ms=0, end_ms=10, follow="")
        assert token.follow == ""


# ==================== AudioTokenList ====================


class TestAudioTokenList:
    def test_duration_ms(self, audio_token_list):
        assert audio_token_list.duration_ms == 400

    def test_duration(self, audio_token_list):
        assert audio_token_list.duration == 0.4

    def test_text(self, audio_token_list):
        assert audio_token_list.text == "你好世界"

    def test_text_with_follow(self):
        tokens = AudioTokenList(
            [
                AudioToken(text="hello", start_ms=0, end_ms=100, follow=" "),
                AudioToken(text="world", start_ms=100, end_ms=200, follow=" "),
            ]
        )
        assert tokens.text == "hello world"

    def test_empty_list(self):
        tokens = AudioTokenList()
        assert tokens.text == ""


# ==================== AudioSpan ====================


class TestAudioSpan:
    def test_create(self, audio_span):
        assert audio_span.start_ms == 0
        assert audio_span.end_ms == 1000

    def test_duration_ms(self, audio_span):
        assert audio_span.duration_ms == 1000

    def test_duration(self, audio_span):
        assert audio_span.duration == 1.0

    def test_text(self, audio_span):
        assert audio_span.text == "你好世界"

    def test_text_no_tokens(self):
        span = AudioSpan(start_ms=0, end_ms=1000, tokens=None)
        assert span.text == ""

    def test_text_setter(self):
        span = AudioSpan(start_ms=0, end_ms=1000)
        span.text = "test"
        assert span.text == "test"

    def test_sample_rate(self, audio_span, sample_rate):
        assert audio_span.sample_rate == sample_rate

    def test_sample_rate_no_waveform(self):
        span = AudioSpan(start_ms=0, end_ms=1000)
        assert span.sample_rate is None

    def test_lt(self):
        s1 = AudioSpan(start_ms=0, end_ms=500)
        s2 = AudioSpan(start_ms=0, end_ms=1000)
        assert s1 < s2
        assert not s2 < s1

    def test_gt(self):
        s1 = AudioSpan(start_ms=0, end_ms=1000)
        s2 = AudioSpan(start_ms=0, end_ms=500)
        assert s1 > s2
        assert not s2 > s1

    def test_getitem(self, audio_span):
        token = audio_span[0]
        assert isinstance(token, AudioToken)
        assert token.text == "你"

    def test_getitem_none_tokens(self):
        span = AudioSpan(start_ms=0, end_ms=1000, tokens=None)
        assert span[0] is None

    def test_len(self, audio_span):
        assert len(audio_span) == 4

    def test_len_no_tokens(self):
        span = AudioSpan(start_ms=0, end_ms=1000, tokens=None)
        assert len(span) == 0


# ==================== AudioSpanList ====================


class TestAudioSpanList:
    def test_duration_ms(self):
        spans = AudioSpanList(
            [
                AudioSpan(start_ms=0, end_ms=500),
                AudioSpan(start_ms=500, end_ms=1500),
            ]
        )
        assert spans.duration_ms == 1500

    def test_duration(self):
        spans = AudioSpanList(
            [
                AudioSpan(start_ms=0, end_ms=1000),
                AudioSpan(start_ms=1000, end_ms=2000),
            ]
        )
        assert spans.duration == 2.0

    def test_padded_duration_ms(self):
        spans = AudioSpanList(
            [
                AudioSpan(start_ms=0, end_ms=300),
                AudioSpan(start_ms=300, end_ms=1300),
            ]
        )
        assert spans.padded_duration_ms == 1000 * 2

    def test_padded_duration_ms_empty(self):
        spans = AudioSpanList()
        assert spans.padded_duration_ms == 0

    def test_max_duration_ms(self):
        spans = AudioSpanList(
            [
                AudioSpan(start_ms=0, end_ms=300),
                AudioSpan(start_ms=300, end_ms=1300),
            ]
        )
        assert spans.max_duration_ms == 1000

    def test_max_duration_ms_empty(self):
        spans = AudioSpanList()
        assert spans.max_duration_ms == 0

    def test_text(self):
        tokens1 = AudioTokenList(
            [AudioToken(text="你好", start_ms=0, end_ms=100, follow="")]
        )
        tokens2 = AudioTokenList(
            [AudioToken(text="世界", start_ms=100, end_ms=200, follow="")]
        )
        spans = AudioSpanList(
            [
                AudioSpan(start_ms=0, end_ms=500, tokens=tokens1),
                AudioSpan(start_ms=500, end_ms=1000, tokens=tokens2),
            ]
        )
        assert spans.text == "你好世界"


# ==================== AudioChannel ====================


class TestAudioChannel:
    def test_create(self, audio_channel):
        assert audio_channel.waveform is not None
        assert audio_channel.segments is not None

    def test_duration(self, audio_channel):
        assert abs(audio_channel.duration - 1.0) < 0.01

    def test_duration_ms(self, audio_channel):
        assert abs(audio_channel.duration_ms - 1000.0) < 10

    def test_start_end_ms(self, audio_channel):
        assert audio_channel.start_ms == 0
        assert abs(audio_channel.end_ms - 1000.0) < 10

    def test_text_from_segments(self, audio_channel):
        assert audio_channel.text == "你好世界"

    def test_text_from_raw_text(self, waveform):
        ch = AudioChannel(waveform=waveform, raw_text="override text")
        assert ch.text == "override text"

    def test_text_setter(self, audio_channel):
        audio_channel.text = "new text"
        assert audio_channel.text == "new text"
        assert audio_channel.raw_text == "new text"

    def test_text_empty(self, waveform):
        ch = AudioChannel(waveform=waveform)
        assert ch.text == ""

    def test_set_waveform(self, sample_rate):
        ch = AudioChannel()
        data = np.zeros(sample_rate, dtype=np.float32)
        ch.set_waveform(data, sample_rate)
        assert ch.waveform is not None
        assert ch.waveform.sample_rate == sample_rate

    def test_set_text(self):
        ch = AudioChannel()
        ch.set_text("hello")
        assert ch.text == "hello"

    def test_resample(self, audio_channel):
        audio_channel.resample(8000)
        assert audio_channel.waveform.sample_rate == 8000

    def test_getitem(self, audio_channel):
        span = audio_channel[0]
        assert isinstance(span, AudioSpan)

    def test_getitem_no_segments(self, waveform):
        ch = AudioChannel(waveform=waveform)
        assert ch[0] is None

    def test_len(self, audio_channel):
        assert len(audio_channel) == 1

    def test_len_no_segments(self, waveform):
        ch = AudioChannel(waveform=waveform)
        assert len(ch) == 0

    def test_add_token(self, waveform):
        ch = AudioChannel(waveform=waveform)
        ch.add_token("测", start_ms=None, end_ms=None)
        assert ch.tokens is not None
        assert len(ch.tokens) == 1
        assert ch.tokens[0].text == "测"

    def test_is_recognized_true(self, audio_channel):
        assert audio_channel.is_recognized() is True

    def test_is_recognized_false_no_segments(self, waveform):
        ch = AudioChannel(waveform=waveform)
        assert ch.is_recognized() is False

    def test_is_recognized_false_no_tokens(self, waveform):
        span = AudioSpan(start_ms=0, end_ms=500, waveform=waveform, tokens=None)
        ch = AudioChannel(waveform=waveform, segments=AudioSpanList([span]))
        assert ch.is_recognized() is False

    def test_text_from_tokens(self, waveform):
        ch = AudioChannel(waveform=waveform)
        ch.add_token("你", follow="")
        ch.add_token("好", follow=" ")
        assert ch.text == "你好 "

    def test_text_priority(self, waveform, audio_token_list):
        """raw_text > sents > segments > tokens"""
        segments = AudioSpanList(
            [AudioSpan(start_ms=0, end_ms=500, tokens=audio_token_list)]
        )
        ch = AudioChannel(waveform=waveform, segments=segments, raw_text="优先")
        assert ch.text == "优先"


# ==================== Audio ====================


class TestAudio:
    def test_create_empty(self):
        audio = Audio()
        assert audio.url is None
        assert audio.channels is None

    def test_create_with_url(self):
        audio = Audio(url="/path/to/audio.wav")
        assert audio.url == "/path/to/audio.wav"

    def test_path_to_str_validator(self, tmp_path):
        from pathlib import Path

        audio = Audio(url=Path("/some/path.wav"))
        assert isinstance(audio.url, str)

    def test_is_loaded_false(self):
        audio = Audio()
        assert audio.is_loaded is False

    def test_is_loaded_true(self, audio_with_channel):
        assert audio_with_channel.is_loaded is True

    def test_is_loaded_empty_channels(self):
        audio = Audio()
        audio.channels = AudioChannelList()
        assert audio.is_loaded is False

    def test_duration_s(self, audio_with_channel):
        assert abs(audio_with_channel.duration_s - 1.0) < 0.01

    def test_duration_ms(self, audio_with_channel):
        assert abs(audio_with_channel.duration_ms - 1000.0) < 10

    def test_append_channel(self):
        audio = Audio()
        audio.append_channel()
        assert audio.channels is not None
        assert len(audio.channels) == 1

    def test_append_channel_with_channel(self, waveform):
        audio = Audio()
        ch = AudioChannel(waveform=waveform)
        audio.append_channel(ch)
        assert len(audio.channels) == 1
        assert audio.channels[0].waveform is not None

    def test_getitem(self, audio_with_channel):
        ch = audio_with_channel[0]
        assert isinstance(ch, AudioChannel)

    def test_getitem_none(self):
        audio = Audio()
        assert audio[0] is None

    def test_len(self, audio_with_channel):
        assert len(audio_with_channel) == 1

    def test_len_no_channels(self):
        assert len(Audio()) == 0

    def test_audio_format(self):
        assert Audio(url="test.wav").audio_format == "wav"
        assert Audio(url="test.mp3").audio_format == "mp3"
        assert Audio(url="test.amr").audio_format == "amr"
        assert Audio(url="test.flac").audio_format == "flac"
        assert Audio(url="test.m4a").audio_format == "m4a"
        assert Audio().audio_format is None

    def test_set_hotwords(self, audio_with_channel):
        audio_with_channel.set_hotwords(["hello", "world"])
        assert audio_with_channel.hotwords == frozenset(["hello", "world"])

    def test_resample(self, audio_with_channel):
        audio_with_channel.resample(8000)
        assert audio_with_channel.channels[0].waveform.sample_rate == 8000

    def test_resample_no_channels(self):
        audio = Audio()
        audio.resample(8000)

    def test_clear_waveform(self, audio_with_channel):
        audio_with_channel.clear_waveform()
        assert audio_with_channel.channels[0].waveform is None

    def test_clear_text(self, waveform, sample_rate):
        ch = AudioChannel(waveform=waveform, raw_text="some text")
        audio = Audio(sample_rate=sample_rate)
        audio.channels = AudioChannelList([ch])
        audio.clear_text()
        assert audio.channels[0].raw_text is None

    def test_clear_bad(self):
        audio = Audio(is_bad=True, bad_reason="broken", bad_component="loader")
        audio.clear_bad()
        assert audio.is_bad is None
        assert audio.bad_reason is None
        assert audio.bad_component is None

    def test_chunk(self, audio_with_channel):
        chunks = audio_with_channel.chunk(chunk_size_ms=200)
        assert isinstance(chunks, AudioChunkList)
        assert len(chunks) > 0
        assert chunks[-1].is_last is True
        assert chunks[0].is_start is True

    def test_chunk_small_audio(self, sample_rate):
        data = np.zeros(sample_rate // 10, dtype=np.float32)
        wf = Waveform(data=data, sample_rate=sample_rate)
        ch = AudioChannel(waveform=wf)
        audio = Audio(sample_rate=sample_rate, duration=0.1)
        audio.channels = AudioChannelList([ch])
        chunks = audio.chunk(chunk_size_ms=200)
        assert len(chunks) == 1
        assert chunks[0].is_last is True
        assert chunks[0].is_start is True

    def test_chunk_no_channels_raises(self):
        audio = Audio()
        with pytest.raises(ValueError):
            audio.chunk(100)

    @pytest.mark.xfail(
        reason="AudioChannel 未定义 bad_component 字段，align_channel_bad 访问该属性会报错"
    )
    def test_align_channel_bad(self, waveform, sample_rate):
        ch = AudioChannel(
            waveform=waveform,
            is_bad=True,
            bad_reason="noise",
            bad_component="loader",
        )
        audio = Audio(sample_rate=sample_rate)
        audio.channels = AudioChannelList([ch])
        audio.align_channel_bad()
        assert audio.is_bad is True

    def test_align_channel_not_bad(self, waveform, sample_rate):
        ch = AudioChannel(waveform=waveform)
        audio = Audio(sample_rate=sample_rate)
        audio.channels = AudioChannelList([ch])
        audio.align_channel_bad()
        assert audio.is_bad is None


# ==================== AudioChunk ====================


class TestAudioChunk:
    def test_create(self, audio_chunk):
        assert audio_chunk.stream_id == "test-stream-id"
        assert audio_chunk.waveform is not None

    def test_text_empty(self, audio_chunk):
        assert audio_chunk.text == ""

    def test_text_with_tokens(self, waveform):
        tokens = AudioTokenList(
            [
                AudioToken(text="你", start_ms=0, end_ms=100, follow=""),
                AudioToken(text="好", start_ms=100, end_ms=200, follow=""),
            ]
        )
        chunk = AudioChunk(stream_id="s", waveform=waveform, tokens=tokens)
        assert chunk.text == "你好"

    def test_duration_ms(self, audio_chunk):
        assert abs(audio_chunk.duration_ms - 1000.0) < 10

    def test_duration_s(self, audio_chunk):
        assert abs(audio_chunk.duration_s - 1.0) < 0.01

    def test_is_final(self):
        wf = Waveform(data=np.zeros(100, dtype=np.float32), sample_rate=16000)
        chunk = AudioChunk(stream_id="s", waveform=wf, is_last=True)
        assert chunk.is_final is True

    def test_clear(self, audio_chunk):
        """AudioChunk.waveform 是必填字段，不能设为 None，验证 clear 会报错"""
        from pydantic import ValidationError

        with pytest.raises(ValidationError):
            audio_chunk.clear()


# ==================== AudioChunkList ====================


class TestAudioChunkList:
    def _make_chunks(self, sample_rate=16000):
        chunks = AudioChunkList()
        for i in range(3):
            data = np.zeros(sample_rate // 2, dtype=np.float32)
            wf = Waveform(data=data, sample_rate=sample_rate)
            chunks.append(
                AudioChunk(
                    stream_id="s",
                    waveform=wf,
                    start_ms=i * 500,
                    end_ms=(i + 1) * 500,
                )
            )
        return chunks

    def test_concat(self):
        chunks = self._make_chunks()
        audio = chunks.concat()
        assert isinstance(audio, Audio)
        assert audio.channels is not None
        assert len(audio.channels) == 1

    def test_concat_empty(self):
        chunks = AudioChunkList()
        audio = chunks.concat()
        assert isinstance(audio, Audio)
        assert audio.channels is None

    def test_concat_to_chunk(self):
        chunks = self._make_chunks()
        merged = chunks.concat_to_chunk()
        assert isinstance(merged, AudioChunk)
        assert merged.start_ms == 0
        assert merged.end_ms == 1500

    def test_concat_to_chunk_single(self):
        wf = Waveform(data=np.zeros(100, dtype=np.float32), sample_rate=16000)
        chunks = AudioChunkList([AudioChunk(stream_id="s", waveform=wf)])
        merged = chunks.concat_to_chunk()
        assert merged is chunks[0]

    def test_concat_to_audio(self):
        chunks = self._make_chunks()
        audio = chunks.concat_to_audio()
        assert isinstance(audio, Audio)
        assert audio.channels is not None


# ==================== AudioList ====================


class TestAudioList:
    def test_from_urls(self):
        from fasr.data.audio import AudioList

        urls = ["/path/a.wav", "/path/b.wav"]
        audios = AudioList.from_urls(urls, load=False)
        assert len(audios) == 2
        assert audios[0].url == "/path/a.wav"

    def test_from_urls_single(self):
        from fasr.data.audio import AudioList

        audios = AudioList.from_urls("/path/a.wav", load=False)
        assert len(audios) == 1

    def test_str(self):
        from fasr.data.audio import AudioList

        audios = AudioList([Audio(url="a.wav"), Audio(url="b.wav", is_bad=True)])
        s = str(audios)
        assert "2 audios" in s
        assert "1 bad" in s

    def test_bad_audios(self):
        from fasr.data.audio import AudioList

        audios = AudioList(
            [
                Audio(url="a.wav"),
                Audio(url="b.wav", is_bad=True),
                Audio(url="c.wav", is_bad=True),
            ]
        )
        assert len(audios.bad_audios) == 2
        assert len(audios.good_audios) == 1

    def test_has_bad_audio(self):
        from fasr.data.audio import AudioList

        audios = AudioList([Audio(url="a.wav", is_bad=True)])
        assert audios.has_bad_audio() is True

        audios2 = AudioList([Audio(url="a.wav")])
        assert audios2.has_bad_audio() is False

    def test_group_by_hotwords(self):
        from fasr.data.audio import AudioList

        a1 = Audio(url="a.wav", hotwords=frozenset(["cat"]))
        a2 = Audio(url="b.wav", hotwords=frozenset(["cat"]))
        a3 = Audio(url="c.wav", hotwords=frozenset(["dog"]))
        a4 = Audio(url="d.wav")
        audios = AudioList([a1, a2, a3, a4])
        groups = audios.group_by_hotwords()
        assert len(groups) == 3
        assert len(groups[frozenset(["cat"])]) == 2

    def test_copy(self, waveform, sample_rate):
        from fasr.data.audio import AudioList

        ch = AudioChannel(waveform=waveform)
        audio = Audio(sample_rate=sample_rate, duration=1.0)
        audio.channels = AudioChannelList([ch])
        audios = AudioList([audio])
        copied = audios.copy()
        assert len(copied) == 1
        assert copied[0] is not audios[0]

    def test_sub(self):
        from fasr.data.audio import AudioList

        a1 = Audio(url="a.wav")
        a2 = Audio(url="b.wav")
        list1 = AudioList([a1, a2])
        list2 = AudioList([a1])
        diff = list1 - list2
        assert len(diff) == 1
