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,AssistantAgent


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,
                is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
                description=description,
            )
        return general_agent
    except Exception as e:
        return e

def group_manager_agent(group_chat):
    gpp_v = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    if gpp_v == "gpt-4o":
        max_tk = 16000
    else:
        max_tk = 4000
    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,
        "cache_seed":None,
        "max_tokens":max_tk,
        "stream" : False,
        "temperature": 0.0
    }

    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 translate_agent():
    gpp_v = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    if gpp_v == "gpt-4o":
        max_tk = 16000
    else:
        max_tk = 4000
    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,
        "cache_seed":None,
        "max_tokens":max_tk,
        "stream" : False,
        "temperature": 0.0
    }
    try:
        translate_agent = ConversableAgent(
                name="translate_agent",
                system_message=""" You are a helpful AI assistant. 
                your job is to translate the input message into the target language""",
                llm_config=llm_confg,
                is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
                human_input_mode="NEVER", 
                description="assistant Agent"
            )
        return assistant_agent
    except Exception as e:
        return e

def assistant_agent():
    gpp_v = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    if gpp_v == "gpt-4o":
        max_tk = 16000
    else:
        max_tk = 4000
    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,
        "cache_seed":None,
        "max_tokens":max_tk,
        "stream" : False,
        "temperature": 0.0
    }
    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,
                is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
                human_input_mode="NEVER", 
                description="assistant Agent"
            )
        return assistant_agent
    except Exception as e:
        return e
    
def writer_assistant_agent():
    gpp_v = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    if gpp_v == "gpt-4o":
        max_tk = 16000
    else:
        max_tk = 4000
    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,
        "cache_seed":None,
        "max_tokens":max_tk,
        "stream" : False,
        "temperature": 0.0
    }
    try:
        assistant_agent = ConversableAgent(
                name="writer_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 'TERMINATE' when it is good enough,otherwise, ask eamil writer agent to generate the email again with your analysis result""",
                llm_config=llm_confg,
                is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
                human_input_mode="NEVER", 
                description="writer assistant Agent"
            )
        return assistant_agent
    except Exception as e:
        return e


def security_agent():
    gpp_v = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    if gpp_v == "gpt-4o":
        max_tk = 16000
    else:
        max_tk = 4000
    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,
        "cache_seed":None,
        "max_tokens":max_tk,
        "stream" : False,
        "temperature": 0.0
    }
    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,
            is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
            human_input_mode="NEVER", 
            description="security Agent",

        )
        return security_agent
    except Exception as e:
        return e
    
def retrieval_agent():
    gpp_v = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    if gpp_v == "gpt-4o":
        max_tk = 16000
    else:
        max_tk = 4000
    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,
        "cache_seed":None,
        "max_tokens":max_tk,
        "stream" : False,
        "temperature": 0.0
    }
    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,
                is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
                human_input_mode="NEVER", 
                description="retrieval Agent"

            )
        return RAG_agent
    except Exception as e:
        return e
    
def writer_agent():
    gpp_v = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    if gpp_v == "gpt-4o":
        max_tk = 16000
    else:
        max_tk = 4000
    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,
        "cache_seed":None,
        "max_tokens":max_tk,
        "stream" : False,
        "temperature": 0.0
    }
    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,
                is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
                human_input_mode="NEVER", 
                description="email writer Agent"
            )
        return writer_ag

    except Exception as e:
        return e
    


def writer_direct_agent():
    gpp_v = os.environ.get("AZURE_OPENAI_MODEL_NAME")
    if gpp_v == "gpt-4o":
        max_tk = 16000
    else:
        max_tk = 4000
    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,
        "cache_seed":None,
        "max_tokens":max_tk,
        "stream" : False,
        "temperature": 0.0
    }
    try:
        writer_ag = ConversableAgent(
                name="email_writer_agent",
                system_message="You are a helpful AI assistant. "
                "Your job is to generate an email for replying customer's email",
                llm_config=llm_confg,
                human_input_mode="NEVER", 
                description="email writer Agent",
                is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
            )
        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",
            description="tool_executor"
        )
        return tool_executor_proxy
    except Exception as e:
        return e
    


