"""
Test Qwen3-TTS voice cloning: load model, extract voice from reference audio, synthesize.
Usage: python test_modules/test_qwen3_tts_clone.py
"""
import sys
import time
import numpy as np
import soundfile as sf

print("=" * 60)
print("Testing Qwen3-TTS Voice Cloning")
print("=" * 60)

# Generate a synthetic reference audio (sine wave simulating speech duration)
# In production this would be the parent's recorded voice
print("Creating synthetic reference audio for testing...")
sr = 16000
duration = 5  # 5 seconds
t = np.linspace(0, duration, sr * duration, dtype=np.float32)
# Simple multi-frequency signal (not real speech, but tests the pipeline)
ref_audio = 0.3 * np.sin(2 * np.pi * 200 * t) + 0.2 * np.sin(2 * np.pi * 400 * t)
ref_audio_path = "sample_sounds/test_ref_audio.wav"
sf.write(ref_audio_path, ref_audio, sr)
print(f"[OK] Reference audio created: {ref_audio_path} ({duration}s)")

print()
print("Loading Qwen3-TTS model (this downloads ~1.7GB on first run)...")
start = time.time()

try:
    import torch
    from qwen_tts import Qwen3TTSModel

    model = Qwen3TTSModel.from_pretrained(
        "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
        device_map="cuda" if torch.cuda.is_available() else "cpu",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    )
    elapsed = time.time() - start
    print(f"[OK] Qwen3-TTS loaded in {elapsed:.1f}s")
except Exception as e:
    print(f"[FAIL] Model loading failed: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# Test 1: Voice cloning with x_vector_only_mode (speaker embedding only, no ref_text needed)
print()
print("Test 1: Voice clone with x_vector_only_mode=True...")
try:
    start = time.time()
    audio_list, sample_rate = model.generate_voice_clone(
        text="Hello! I am reading a bedtime story for you tonight.",
        language="en",
        ref_audio=ref_audio_path,
        x_vector_only_mode=True,
    )
    elapsed = time.time() - start
    total_samples = sum(len(a) for a in audio_list)
    duration_out = total_samples / sample_rate
    print(f"[OK] Voice clone (x_vector) in {elapsed:.1f}s, output: {duration_out:.1f}s at {sample_rate}Hz")

    # Save output
    output = np.concatenate(audio_list)
    sf.write("sample_sounds/test_clone_xvector.wav", output, sample_rate)
    print(f"[OK] Saved to sample_sounds/test_clone_xvector.wav")
except Exception as e:
    print(f"[FAIL] x_vector clone failed: {e}")
    import traceback
    traceback.print_exc()

# Test 2: Create and reuse voice clone prompt (for caching)
print()
print("Test 2: Create reusable voice clone prompt...")
try:
    start = time.time()
    prompt_items = model.create_voice_clone_prompt(
        ref_audio=ref_audio_path,
        x_vector_only_mode=True,
    )
    elapsed = time.time() - start
    print(f"[OK] Voice clone prompt created in {elapsed:.1f}s")

    # Reuse cached prompt for synthesis
    start = time.time()
    audio_list2, sr2 = model.generate_voice_clone(
        text="Once upon a time, there was a little rabbit named Peter.",
        language="en",
        voice_clone_prompt=prompt_items,
    )
    elapsed = time.time() - start
    total_samples2 = sum(len(a) for a in audio_list2)
    duration_out2 = total_samples2 / sr2
    print(f"[OK] Synthesis from cached prompt in {elapsed:.1f}s, output: {duration_out2:.1f}s")

    output2 = np.concatenate(audio_list2)
    sf.write("sample_sounds/test_clone_cached.wav", output2, sr2)
    print(f"[OK] Saved to sample_sounds/test_clone_cached.wav")
except Exception as e:
    print(f"[FAIL] Cached prompt failed: {e}")
    import traceback
    traceback.print_exc()

print()
print("=" * 60)
print("[OK] Qwen3-TTS voice cloning tests complete")
print("=" * 60)