import os
import sys
import base64
import io
import types
import wave
import pytest
from pathlib import Path
from unittest.mock import MagicMock, patch
import numpy as np

BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))


def _make_mlx_stub():
    """Create minimal MLX stub for testing without Apple Silicon"""
    mlx = types.SimpleNamespace()
    mlx.core = types.SimpleNamespace()
    mx_array = type('mx.array', (), {'item': lambda self: 1})
    mlx.core.array = mx_array

    def mock_load(path):
        return MagicMock()
    mlx.core.load = mock_load  # type: ignore

    mlx.nn = types.SimpleNamespace()
    return mlx


def _make_mlx_audio_stub():
    """Create minimal mlx-audio stub"""
    stt = types.SimpleNamespace()
    stt.utils = types.SimpleNamespace()

    def mock_load(path):  # type: ignore
        model = MagicMock()
        output = types.SimpleNamespace()
        output.text = "识别结果"
        output.language = "zh-CN"
        model.generate = MagicMock(return_value=output)
        return model

    stt.utils.load = mock_load  # type: ignore

    qwen3_asr_mod = types.SimpleNamespace()
    qwen3_asr_mod.Qwen3ASRModel = type('Qwen3ASRModel', (), {})
    qwen3_asr_mod.ForcedAlignerModel = type('ForcedAlignerModel', (), {})
    stt.models = types.SimpleNamespace()  # type: ignore
    stt.models.qwen3_asr = qwen3_asr_mod  # type: ignore

    audio = types.SimpleNamespace()
    audio.stt = stt  # type: ignore
    return audio


def _reload_tts_asr_with_mocks():
    """Reload tts_asr with mocked MLX dependencies"""
    for mod_name in list(sys.modules.keys()):
        if 'tts_asr' in mod_name or 'mlx' in mod_name:
            del sys.modules[mod_name]

    mlx_stub = _make_mlx_stub()
    sys.modules['mlx'] = mlx_stub  # type: ignore
    sys.modules['mlx.core'] = mlx_stub.core  # type: ignore
    sys.modules['mlx.nn'] = mlx_stub.nn  # type: ignore

    audio_stub = _make_mlx_audio_stub()
    sys.modules['mlx-audio'] = audio_stub  # type: ignore
    sys.modules['mlx_audio'] = audio_stub  # type: ignore
    sys.modules['mlx_audio.stt'] = audio_stub.stt  # type: ignore
    sys.modules['mlx_audio.stt.utils'] = audio_stub.stt.utils  # type: ignore
    sys.modules['mlx_audio.stt.models'] = audio_stub.stt.models  # type: ignore
    sys.modules['mlx_audio.stt.models.qwen3_asr'] = audio_stub.stt.models.qwen3_asr  # type: ignore

    import tts_asr
    return tts_asr, audio_stub


@pytest.fixture(autouse=True)
def _clean_env():
    """Clean ASR-related env vars before/after each test"""
    saved = {}
    for k in ['HF_ENDPOINT']:
        saved[k] = os.environ.get(k)
        if k in os.environ:
            del os.environ[k]
    yield
    for k, v in saved.items():
        if v is not None:
            os.environ[k] = v  # type: ignore


def _make_wav_bytes(sr=16000, duration_sec=1.0, channels=1):
    """Helper: generate WAV bytes as base64"""
    samples = int(sr * duration_sec)
    audio = np.random.randint(-32768, 32767, size=samples * channels, dtype=np.int16)
    buf = io.BytesIO()
    with wave.open(buf, 'wb') as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(2)
        wf.setframerate(sr)
        wf.writeframes(audio.tobytes())
    return base64.b64encode(buf.getvalue()).decode()


class TestASRLazyLoading:
    """测试 ASR 模型懒加载"""

    def test_ensure_asr_loads_on_call(self):
        tts, audio_stub = _reload_tts_asr_with_mocks()
        assert tts._asr_model is None

        model = tts._ensure_asr_model()
        assert model is not None

    def test_ensure_align_loads_on_call(self):
        tts, audio_stub = _reload_tts_asr_with_mocks()
        assert tts._align_model is None

        model = tts._ensure_align_model()
        assert model is not None


class TestASREndpoint:
    """测试 ASR 端点逻辑"""

    def test_asr_basic_recognition(self, fastapi_testclient=None):
        """ASR 端点应正确返回识别结果"""
        tts, _ = _reload_tts_asr_with_mocks()

        # Mock the model to return known values
        tts._asr_model = MagicMock()
        output = types.SimpleNamespace()
        output.text = "你好世界"
        output.language = "zh-CN"
        tts._asr_model.generate.return_value = output

        wav_b64 = _make_wav_bytes()
        req = tts.ASRRequest(audio_base64=wav_b64)

        # Call generate directly (simulating endpoint logic)
        audio_bytes = base64.b64decode(req.audio_base64)
        wav_buffer = io.BytesIO(audio_bytes)
        with wave.open(wav_buffer, 'rb') as wf:
            raw = wf.readframes(wf.getnframes())
            arr = np.frombuffer(raw, dtype=np.int16)
            arr = arr.astype(np.float32) / 32768.0

        result = tts._asr_model.generate(arr, language=req.language)
        assert result.text == "你好世界"

    def test_asr_stereo_to_mono(self):
        """立体声音频应被正确转换为单声道"""
        wav_b64 = _make_wav_bytes(channels=2)

        audio_bytes = base64.b64decode(wav_b64)
        wav_buffer = io.BytesIO(audio_bytes)
        with wave.open(wav_buffer, 'rb') as wf:
            assert wf.getnchannels() == 2
            n_frames = wf.getnframes()
            raw_data = wf.readframes(n_frames)
            audio_array = np.frombuffer(raw_data, dtype=np.int16)

        # Convert to mono
        audio_array = np.mean(audio_array.reshape(-1, 2), axis=1)
        assert audio_array.ndim == 1

    def test_asr_resample_to_16k(self):
        """非 16kHz 音频应被重采样"""
        wav_b64 = _make_wav_bytes(sr=48000, duration_sec=0.5)

        audio_bytes = base64.b64decode(wav_b64)
        wav_buffer = io.BytesIO(audio_bytes)
        with wave.open(wav_buffer, 'rb') as wf:
            assert wf.getframerate() == 48000

    def test_asr_44100_resample(self):
        """44.1kHz 常见采样率应被重采样到 16k"""
        wav_b64 = _make_wav_bytes(sr=44100, duration_sec=1.0)

        audio_bytes = base64.b64decode(wav_b64)
        wav_buffer = io.BytesIO(audio_bytes)
        with wave.open(wav_buffer, 'rb') as wf:
            framerate = wf.getframerate()
            n_frames = wf.getnframes()
            raw_data = wf.readframes(n_frames)
            audio_array = np.frombuffer(raw_data, dtype=np.int16)

        # Simulate resample calculation
        if framerate != 16000:
            n_samples = int(len(audio_array) * 16000 / framerate)
        else:
            n_samples = len(audio_array)

        expected_16k_samples = int(1.0 * 16000)
        assert abs(n_samples - expected_16k_samples) < 2


class TestASRModelDownload:
    """测试 ASR 模型下载路径"""

    def test_load_asr_from_path_success(self):
        tts, _ = _reload_tts_asr_with_mocks()

        with patch('backend.tts_asr.snapshot_download', return_value='/fake/asr'):  # type: ignore
            tts._load_asr_models()

        assert tts._asr_model is not None

    def test_load_asr_skips_without_mlx(self):
        """不注入 MLX stub 时应跳过 ASR"""
        for mod_name in list(sys.modules.keys()):
            if 'tts_asr' in mod_name or 'mlx' in mod_name:
                del sys.modules[mod_name]

        import tts_asr  # noqa: F811
        assert tts_asr.Qwen3ASRModel is None

    def test_load_align_from_path(self):
        tts, _ = _reload_tts_asr_with_mocks()

        with patch('backend.tts_asr.snapshot_download', return_value='/fake/align'):  # type: ignore
            tts._load_asr_models()

        assert tts._align_model is not None


class TestModelConstants:
    """测试模型 ID 常量"""

    def test_asr_model_id(self):
        tts, _ = _reload_tts_asr_with_mocks()
        assert "aufklarer" in tts.ASR_MODEL_ID_MS

    def test_align_model_id(self):
        tts, _ = _reload_tts_asr_with_mocks()
        assert "ForcedAligner" in tts.ALIGN_MODEL_ID_MS

    def test_tts_model_id(self):
        tts, _ = _reload_tts_asr_with_mocks()
        assert "Qwen3-TTS" in tts.MODEL_ID_MS


class TestHFEndpointMirror:
    """测试镜像站配置"""

    def test_hf_endpoint_set(self):
        tts, _ = _reload_tts_asr_with_mocks()
        assert os.environ.get("HF_ENDPOINT") == "https://hf-mirror.com"

    def test_hf_endpoint_default(self):
        """即使环境变量未设置，模块也应默认设置镜像"""
        for mod_name in list(sys.modules.keys()):
            if 'tts_asr' in mod_name or 'mlx' in mod_name:
                del sys.modules[mod_name]

        if "HF_ENDPOINT" in os.environ:
            del os.environ["HF_ENDPOINT"]

        import tts_asr  # noqa: F811
        assert os.environ.get("HF_ENDPOINT") == "https://hf-mirror.com"
