import os
from typing import Dict, Any
from dotenv import load_dotenv
from openai import OpenAI, AzureOpenAI

load_dotenv()

# Initialize clients only if their required environment variables are present
engine_map = {}

# Import here to avoid circular imports - will be populated in code below
from agent_config import BaseAgentConfig, SITAASAgentConfig

# Azure OpenAI client setup
if all(os.getenv(var) for var in ["AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_API_VERSION"]):
    azure_client = AzureOpenAI(
        api_key=os.getenv("AZURE_OPENAI_API_KEY"),
        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
        api_version=os.getenv("AZURE_OPENAI_API_VERSION")
    )
    engine_map.update({
        "azure-gpt-4o-mini": {
            "client": azure_client, 
            "model": "gpt-4o-mini",   
            "config": BaseAgentConfig("azure-gpt-4o-mini")
        },
        "azure-gpt-4o": {
            "client": azure_client, 
            "model": "gpt-4o",   
            "config": BaseAgentConfig("azure-gpt-4o")
        },
    })

# OpenAI client setup
if os.getenv("OPENAI_API_KEY"):
    openai_client = OpenAI(
        api_key=os.getenv("OPENAI_API_KEY"),
    )
    engine_map.update({
        "openai-gpt-4o-mini": {
            "client": openai_client, 
            "model": "gpt-4o-mini",   
            "config": BaseAgentConfig("openai-gpt-4o-mini")
        },
        "openai-gpt-4o": {
            "client": openai_client, 
            "model": "gpt-4o",   
            "config": BaseAgentConfig("openai-gpt-4o")
        },
    })

# x.ai client setup
if os.getenv("XAI_API_KEY"):
    xai_client = OpenAI(
        api_key=os.getenv("XAI_API_KEY"),
        base_url="https://api.x.ai/v1",
    )
    engine_map["grok-2-latest"] = {
        "client": xai_client, 
        "model": "grok-2-latest",   
        "config": BaseAgentConfig("grok-2-latest")
    }

# Add the SITAAS agent to the engine_map if required env vars are available
if all(os.getenv(var) for var in ["SEARCH_ENDPOINT", "SEARCH_KEY"]):
    engine_map["sitaas-agent"] = {
        "client": azure_client,
        "model": "gpt-4.1-mini",
        "config": SITAASAgentConfig("gpt-4.1-mini")  # Use the SITAAS agent configuration
    }

def get_client(engine: str) -> OpenAI:
    """Get the appropriate client for the given engine"""
    try:
        return engine_map[engine]["client"]
    except KeyError:
        raise ValueError(f"Unsupported engine in get_client: {engine}")
    
def get_model(engine: str) -> str:
    """Get the model name for the given engine"""
    try:
        return engine_map[engine]["model"]
    except KeyError:
        raise ValueError(f"Unsupported engine in get_model: {engine}")

def get_config(engine: str) -> Any:
    """Get the agent configuration for the given engine"""
    try:
        return engine_map[engine]["config"]
    except KeyError:
        raise ValueError(f"Unsupported engine in get_config: {engine}")