#!/usr/bin/env python3
"""Test script to verify pre-downloading XTTS model and using local path."""

import os
import sys
from pathlib import Path

# Set env vars BEFORE TTS import - ALL of them needed to suppress prompts
os.environ['TTS_HOME'] = str(Path('E:/tmp/tts_models').absolute())
os.environ['TTS_CPML'] = '1'
os.environ['TTS_SKIP_TOS'] = '1'
os.environ['TTS_DISABLE_WEB_VERSION_PROMPT'] = '1'
os.environ['COQUI_TOS_AGREED'] = '1'

# Also try monkey-patching input to auto-answer
import builtins
original_input = builtins.input

def auto_input(prompt=""):
    print(prompt + "y")  # Print the prompt with auto-answer
    return "y"

builtins.input = auto_input

print("=" * 60)
print("XTTS Model Pre-Download Test")
print("=" * 60)

# Step 1: Download model to local path
print("\n[Step 1] Downloading XTTS model...")
try:
    from TTS.api import TTS
    
    # This will auto-download to TTS_HOME
    print("Initializing TTS (will auto-download model)...")
    model = TTS(
        model_name="tts_models/multilingual/multi-dataset/xtts_v2",
        gpu=False
    )
    print("✓ Model downloaded successfully!")
    
    # Find where it was downloaded
    tts_home = os.environ.get('TTS_HOME', os.path.expanduser('~/.tts_models'))
    model_path = Path(tts_home) / "tts_models--multilingual--multi-dataset--xtts_v2"
    
    if model_path.exists():
        print(f"✓ Model location: {model_path}")
        print(f"✓ Model files:")
        for item in sorted(model_path.iterdir()):
            print(f"  - {item.name}")
    
    # Step 2: Test loading from local path
    print("\n[Step 2] Testing local path loading...")
    del model  # Free memory
    
    # Try loading with explicit paths
    config_path = model_path / "config.json"
    
    if config_path.exists():
        print(f"✓ Config found at: {config_path}")
        print("\n[Step 3] Loading model from local path...")
        try:
            model2 = TTS(
                model_name="tts_models/multilingual/multi-dataset/xtts_v2",
                model_path=str(model_path),
                config_path=str(config_path),
                gpu=False
            )
            print("✓ Model loaded successfully from local path!")
            print("\n" + "=" * 60)
            print("SUCCESS: Pre-downloaded model loading works!")
            print("=" * 60)
        except Exception as e:
            print(f"✗ Failed to load from local path: {e}")
            print("\nTrying fallback: Loading with just model_name...")
            model2 = TTS(
                model_name="tts_models/multilingual/multi-dataset/xtts_v2",
                gpu=False
            )
            print("✓ Fallback load successful (auto-download approach)")
    else:
        print(f"✗ Config not found at {config_path}")
        print("Model structure:")
        for item in model_path.iterdir():
            print(f"  {item}")
            
except KeyboardInterrupt:
    print("\n\n✗ Download interrupted by user")
    sys.exit(1)
except Exception as e:
    print(f"✗ Error: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)
finally:
    # Restore original input
    builtins.input = original_input