#!/usr/bin/env python3
"""
Create Compatible Environment Script
===================================

Creates a compatible Python environment for ECG-FM testing.
This addresses the version incompatibilities found in the lightweight test.
"""

import subprocess
import sys
import os

def check_python_version():
    """Check current Python version"""
    version = sys.version_info
    print(f"🐍 Current Python: {version.major}.{version.minor}.{version.micro}")
    
    if version.major == 3 and 8 <= version.minor <= 11:
        print("✅ Python version is compatible with HF Spaces")
        return True
    else:
        print("❌ Python version NOT compatible with HF Spaces")
        print("   Need Python 3.8-3.11, have 3.13.3")
        return False

def create_virtual_env():
    """Create a virtual environment with compatible Python"""
    print("\n🔧 Creating Virtual Environment")
    print("=" * 40)
    
    # Check if virtualenv is available
    try:
        result = subprocess.run([sys.executable, '-m', 'pip', 'install', 'virtualenv'], 
                              capture_output=True, text=True)
        if result.returncode == 0:
            print("✅ virtualenv installed")
        else:
            print("❌ Failed to install virtualenv")
            return False
    except Exception as e:
        print(f"❌ Error installing virtualenv: {e}")
        return False
    
    # Create virtual environment
    env_name = "ecg_fm_env"
    try:
        cmd = [sys.executable, '-m', 'virtualenv', env_name]
        print(f"Creating environment: {env_name}")
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print(f"✅ Virtual environment created: {env_name}")
            return True
        else:
            print("❌ Failed to create virtual environment")
            print(result.stderr)
            return False
    except Exception as e:
        print(f"❌ Error creating environment: {e}")
        return False

def install_compatible_packages():
    """Install compatible package versions"""
    print("\n📦 Installing Compatible Packages")
    print("=" * 40)
    
    # These versions are compatible with each other
    packages = [
        "numpy==1.24.3",
        "torch==1.13.1",
        "torchvision==0.14.1", 
        "torchaudio==0.13.1",
        "omegaconf==2.3.0",
        "fastapi==0.104.1",
        "uvicorn[standard]==0.24.0",
        "huggingface-hub==0.19.4",
        "transformers==4.21.0",
        "einops==0.7.0",
        "pyyaml==6.0.1"
    ]
    
    # Get the virtual environment Python
    if os.name == 'nt':  # Windows
        python_path = os.path.join("ecg_fm_env", "Scripts", "python.exe")
    else:  # Unix/Linux
        python_path = os.path.join("ecg_fm_env", "bin", "python")
    
    if not os.path.exists(python_path):
        print("❌ Virtual environment Python not found")
        return False
    
    print(f"Using Python: {python_path}")
    
    all_good = True
    
    for package in packages:
        print(f"Installing {package}...")
        try:
            cmd = [python_path, '-m', 'pip', 'install', package]
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
            
            if result.returncode == 0:
                print(f"  ✅ {package} installed")
            else:
                print(f"  ❌ {package} failed")
                print(f"     Error: {result.stderr}")
                all_good = False
                
        except subprocess.TimeoutExpired:
            print(f"  ⚠️  {package} timed out")
            all_good = False
        except Exception as e:
            print(f"  ❌ {package} error: {e}")
            all_good = False
    
    return all_good

def test_compatible_environment():
    """Test the compatible environment"""
    print("\n🧪 Testing Compatible Environment")
    print("=" * 40)
    
    if os.name == 'nt':  # Windows
        python_path = os.path.join("ecg_fm_env", "Scripts", "python.exe")
    else:  # Unix/Linux
        python_path = os.path.join("ecg_fm_env", "bin", "python")
    
    if not os.path.exists(python_path):
        print("❌ Virtual environment not found")
        return False
    
    # Test imports
    test_script = f"""
import sys
print(f"Python: {{sys.version}}")

try:
    import numpy
    print(f"NumPy: {{numpy.__version__}}")
except ImportError as e:
    print(f"NumPy: ❌ {{e}}")

try:
    import torch
    print(f"PyTorch: {{torch.__version__}}")
except ImportError as e:
    print(f"PyTorch: ❌ {{e}}")

try:
    import omegaconf
    print(f"OmegaConf: {{omegaconf.__version__}}")
except ImportError as e:
    print(f"OmegaConf: ❌ {{e}}")

try:
    import fastapi
    print(f"FastAPI: {{fastapi.__version__}}")
except ImportError as e:
    print(f"FastAPI: ❌ {{e}}")

try:
    import transformers
    print(f"Transformers: {{transformers.__version__}}")
except ImportError as e:
    print(f"Transformers: ❌ {{e}}")
"""
    
    try:
        result = subprocess.run([python_path, '-c', test_script], 
                              capture_output=True, text=True)
        
        if result.returncode == 0:
            print("✅ Environment test successful!")
            print("\n📋 Package Versions:")
            print(result.stdout)
            return True
        else:
            print("❌ Environment test failed!")
            print(result.stderr)
            return False
            
    except Exception as e:
        print(f"❌ Test error: {e}")
        return False

def main():
    """Main function"""
    print("🚀 Creating Compatible Environment for ECG-FM")
    print("=" * 60)
    print("This will create a virtual environment with compatible versions")
    print()
    
    # Step 1: Check Python version
    if not check_python_version():
        print("\n⚠️  Warning: Python 3.13.3 may cause issues")
        print("   Virtual environment will help isolate dependencies")
    
    # Step 2: Create virtual environment
    if not create_virtual_env():
        print("\n❌ Failed to create virtual environment")
        return 1
    
    # Step 3: Install compatible packages
    if not install_compatible_packages():
        print("\n❌ Failed to install compatible packages")
        return 1
    
    # Step 4: Test environment
    if not test_compatible_environment():
        print("\n❌ Environment test failed")
        return 1
    
    print("\n🎉 COMPATIBLE ENVIRONMENT CREATED!")
    print("✅ All packages installed with compatible versions")
    print("🚀 Ready for local testing before HF upload")
    
    # Instructions
    print("\n📋 How to use:")
    if os.name == 'nt':  # Windows
        print("   Activate: ecg_fm_env\\Scripts\\activate")
        print("   Python: ecg_fm_env\\Scripts\\python.exe")
    else:  # Unix/Linux
        print("   Activate: source ecg_fm_env/bin/activate")
        print("   Python: ecg_fm_env/bin/python")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())