import os
import gradio as gr
from autogen import AssistantAgent
from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent
from dotenv import load_dotenv

load_dotenv()  # Load environment variables

gpt4 = {
    "config_list": [{"model": "gpt-4", "api_key": os.environ["OPENAI_API_KEY"]}],
}

# Initialize agents
assistant = AssistantAgent(
    name="assistant",
    system_message="You are a helpful assistant. Give the user information with a set of logical steps to follow. In addition provide helpful url resources from ragproxyagent.",
    llm_config=gpt4,
)

ragproxyagent = RetrieveUserProxyAgent(
    name="ragproxyagent",
    human_input_mode="ALWAYS",
    retrieve_config={
        "task": "qa",
        "docs_path": "./scraped_results/",
        "context_max_tokens": 5000,
        "overwrite": False,
        "get_or_create": True,
        "return_source": True,
    },
    code_execution_config=False,
)

# Function to reset agents
def _reset_agents():
    ragproxyagent.reset()
    assistant.reset()

import time


def chat_with_agents(user_input, history):
    _reset_agents()  # Reset agents before each new conversation

    # Step 1: RAG Agent retrieves context (doc-based information)
    chat_result = ragproxyagent.initiate_chat(
        assistant,
        message=ragproxyagent.message_generator,
        problem=user_input,
        n_results=3,
        max_turns=1,
    )

    print("Raw RAG Response:", chat_result)  # Debugging

    # Step 2: Extract assistant's final summarized response
    response = "Sorry, I couldn't generate a response."

    # Check if the response contains the summary
    if hasattr(chat_result, 'summary') and chat_result.summary:
        response = chat_result.summary

   # print("Assistant's Final Response:", response)  # Debugging

    # Step 3: Append user and assistant messages
   # history.append({"role": "user", "content": user_input})
   # history.append({"role": "assistant", "content": response})

    # Return only the last 2 messages (user + assistant) for Gradio
    return response

# Use "messages" type for proper chat UI
chat_interface = gr.ChatInterface(
    fn=chat_with_agents,
    type="messages"
)

if __name__ == "__main__":
    chat_interface.launch()