predictive_analysis/diagrams/setup.py

"""
Setup Script for Architecture Diagrams
Installs dependencies and sets up the environment for diagram generation
"""

import subprocess
import sys
import os
import platform
from pathlib import Path

def run_command(command, description=""):
    """Run a shell command and handle errors"""
    try:
        if description:
            print(f"🔄 {description}...")
        
        result = subprocess.run(command, shell=True, capture_output=True, text=True)
        
        if result.returncode == 0:
            if description:
                print(f"✅ {description} completed successfully")
            return True
        else:
            print(f"❌ Error: {result.stderr}")
            return False
            
    except Exception as e:
        print(f"❌ Exception: {e}")
        return False

def check_python_version():
    """Check if Python version is compatible"""
    version = sys.version_info
    if version.major >= 3 and version.minor >= 8:
        print(f"✅ Python {version.major}.{version.minor}.{version.micro} is compatible")
        return True
    else:
        print(f"❌ Python {version.major}.{version.minor}.{version.micro} is not compatible (requires Python 3.8+)")
        return False

def check_graphviz():
    """Check if Graphviz is installed"""
    print("🔍 Checking Graphviz installation...")
    
    if run_command("dot -V", ""):
        print("✅ Graphviz is already installed")
        return True
    else:
        print("❌ Graphviz not found")
        return False

def install_graphviz():
    """Install Graphviz based on the operating system"""
    system = platform.system().lower()
    
    print(f"🔧 Installing Graphviz for {system}...")
    
    if system == "windows":
        print("📋 Windows detected. Please install Graphviz manually:")
        print("   1. Download from: https://graphviz.org/download/")
        print("   2. Or use Chocolatey: choco install graphviz")
        print("   3. Or use Scoop: scoop install graphviz")
        print("   4. Add Graphviz bin directory to PATH")
        return False
        
    elif system == "darwin":  # macOS
        if run_command("which brew", ""):
            return run_command("brew install graphviz", "Installing Graphviz via Homebrew")
        else:
            print("❌ Homebrew not found. Please install Graphviz manually:")
            print("   1. Install Homebrew: https://brew.sh/")
            print("   2. Then run: brew install graphviz")
            return False
            
    elif system == "linux":
        # Try different package managers
        if run_command("which apt-get", ""):
            return run_command("sudo apt-get update && sudo apt-get install -y graphviz", 
                             "Installing Graphviz via apt-get")
        elif run_command("which yum", ""):
            return run_command("sudo yum install -y graphviz", 
                             "Installing Graphviz via yum")
        elif run_command("which dnf", ""):
            return run_command("sudo dnf install -y graphviz", 
                             "Installing Graphviz via dnf")
        elif run_command("which pacman", ""):
            return run_command("sudo pacman -S graphviz", 
                             "Installing Graphviz via pacman")
        else:
            print("❌ No supported package manager found. Please install Graphviz manually.")
            return False
    else:
        print(f"❌ Unsupported operating system: {system}")
        return False

def install_python_dependencies():
    """Install Python dependencies for diagrams"""
    print("📦 Installing Python dependencies...")
    
    # Create diagrams directory if it doesn't exist
    Path("diagrams").mkdir(exist_ok=True)
    
    # Install from requirements file
    if os.path.exists("diagrams/requirements.txt"):
        return run_command(f"{sys.executable} -m pip install -r diagrams/requirements.txt",
                          "Installing Python packages")
    else:
        # Fallback to direct installation
        packages = ["diagrams>=0.23.4", "graphviz>=0.20.1"]
        for package in packages:
            if not run_command(f"{sys.executable} -m pip install {package}",
                              f"Installing {package}"):
                return False
        return True

def verify_installation():
    """Verify that everything is installed correctly"""
    print("🔍 Verifying installation...")
    
    # Check Python imports
    try:
        import diagrams
        print(f"✅ Diagrams library version: {diagrams.__version__}")
    except ImportError:
        print("❌ Failed to import diagrams library")
        return False
    
    # Check Graphviz again
    if not check_graphviz():
        return False
    
    # Test diagram generation
    try:
        print("🧪 Testing diagram generation...")
        
        test_code = '''
from diagrams import Diagram, Node
from diagrams.generic.compute import Rack

with Diagram("Test Diagram", filename="diagrams/test_diagram", show=False):
    test_node = Rack("Test Node")
'''
        
        # Write test file
        with open("diagrams/test_setup.py", "w") as f:
            f.write(test_code)
        
        # Run test
        result = subprocess.run([sys.executable, "diagrams/test_setup.py"], 
                              capture_output=True, text=True)
        
        if result.returncode == 0 and os.path.exists("diagrams/test_diagram.png"):
            print("✅ Test diagram generated successfully")
            # Clean up test files
            os.remove("diagrams/test_setup.py")
            os.remove("diagrams/test_diagram.png")
            return True
        else:
            print(f"❌ Test diagram generation failed: {result.stderr}")
            return False
            
    except Exception as e:
        print(f"❌ Test generation exception: {e}")
        return False

def main():
    """Main setup function"""
    print("🚀 Architecture Diagrams Setup")
    print("=" * 40)
    
    # Check Python version
    if not check_python_version():
        sys.exit(1)
    
    # Check and install Graphviz
    if not check_graphviz():
        print("📥 Graphviz installation required...")
        if not install_graphviz():
            print("\n⚠️  Please install Graphviz manually and run this script again.")
            sys.exit(1)
    
    # Install Python dependencies
    if not install_python_dependencies():
        print("❌ Failed to install Python dependencies")
        sys.exit(1)
    
    # Verify installation
    if not verify_installation():
        print("❌ Installation verification failed")
        sys.exit(1)
    
    print("\n" + "=" * 40)
    print("🎉 Setup completed successfully!")
    print("\nNext steps:")
    print("  1. Generate all diagrams: python diagrams/generate_all_diagrams.py")
    print("  2. Generate specific diagram: cd diagrams && python 01_data_flow_architecture.py")
    print("  3. View generated PNG files in the diagrams/ directory")
    print("\n📚 See diagrams/README.md for detailed usage instructions")

if __name__ == "__main__":
    main()