import os
from dotenv import load_dotenv, find_dotenv
from simple_salesforce import Salesforce

def get_credentials(context):
    """
    Get credentials for a given context from the .env file
    
    Args:
        context (str): Context name (e.g., 'qa2', 'prod')
    
    Returns:
        dict: Credentials dictionary with username, password, and security_token
    """
    context = context.upper()
            
    # Initialize credentials dictionary
    credentials = {
        'USERNAME': None,
        'PASSWORD': None,
        'SECURITY_TOKEN': None
    }
    
    if context != 'PROD':
        credentials['DOMAIN'] = 'test'

    # Load the .env file
    env_file = find_dotenv(".env")
    load_dotenv(env_file)
    
    # Load all environment variables
    env_vars = os.environ
    
    for key, value in env_vars.items():
        if f'{context}_SF_' in key:
            credential_key = key.split(f'{context}_SF_')[-1].upper()
            credentials[credential_key] = value
            
    return credentials

def get_sf_connection(context):
    """
    Create Salesforce connection based on context
    
    Args:
        context (str): Context name (e.g., 'qa2', 'prod')
    
    Returns:
        Salesforce: Authenticated Salesforce connection
    """
    credentials = get_credentials(context)
    
    if not all(credentials.values()):
        raise ValueError(f"Missing credentials for context: {context}")
    
    if context.lower() == 'prod':
        return Salesforce(
            username=credentials['USERNAME'],
            password=credentials['PASSWORD'],
            security_token=credentials['SECURITY_TOKEN'],
            version='62.0'
        )
    else:
        return Salesforce(
            username=credentials['USERNAME'],
            password=credentials['PASSWORD'],
            security_token=credentials['SECURITY_TOKEN'],
            domain=credentials['DOMAIN'],
            version='62.0'
        )