import openai
import asyncio
import os
from openai import AzureOpenAI,AsyncAzureOpenAI
from dotenv import load_dotenv
import enum,json,io
import datetime

from autogen import ConversableAgent,GroupChatManager,GroupChat



try:
    load_dotenv()
except Exception as e:
    print(f"Error loading .env file: {e}")



def general_agent(agent_name:str,system_message:str,llm_config:dict,description:str,max_consecutive_auto_reply:int=1,human_input_mode:str="NEVER"):
    try:
        general_agent = ConversableAgent(
                name=agent_name,
                system_message=system_message,
                llm_config=llm_config,
                max_consecutive_auto_reply=max_consecutive_auto_reply,
                human_input_mode=human_input_mode, 
                description=description
            )
        return general_agent
    except Exception as e:
        return e

def group_manager_agent(group_chat):
    model = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    api_key = os.environ.get("AZURE_OPENAI_API_KEY")
    base_url = os.environ.get("AZURE_OPENAI_ENDPOINT")
    api_type = os.environ.get("OPENAI_API_TYPE")
    api_version = os.environ.get("OPENAI_API_VERSION")
    #config llm 
    llm_confg = {
        "model":model,
        "api_key":api_key,
        "base_url":base_url,
        "api_type":api_type,
        "api_version":api_version
    }

    try:
        group_manager = GroupChatManager(
        groupchat=group_chat,
        llm_config = llm_confg,
        code_execution_config=False,
        is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],)
        return group_manager
    except Exception as e:
        return e

def assistant_agent():
    model = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    api_key = os.environ.get("AZURE_OPENAI_API_KEY")
    base_url = os.environ.get("AZURE_OPENAI_ENDPOINT")
    api_type = os.environ.get("OPENAI_API_TYPE")
    api_version = os.environ.get("OPENAI_API_VERSION")
    #config llm 
    llm_confg = {
        "model":model,
        "api_key":api_key,
        "base_url":base_url,
        "api_type":api_type,
        "api_version":api_version
    }
    try:
        assistant_agent = ConversableAgent(
                name="assistant_agent",
                system_message=""" You are a helpful AI assistant. 
                You can help with checking if the email be generated by email writer agent is good enough to answer customer's request in customer email
                Return 'good to next step' when it is good enough,otherwise, ask eamil writer agent to generate the email again""",
                llm_config=llm_confg,
                max_consecutive_auto_reply=1,
                human_input_mode="NEVER", 
                description="assistant Agent"
            )
        return assistant_agent
    except Exception as e:
        return e


def security_agent():
    model = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    api_key = os.environ.get("AZURE_OPENAI_API_KEY")
    base_url = os.environ.get("AZURE_OPENAI_ENDPOINT")
    api_type = os.environ.get("OPENAI_API_TYPE")
    api_version = os.environ.get("OPENAI_API_VERSION")
    #config llm 
    llm_confg = {
        "model":model,
        "api_key":api_key,
        "base_url":base_url,
        "api_type":api_type,
        "api_version":api_version
    }
    try:
        security_agent = ConversableAgent(
            name="security_assistant",
            system_message="You are a helpful AI assistant. "
            "You can help to only suggest tool - security_checking call to get related customer information",
            llm_config=llm_confg,
            max_consecutive_auto_reply=1,
            human_input_mode="NEVER", 
            description="security Agent"

        )
        return security_agent
    except Exception as e:
        return e
    
def retrieval_agent():
    model = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    api_key = os.environ.get("AZURE_OPENAI_API_KEY")
    base_url = os.environ.get("AZURE_OPENAI_ENDPOINT")
    api_type = os.environ.get("OPENAI_API_TYPE")
    api_version = os.environ.get("OPENAI_API_VERSION")
    #config llm 
    llm_confg = {
        "model":model,
        "api_key":api_key,
        "base_url":base_url,
        "api_type":api_type,
        "api_version":api_version
    }
    try:
        RAG_agent = ConversableAgent(
                name="retrieval_assistant",
                system_message="You are a helpful AI assistant. "
                "You can help to only suggest tool - sales_data_retrieval call to get related sales transaction data. ",
                llm_config=llm_confg,
                max_consecutive_auto_reply=1,
                human_input_mode="NEVER", 
                description="retrieval Agent"

            )
        return RAG_agent
    except Exception as e:
        return e
    
def writer_agent():
    model = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    api_key = os.environ.get("AZURE_OPENAI_API_KEY")
    base_url = os.environ.get("AZURE_OPENAI_ENDPOINT")
    api_type = os.environ.get("OPENAI_API_TYPE")
    api_version = os.environ.get("OPENAI_API_VERSION")
    #config llm 
    llm_confg = {
        "model":model,
        "api_key":api_key,
        "base_url":base_url,
        "api_type":api_type,
        "api_version":api_version
    }
    try:
        writer_ag = ConversableAgent(
                name="email_writer_assistant",
                system_message="You are a helpful AI assistant. "
                "You can help to suggest tool - email_writer to generate email. MUST to use the raw input parameters as the tool parameters!",
                llm_config=llm_confg,
                max_consecutive_auto_reply=1,
                human_input_mode="NEVER", 
                description="email writer Agent"
            )
        return writer_ag

    except Exception as e:
        return e
    
    
def tool_executor_agent():
    try:
        tool_executor_proxy = ConversableAgent(
            name = "tool_executor_proxy",
            llm_config=False,
            is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
            human_input_mode="NEVER"
        )
        return tool_executor_proxy
    except Exception as e:
        return e
    


