import yaml
import os
import sys
import platform

def check_python():
    print(f"Python Version: {platform.python_version()}")
    print("Python executable:", sys.executable)  # Print the path to the Python executable

def load_environment_yaml(path_to_yaml):
    """Load the list of dependencies from the environment.yml file."""
    try:
        with open(path_to_yaml, 'r') as file:
            env_yaml = yaml.safe_load(file)
    except FileNotFoundError:
        print(f"Error: File not found {path_to_yaml}")
        sys.exit(1)  # Exit the script if the file is not found
    dependencies = env_yaml.get('dependencies', [])
    return [dep.split('=')[0] for dep in dependencies if isinstance(dep, str)]

def check_imports(dependencies):
    """Check if each dependency in the list can be imported and print their versions."""
    for dependency in dependencies:
        if dependency == "python":  # "python" is not a module
            continue
        try:
            module = __import__(dependency)
            version = getattr(module, '__version__', 'version not specified')
            print(f"{dependency} is available, version: {version}")
        except ImportError:
            print(f"{dependency} is not available")


def check_environment_variables():
    print("Environment Variables:")
    relevant_vars = ['SLURM_JOB_ID', 'SLURM_JOB_NODELIST', 'SLURM_NPROCS', 'SLURM_SUBMIT_DIR']
    for var in relevant_vars:
        value = os.getenv(var)
        if value is not None:
            print(f"{var}: {value}")
        else:
            print(f"{var} is not set.")



if __name__ == "__main__":
    print("Starting Environment Test...")
    print("PYTHONPATH:", os.getenv("PYTHONPATH"))  # Print the PYTHONPATH environment variable
    # Use the hard-coded absolute path
    
    # Discover the current path 
    current_path = os.getcwd()
    print("Current path:", current_path)
    
    # Build path to required env 
    dependencies_dir = os.path.dirname(current_path)
    dependencies_dir = os.path.join(dependencies_dir, "env")
    env_yaml_path = os.path.join(dependencies_dir, "environment.yml")
    
    dependencies = load_environment_yaml(env_yaml_path)
    check_python()
    check_imports(dependencies)
    check_environment_variables()
