import os
import streamlit as st
from crewai import Agent, Task, Crew, Process
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
from crewai_tools import SerperDevTool

# Load environment variables from .env file
load_dotenv()

# Set up Streamlit page configuration
st.set_page_config(page_title="SEO Crew Builder", layout="wide")

# Streamlit UI
st.title("SEO Crew Builder")

# Get the API keys from environment variables
openai_api_key = os.getenv("OPENAI_API_KEY")
serper_api_key = os.getenv("SERPER_API_KEY")

if openai_api_key and serper_api_key:
    # Create a ChatOpenAI instance with GPT-4
    gpt4_model = ChatOpenAI(model_name="gpt-4")

    # Create a SerperDevTool instance
    search_tool = SerperDevTool()

    # Predefined SEO agent roles with default goals, backstories, and tasks
    seo_roles = {
        "SEO Strategist": {
            "goal": "Develop comprehensive SEO strategies to improve website visibility and rankings",
            "backstory": "An experienced SEO professional with a track record of boosting organic traffic for various industries.",
            "task": "Conduct a comprehensive analysis of our website's current SEO performance and develop a strategy to improve our rankings for key target keywords.",
            "expected_output": "A detailed SEO strategy report with actionable recommendations."
        },
        "Content Optimizer": {
            "goal": "Analyze and optimize website content to align with SEO best practices and target keywords",
            "backstory": "A skilled writer and editor with deep knowledge of SEO content creation and optimization techniques.",
            "task": "Review our top 5 underperforming blog posts and provide specific recommendations for optimizing them to improve search engine rankings.",
            "expected_output": "A report with optimization recommendations for 5 blog posts."
        },
        "Technical SEO Specialist": {
            "goal": "Identify and resolve technical SEO issues to improve website performance and crawlability",
            "backstory": "A tech-savvy SEO expert with a strong background in web development and search engine algorithms.",
            "task": "Perform a technical SEO audit of our website, identifying any issues that may be hindering our search performance, and provide a prioritized list of fixes.",
            "expected_output": "A comprehensive technical SEO audit report with a prioritized list of fixes."
        }
    }

    # Function to create an agent
    def create_agent(role, goal, backstory):
        return Agent(
            role=role,
            goal=goal,
            backstory=backstory,
            verbose=True,
            allow_delegation=False,
            tools=[search_tool],
            llm=gpt4_model
        )

    # Function to create a task
    def create_task(description, expected_output, agent):
        return Task(
            description=description,
            expected_output=expected_output,
            agent=agent
        )

    # Sidebar for agent creation
    st.sidebar.header("Create SEO Agents")
    agents = []
    tasks = []
    for i, (default_role, details) in enumerate(seo_roles.items()):
        with st.sidebar.expander(f"Agent {i+1}"):
            role = st.text_input(f"Role for Agent {i+1}", value=default_role, key=f"role_{i}")
            goal = st.text_input(f"Goal for Agent {i+1}", value=details['goal'], key=f"goal_{i}")
            backstory = st.text_area(f"Backstory for Agent {i+1}", value=details['backstory'], key=f"backstory_{i}")
            
            if role and goal and backstory:
                agent = create_agent(role, goal, backstory)
                agents.append(agent)
                
                # Create task for this agent
                task_description = st.text_area(f"Task for {role}", value=details['task'], key=f"task_{i}")
                expected_output = st.text_input(f"Expected Output for {role}", value=details['expected_output'], key=f"output_{i}")
                if task_description and expected_output:
                    tasks.append(create_task(task_description, expected_output, agent))

    # Process selection
    process = st.radio("Select Process", ["Sequential", "Hierarchical"])
    process_enum = Process.sequential if process == "Sequential" else Process.hierarchical

    # Button to start the analysis
    if st.button("Start SEO Analysis"):
        if agents and tasks:
            with st.spinner("SEO Crew is working on the tasks..."):
                crew = Crew(
                    agents=agents,
                    tasks=tasks,
                    verbose=True,
                    process=process_enum
                )
                result = crew.kickoff()
            
            st.success("SEO Analysis complete!")
            
            # Display the result
            st.markdown("## SEO Analysis Result")
            st.markdown(str(result))
            
            # Add a download button for the full report
            st.download_button(
                label="Download Full SEO Report",
                data=str(result),
                file_name="seo_analysis_report.txt",
                mime="text/plain"
            )
        else:
            st.error("Please create at least one agent and one task before starting the analysis.")
else:
    st.error("API keys not found in .env file. Please add OPENAI_API_KEY and SERPER_API_KEY and restart the app.")

# Instructions for using the app
st.sidebar.header("How to Use")
st.sidebar.info(
    "1. Customize SEO agents in the sidebar by modifying their roles, goals, and backstories.\n"
    "2. Modify the tasks for each agent if needed.\n"
    "3. Select the process type (Sequential or Hierarchical).\n"
    "4. Click 'Start SEO Analysis' to run the crew.\n"
    "5. View the results and download the report if desired."
)