#!/usr/bin/env python3
"""
Fix OpenBLAS threading issues for RuralPoint Django application.
This script addresses the resource exhaustion problems that prevent Django from starting.
"""

import os
import sys
import subprocess
from pathlib import Path

def set_environment_variables():
    """Set environment variables to limit OpenBLAS thread usage."""
    print("🔧 Setting OpenBLAS environment variables...")
    
    # Set OpenBLAS to use single thread to prevent resource exhaustion
    os.environ['OPENBLAS_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['BLAS_NUM_THREADS'] = '1'
    os.environ['LAPACK_NUM_THREADS'] = '1'
    
    print("✓ Environment variables set:")
    print(f"  OPENBLAS_NUM_THREADS: {os.environ.get('OPENBLAS_NUM_THREADS')}")
    print(f"  MKL_NUM_THREADS: {os.environ.get('MKL_NUM_THREADS')}")
    print(f"  OMP_NUM_THREADS: {os.environ.get('OMP_NUM_THREADS')}")

def test_django_import():
    """Test if Django can import without OpenBLAS errors."""
    print("\n🧪 Testing Django import...")
    
    try:
        # Set Django settings module
        os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'ruralpoint_system.settings')
        
        # Import Django and setup
        import django
        django.setup()
        
        print("✓ Django imported successfully!")
        return True
        
    except Exception as e:
        print(f"❌ Django import failed: {e}")
        return False

def test_django_commands():
    """Test basic Django management commands."""
    print("\n🔍 Testing Django commands...")
    
    commands_to_test = [
        ['python', 'manage.py', 'check', '--deploy'],
        ['python', 'manage.py', 'showmigrations'],
        ['python', 'manage.py', 'collectstatic', '--dry-run'],
    ]
    
    for cmd in commands_to_test:
        print(f"\nTesting: {' '.join(cmd)}")
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            if result.returncode == 0:
                print(f"✓ Command successful")
            else:
                print(f"⚠️ Command failed with return code {result.returncode}")
                if result.stderr:
                    print(f"Error: {result.stderr[:200]}...")
        except subprocess.TimeoutExpired:
            print("⚠️ Command timed out")
        except Exception as e:
            print(f"❌ Command error: {e}")

def create_environment_script():
    """Create a script to set environment variables for future runs."""
    print("\n📝 Creating environment setup script...")
    
    script_content = """#!/bin/bash
# RuralPoint OpenBLAS Environment Setup
# Run this script before starting Django

export OPENBLAS_NUM_THREADS=1
export MKL_NUM_THREADS=1
export OMP_NUM_THREADS=1
export BLAS_NUM_THREADS=1
export LAPACK_NUM_THREADS=1

echo "OpenBLAS environment variables set:"
echo "OPENBLAS_NUM_THREADS: $OPENBLAS_NUM_THREADS"
echo "MKL_NUM_THREADS: $MKL_NUM_THREADS"
echo "OMP_NUM_THREADS: $OMP_NUM_THREADS"

# Optional: Start Django with these settings
# python manage.py runserver
"""
    
    script_path = Path("set_openblas_env.sh")
    with open(script_path, 'w') as f:
        f.write(script_content)
    
    # Make it executable
    script_path.chmod(0o755)
    print(f"✓ Created {script_path}")

def create_python_environment_script():
    """Create a Python script to set environment variables."""
    print("\n📝 Creating Python environment setup script...")
    
    script_content = '''#!/usr/bin/env python3
# RuralPoint OpenBLAS Environment Setup for Python
# Import this script before importing Django

import os

def setup_openblas_environment():
    """Set OpenBLAS environment variables to prevent resource exhaustion."""
    os.environ['OPENBLAS_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['BLAS_NUM_THREADS'] = '1'
    os.environ['LAPACK_NUM_THREADS'] = '1'
    
    print("OpenBLAS environment variables set:")
    print(f"OPENBLAS_NUM_THREADS: {os.environ.get('OPENBLAS_NUM_THREADS')}")
    print(f"MKL_NUM_THREADS: {os.environ.get('MKL_NUM_THREADS')}")
    print(f"OMP_NUM_THREADS: {os.environ.get('OMP_NUM_THREADS')}")

if __name__ == "__main__":
    setup_openblas_environment()
'''
    
    script_path = Path("openblas_env.py")
    with open(script_path, 'w') as f:
        f.write(script_content)
    
    print(f"✓ Created {script_path}")

def main():
    """Main function to fix OpenBLAS issues."""
    print("🚀 RuralPoint OpenBLAS Issue Fixer")
    print("=" * 40)
    
    # Check if we're in the right directory
    if not Path("manage.py").exists():
        print("❌ Error: manage.py not found. Please run this script from the RuralPoint project root.")
        sys.exit(1)
    
    # Set environment variables
    set_environment_variables()
    
    # Test Django import
    if test_django_import():
        # Test Django commands
        test_django_commands()
        
        # Create helper scripts
        create_environment_script()
        create_python_environment_script()
        
        print("\n🎉 OpenBLAS issues fixed successfully!")
        print("\n📋 Next steps:")
        print("1. Run 'source set_openblas_env.sh' before starting Django")
        print("2. Or import 'openblas_env.py' in your Python scripts")
        print("3. Try running Django commands again")
        
    else:
        print("\n❌ Failed to fix OpenBLAS issues.")
        print("Please check the error messages above.")

if __name__ == "__main__":
    main()