import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Function to check if an API key is set
def check_api_key(key_name):
    key_value = os.getenv(key_name)
    if not key_value:
        print(f"Warning: {key_name} is not set in the environment.")
        return False
    return True

# Check API keys for various providers
openai_key_set = check_api_key("OPENAI_API_KEY")
azure_key_set = check_api_key("AZURE_OPENAI_API_KEY")
cohere_key_set = check_api_key("COHERE_API_KEY")
anthropic_key_set = check_api_key("ANTHROPIC_API_KEY")
google_key_set = check_api_key("GOOGLE_API_KEY")
huggingface_key_set = check_api_key("HUGGINGFACE_API_KEY")
togetherai_key_set = check_api_key("TOGETHERAI_API_KEY")
replicate_key_set = check_api_key("REPLICATE_API_TOKEN")
mistral_key_set = check_api_key("MISTRAL_API_KEY")
anyscale_key_set = check_api_key("ANYSCALE_API_KEY")
fireworks_key_set = check_api_key("FIREWORKS_API_KEY")
bedrock_key_set = check_api_key("AWS_ACCESS_KEY_ID") #Bedrock requires more than one key, but this is a start
ai21_key_set = check_api_key("AI21_API_KEY")
deepinfra_key_set = check_api_key("DEEPINFRA_API_KEY")
cerebrium_ai_key_set = check_api_key("CEREBRIUM_AI_API_KEY")
octoai_key_set = check_api_key("OCTOAI_API_KEY")
stability_key_set = check_api_key("STABILITY_API_KEY")

# Additional Azure checks
if azure_key_set:
    azure_base_set = check_api_key("AZURE_OPENAI_API_BASE")
    azure_version_set = check_api_key("AZURE_OPENAI_API_VERSION")
    if not (azure_base_set and azure_version_set):
        print("Warning: AZURE_OPENAI_API_BASE and AZURE_OPENAI_API_VERSION must also be set for Azure.")

#Additional OctoAI Checks
if octoai_key_set:
    octoai_endpoint_set = check_api_key("OCTOAI_ENDPOINT")
    if not octoai_endpoint_set:
        print("Warning: OCTOAI_ENDPOINT must also be set for OctoAI.")

# Additional Bedrock checks
if bedrock_key_set:
    bedrock_secret_set = check_api_key("AWS_SECRET_ACCESS_KEY")
    bedrock_region_set = check_api_key("AWS_REGION_NAME")
    if not (bedrock_secret_set and bedrock_region_set):
        print("Warning: AWS_SECRET_ACCESS_KEY and AWS_REGION_NAME must also be set for Bedrock.")


print("Environment setup check complete.")