import pytest
import numpy as np
import torch
from unittest.mock import MagicMock, patch
from vieneu.turbo import TurboVieNeuTTS, TurboGPUVieNeuTTS

@pytest.fixture
def mock_onnx_session():
    session = MagicMock()
    # Mock return value for decoder/encoder run (returns a list of outputs)
    session.run.return_value = [np.zeros((1, 1, 48000), dtype=np.float32)]
    return session

@pytest.fixture
def mock_llama_instance():
    llama = MagicMock()
    llama.return_value = {
        "choices": [{"text": "<|speech_1|><|speech_2|><|speech_3|>"}]
    }
    return llama

@patch("onnxruntime.InferenceSession")
@patch("llama_cpp.Llama")
@patch("huggingface_hub.hf_hub_download", return_value="dummy_path")
def test_turbo_gguf_init(mock_hf, mock_llama, mock_ort):
    tts = TurboVieNeuTTS(backbone_repo="dummy", device="cpu")
    assert tts.backbone is not None
    assert tts.decoder_sess is not None
    mock_llama.assert_called_once()

@patch("onnxruntime.InferenceSession")
@patch("llama_cpp.Llama")
@patch("huggingface_hub.hf_hub_download", return_value="dummy_path")
def test_turbo_gguf_infer(mock_hf, mock_llama, mock_ort, mock_onnx_session, mock_llama_instance):
    mock_ort.return_value = mock_onnx_session
    mock_llama.return_value = mock_llama_instance
    
    tts = TurboVieNeuTTS(backbone_repo="dummy", device="cpu")
    tts._preset_voices = {"test": {"codes": np.zeros(128), "text": "test"}}
    
    with patch("vieneu_utils.phonemize_text.phonemize_text", return_value="p-h-o-n-e-m-e-s"):
        audio = tts.infer("Xin chào", voice="test")
        assert isinstance(audio, np.ndarray)
        assert len(audio) > 0
        mock_llama_instance.assert_called()

@patch("onnxruntime.InferenceSession")
@patch("transformers.AutoTokenizer.from_pretrained")
@patch("transformers.AutoModelForCausalLM.from_pretrained")
def test_turbo_gpu_standard_init(mock_model, mock_tokenizer, mock_ort):
    # Mock model and tokenizer
    mock_tokenizer.return_value = MagicMock()
    mock_model_instance = MagicMock()
    mock_model_instance.to.return_value = mock_model_instance
    mock_model.return_value = mock_model_instance

    tts = TurboGPUVieNeuTTS(backbone_repo="dummy", device="cuda", backend="standard")
    assert tts.backend == "standard"
    assert tts.backbone is not None
    assert tts.device == "cuda"

@patch("onnxruntime.InferenceSession")
@patch("lmdeploy.pipeline")
def test_turbo_gpu_lmdeploy_init(mock_pipeline, mock_ort):
    mock_pipeline_instance = MagicMock()
    mock_pipeline.return_value = mock_pipeline_instance
    
    tts = TurboGPUVieNeuTTS(backbone_repo="dummy", device="cuda", backend="lmdeploy")
    assert tts.backend == "lmdeploy"
    assert tts.backbone is not None

@patch("onnxruntime.InferenceSession")
@patch("transformers.AutoTokenizer.from_pretrained")
@patch("transformers.AutoModelForCausalLM.from_pretrained")
def test_turbo_gpu_infer(mock_model, mock_tokenizer, mock_ort, mock_onnx_session):
    mock_ort.return_value = mock_onnx_session
    
    # Mock standard Transformers path
    mock_tokenizer_instance = MagicMock()
    mock_tokenizer_instance.return_value = {"input_ids": torch.zeros((1, 5), dtype=torch.long)}
    mock_tokenizer_instance.decode.return_value = "<|speech_100|><|speech_101|>"
    mock_tokenizer.return_value = mock_tokenizer_instance
    
    mock_model_instance = MagicMock()
    mock_model_instance.generate.return_value = torch.zeros((1, 20), dtype=torch.long)
    mock_model_instance.to.return_value = mock_model_instance
    mock_model.return_value = mock_model_instance

    tts = TurboGPUVieNeuTTS(backbone_repo="dummy", device="cpu", backend="standard")
    tts._preset_voices = {"test": {"codes": np.zeros(128), "text": "test"}}
    
    with patch("vieneu_utils.phonemize_text.phonemize_text", return_value="p-h-o-n-e-m-e-s"):
        audio = tts.infer("Xin chào", voice="test")
        assert isinstance(audio, np.ndarray)
        assert len(audio) > 0

@patch("onnxruntime.InferenceSession")
@patch("llama_cpp.Llama")
@patch("huggingface_hub.hf_hub_download", return_value="dummy_path")
def test_turbo_voice_cloning_encode(mock_hf, mock_llama, mock_ort, mock_onnx_session):
    mock_ort.return_value = mock_onnx_session
    # Mock encoder return
    mock_onnx_session.run.return_value = [np.zeros((1, 128), dtype=np.float32)]
    
    tts = TurboVieNeuTTS(backbone_repo="dummy", device="cpu")
    with patch("librosa.load", return_value=(np.zeros(24000), 24000)):
        emb = tts.encode_reference("dummy.wav")
        assert isinstance(emb, np.ndarray)
        assert emb.shape == (1, 128)
