#!/usr/bin/env python3
"""
Simple script to test if your specific Python environment works correctly
"""

import subprocess
import os
import asyncio

# Your specific environment path
ENV_PATH = "/var/www/eduai.edurigo.com/doc_train/edurigo_ai/Puru/ai4bharat/ai4bharat_env"

def get_python_executable():
    """Get the Python executable from the virtual environment"""
    return os.path.join(ENV_PATH, "bin", "python")

async def test_environment():
    """Test if the environment works"""
    
    python_exe = get_python_executable()
    
    print(f"Testing environment: {ENV_PATH}")
    print(f"Python executable: {python_exe}")
    print("-" * 50)
    
    # Test 1: Check if Python executable exists
    if not os.path.exists(python_exe):
        print("❌ FAILED: Python executable not found!")
        print(f"   Expected at: {python_exe}")
        return False
    else:
        print("✅ Python executable found")
    
    # Test 2: Check if Python runs
    try:
        process = await asyncio.create_subprocess_exec(
            python_exe, "-c", "import sys; print(f'Python version: {sys.version}'); print(f'Python path: {sys.executable}')",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        stdout, stderr = await process.communicate()
        
        if process.returncode == 0:
            print("✅ Python executes successfully")
            print(f"   Output: {stdout.decode().strip()}")
        else:
            print("❌ FAILED: Python execution failed")
            print(f"   Error: {stderr.decode().strip()}")
            return False
            
    except Exception as e:
        print(f"❌ FAILED: Exception running Python: {e}")
        return False
    
    # Test 3: Check if required packages are installed
    test_script = '''
import sys
try:
    import transformers
    print(f"✅ transformers version: {transformers.__version__}")
except ImportError as e:
    print(f"❌ transformers not found: {e}")

try:
    import torch
    print(f"✅ torch version: {torch.__version__}")
except ImportError as e:
    print(f"❌ torch not found: {e}")

try:
    import soundfile
    print(f"✅ soundfile available")
except ImportError as e:
    print(f"❌ soundfile not found: {e}")

try:
    import numpy
    print(f"✅ numpy version: {numpy.__version__}")
except ImportError as e:
    print(f"❌ numpy not found: {e}")

try:
    import scipy
    print(f"✅ scipy available")
except ImportError as e:
    print(f"❌ scipy not found: {e}")
'''
    
    try:
        process = await asyncio.create_subprocess_exec(
            python_exe, "-c", test_script,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        stdout, stderr = await process.communicate()
        
        print("\n📦 Package availability:")
        print(stdout.decode().strip())
        
        if stderr.decode().strip():
            print(f"\nWarnings/Errors: {stderr.decode().strip()}")
            
    except Exception as e:
        print(f"❌ FAILED: Exception testing packages: {e}")
        return False
    
    # Test 4: Test simple model loading (optional, might take time)
    print("\n🤖 Testing model loading (this may take a while)...")
    
    model_test_script = '''
try:
    from transformers import VitsModel, AutoTokenizer
    print("✅ Successfully imported transformers classes")
    
    # Try to load the model (this will download if not cached)
    model = VitsModel.from_pretrained("facebook/mms-tts-hin")
    tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-hin")
    print("✅ Model and tokenizer loaded successfully")
    
    # Quick test
    inputs = tokenizer("हैलो", return_tensors="pt")
    print("✅ Tokenization works")
    
    print("🎉 ALL TESTS PASSED!")
    
except Exception as e:
    print(f"❌ Model loading failed: {e}")
'''
    
    try:
        process = await asyncio.create_subprocess_exec(
            python_exe, "-c", model_test_script,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        stdout, stderr = await process.communicate()
        
        print(stdout.decode().strip())
        
        if stderr.decode().strip():
            print(f"Stderr: {stderr.decode().strip()}")
            
        if process.returncode == 0:
            print("\n🎉 Environment test completed successfully!")
            return True
        else:
            print("\n❌ Environment test failed during model loading")
            return False
            
    except Exception as e:
        print(f"❌ FAILED: Exception during model test: {e}")
        return False

def main():
    """Main function to run the test"""
    print("🔍 Environment Testing Script")
    print("=" * 50)
    
    # Run the async test
    result = asyncio.run(test_environment())
    
    print("\n" + "=" * 50)
    if result:
        print("🎉 RESULT: Environment is working correctly!")
        print("You can now use this environment for TTS generation.")
    else:
        print("❌ RESULT: Environment has issues!")
        print("\nTo fix:")
        print(f"1. rm -rf {ENV_PATH}")
        print(f"2. python3 -m venv {ENV_PATH}")
        print(f"3. source {ENV_PATH}/bin/activate")
        print("4. pip install transformers torch scipy soundfile numpy")

if __name__ == "__main__":
    main()