import streamlit as st
from crewai import Crew
from langchain_groq import ChatGroq
import os
from textwrap import dedent
from crewai import Agent
from crewai_tools import ScrapeWebsiteTool, SerperDevTool
from crewai import Task

# Initialize the LLM
llm = ChatGroq(
    api_key="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",
    verbose=True,
    model="llama3-8b-8192"
)

# Set environment variable for API key
os.environ["SERPER_API_KEY"] = "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"

# Define the agents and tasks
def web_researcher_agent(topic):
    return Agent(
        role="Expert web researcher",
        goal=dedent(f"""Your goal is to search on the web and scrape websites for relevant high-quality content about the given topic.
                     Use the tools provided to search on web and scrape website.
                     Topic: {topic}"""),
        tools=[ScrapeWebsiteTool(), SerperDevTool()],
        backstory=dedent("""You are proficient at searching for specific topics on the web, selecting those that provide
                            best value and information."""),
        verbose=True,
        allow_delegation=False,
        llm=llm
    )

def writer_agent(topic):
    return Agent(
        role="Expert Writer",
        goal=dedent(f"""Your goal is to write high-quality content about the given topic.
                     Topic: {topic}"""),
        tools=[ScrapeWebsiteTool(), SerperDevTool()],
        backstory=dedent("""You are proficient at writing high-quality content about the given topic."""),
        verbose=True,
        allow_delegation=False,
        llm=llm
    )

def web_researcher_task(topic):
    return Task(
        description=dedent(f"""Get valuable and high-quality information from the web about a given topic.
                              Your goal is to extract high-quality information from the web about a given topic.
                              Topic: {topic}"""),
        expected_output=dedent(f"""The expected output should be well-formatted information about the given topic,
                                  only high-quality information related to the topic.
                                  Topic: {topic}"""),
        agent=web_researcher_agent(topic)
    )

def writer_task(topic):
    return Task(
        description=dedent(f"""Using information extracted by the web researcher agent,
                              your task is to write very detailed and lengthy content on the given topic.
                              Topic: {topic}"""),
        expected_output=dedent(f"""The expected output should be well-formatted and flawless content on the given topic.
                                  Topic: {topic}"""),
        agent=writer_agent(topic)
    )

# Streamlit UI
st.header("Multi-Agent Chat System")

# Initialize chat history
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

# Function to format messages with HTML
def format_message(role, message):
    if role == 'User':
        return f"<div style='padding:10px;'><b style='color: #007bff;'>User:</b> {message}</div>"
    elif role == 'System':
        return f"<div style='padding:10px;'><b style='color: #dc3545;'>System:</b> {message}</div>"

# Display chat history
for entry in st.session_state.chat_history:
    st.markdown(format_message(entry['role'], entry['message']), unsafe_allow_html=True)

# Input and buttons
topic = st.text_input("Enter the topic:")

col1, col2 = st.columns([2, 1])

with col1:
    if st.button("Generate") and topic:
        # Add user input to chat history
        st.session_state.chat_history.append({'role': 'User', 'message': topic})

        # Create and run Crew with agents and tasks
        crew = Crew(
            agents=[web_researcher_agent(topic), writer_agent(topic)],
            tasks=[web_researcher_task(topic), writer_task(topic)],
            verbose=True,
        )
        result = crew.kickoff()

        # Add system response to chat history
        st.session_state.chat_history.append({'role': 'System', 'message': result})

        # Display updated chat history
        for entry in st.session_state.chat_history:
            st.markdown(format_message(entry['role'], entry['message']), unsafe_allow_html=True)

with col2:
    if st.button("Clear Chat"):
       st.session_state.chat_history = []
       st.session_state['dummy_var'] = not st.          session_state.get('dummy_var', False)