import ctypes

import numpy as np
import pytest

from fasr.data.audio import AudioChunk, AudioChunkList
from fasr.data.stream import AudioBytesStream


class TestAudioBytesStream:
    def test_init(self):
        stream = AudioBytesStream(sample_rate=16000, chunk_size_ms=100)
        assert stream._sample_rate == 16000
        expected_bytes = 16000 // 1000 * 100 * ctypes.sizeof(ctypes.c_int16)
        assert stream._bytes_per_frame == expected_bytes

    def test_push_exact_frame(self):
        sr = 16000
        chunk_ms = 100
        stream = AudioBytesStream(sample_rate=sr, chunk_size_ms=chunk_ms)

        samples_per_frame = sr // 1000 * chunk_ms
        data = np.zeros(samples_per_frame, dtype=np.int16).tobytes()
        chunks = stream.push(data)
        assert len(chunks) == 1
        assert isinstance(chunks[0], AudioChunk)

    def test_push_multiple_frames(self):
        sr = 16000
        chunk_ms = 100
        stream = AudioBytesStream(sample_rate=sr, chunk_size_ms=chunk_ms)

        samples_per_frame = sr // 1000 * chunk_ms
        data = np.zeros(samples_per_frame * 3, dtype=np.int16).tobytes()
        chunks = stream.push(data)
        assert len(chunks) == 3

    def test_push_partial_then_complete(self):
        sr = 16000
        chunk_ms = 100
        stream = AudioBytesStream(sample_rate=sr, chunk_size_ms=chunk_ms)

        samples_per_frame = sr // 1000 * chunk_ms
        half = np.zeros(samples_per_frame // 2, dtype=np.int16).tobytes()

        chunks = stream.push(half)
        assert len(chunks) == 0

        chunks = stream.push(half)
        assert len(chunks) == 1

    def test_push_incremental(self):
        sr = 16000
        chunk_ms = 100
        stream = AudioBytesStream(sample_rate=sr, chunk_size_ms=chunk_ms)

        samples_per_frame = sr // 1000 * chunk_ms
        total_chunks = 0
        for _ in range(10):
            small = np.zeros(samples_per_frame // 5, dtype=np.int16).tobytes()
            chunks = stream.push(small)
            total_chunks += len(chunks)
        assert total_chunks == 2

    def test_flush_remaining(self):
        """flush 时 buffer 是 bytearray 而非 ndarray，会导致 Waveform 验证失败（已知 bug）"""
        sr = 16000
        chunk_ms = 100
        stream = AudioBytesStream(sample_rate=sr, chunk_size_ms=chunk_ms)

        samples = sr // 1000 * chunk_ms + 100
        data = np.zeros(samples, dtype=np.int16).tobytes()
        chunks = stream.push(data)
        assert len(chunks) == 1

        from pydantic import ValidationError

        with pytest.raises(ValidationError):
            stream.flush()

    def test_flush_empty(self):
        stream = AudioBytesStream(sample_rate=16000, chunk_size_ms=100)
        chunks = stream.flush()
        assert len(chunks) == 0

    def test_stream_id_consistent(self):
        sr = 16000
        chunk_ms = 100
        stream = AudioBytesStream(sample_rate=sr, chunk_size_ms=chunk_ms)

        samples = sr // 1000 * chunk_ms * 3
        data = np.zeros(samples, dtype=np.int16).tobytes()
        chunks = stream.push(data)

        ids = {c.stream_id for c in chunks}
        assert len(ids) == 1

    def test_waveform_dtype(self):
        sr = 16000
        chunk_ms = 100
        stream = AudioBytesStream(sample_rate=sr, chunk_size_ms=chunk_ms)

        samples = sr // 1000 * chunk_ms
        data = np.ones(samples, dtype=np.int16).tobytes()
        chunks = stream.push(data)
        assert chunks[0].waveform.data.dtype == np.float32

    def test_multichannel(self):
        sr = 16000
        chunk_ms = 100
        num_channels = 2
        stream = AudioBytesStream(
            sample_rate=sr, num_channels=num_channels, chunk_size_ms=chunk_ms
        )

        samples = sr // 1000 * chunk_ms * num_channels
        data = np.zeros(samples, dtype=np.int16).tobytes()
        chunks = stream.push(data)
        assert len(chunks) == 1
