from crewai import Crew
from textwrap import dedent
from agents import ContentAgents
from tasks import ContentTasks
from dotenv import load_dotenv
import json
import streamlit as st
from image_search_tool import image_search_tool

load_dotenv()

class ContentCrew:
    def __init__(self, topic):
        self.topic = topic

    def run(self):
        agents = ContentAgents()
        tasks = ContentTasks()

        planner_agent = agents.planner_agent(self.topic)
        researcher_agent = agents.researcher_agent(self.topic)
        writer_agent = agents.writer_agent(self.topic)

        planner_task = tasks.planner_task(planner_agent, self.topic)
        researcher_task = tasks.researcher_task(researcher_agent, self.topic)
        writer_task = tasks.writer_task(writer_agent, self.topic)

        crew = Crew(
            agents=[planner_agent, researcher_agent, writer_agent],
            tasks=[planner_task, researcher_task, writer_task],
            verbose=True
        )

        result = crew.kickoff()

        # Debugging: Print the structure and content of the result
        print("DEBUG: Result structure and content")
        print(json.dumps(result, indent=2))

        return result

# Function to display the result in Streamlit
def display_result(result):
    st.write("## Research on Topic")
    st.write('-------------------------------')

    if 'content_plan' in result:
        st.write("### Content Plan")
        st.write(result['content_plan'])

    if 'references' in result:
        st.write("### References")
        for ref in result['references']:
            st.markdown(f"- [{ref['title']}]({ref['link']}): {ref['snippet']}")

    if 'images' in result:
        st.write("### Images")
        for idx, image in enumerate(result['images']):
            img_path = f"image_{idx}.png"
            WebSearchTools.download_image(image, img_path)
            st.image(img_path, caption=f"Image {idx+1}")

    if 'content' in result:
        st.markdown(result['content'])


# Streamlit UI
st.title('Content Creator Chatbot')
st.write('-------------------------------')
    
topic = st.text_input("What topic do you want to write about?")

if st.button('Generate Content'):
    if topic:
        # Instantiate ContentCrew and run
        content_crew = ContentCrew(topic)
        result = content_crew.run()

        # Display the results
        st.write(result)
        image_search_tool(topic)
        image_path = 'images/image_1.jpg'
        st.image(image_path, caption='Image1', use_column_width=True)
    else:
        st.error("Please enter a topic.")