from crewai import Agent, Task, Crew, Process
from crewai_tools import SerperDevTool
import gradio as gr
import os
import markdown as md
from datetime import datetime

# Model options
llm_models = [
    "gemini/gemini-1.5-flash",
    "gemini/gemini-1.5-pro",
    "gemini/gemini-pro"
]

selected_model = llm_models[0]


def set_model(selected_model_name):
    global selected_model
    selected_model = selected_model_name
    # return f"Model set to: {selected_model}"


def toggle_serper_input(choice):
    return gr.Textbox(visible=(choice == "Yes"))


def run_crew(gemini_api_key, search_choice, serper_api_key, topic):
    try:
        # Validate inputs
        if not gemini_api_key:
            raise ValueError("Gemini API key is required")

        os.environ['GEMINI_API_KEY'] = gemini_api_key

        # Configure search tool
        search_tool = None
        if search_choice == "Yes":
            if not serper_api_key:
                raise ValueError("Serper API key is required for online search")
            os.environ['SERPER_API_KEY'] = serper_api_key
            search_tool = SerperDevTool()

        # Create agents
        researcher = Agent(
            role="Online Research Specialist",
            goal="Aggregate comprehensive information on {topic}",
            verbose=True,
            backstory="Expert research analyst with data sourcing expertise",
            tools=[search_tool] if search_tool else [],
            llm=selected_model,
            allow_delegation=True
        )

        content_writer = Agent(
            role="Expert Content Writer",
            goal="Create SEO-optimized content on {topic}",
            verbose=True,
            backstory="Professional writer with digital journalism background",
            tools=[],
            llm=selected_model,
            allow_delegation=False
        )

        # Create tasks
        research_task = Task(
            description=f"Conduct SEO research on '{topic}'",
            expected_output="Detailed research report with SEO recommendations",
            tools=[search_tool] if search_tool else [],
            agent=researcher
        )

        writer_task = Task(
            description=f"Write SEO-optimized article on '{topic}'",
            expected_output="Polished article draft ready for publication",
            agent=content_writer,
            output_file="content.md"
        )

        # Create and run crew
        crew = Crew(
            agents=[researcher, content_writer],
            tasks=[research_task, writer_task],
            process=Process.sequential,
            verbose=True,
            max_rpm=100,
            share_crew=True,
            output_log_file=True
        )

        result = crew.kickoff(inputs={'topic': topic})

        # Return results
        with open("content.md", "r") as f:
            content = f.read()
        with open("logs.txt", 'r') as f:
            logs = f.read()
        os.remove("logs.txt")
        print(f"{datetime.now()}::{topic}-->{content}")
        return content, logs

    except Exception as e:
        return f"Error: {str(e)}", str(e)


# UI Setup
with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]),
               css='footer {visibility: hidden}') as demo:
    gr.Markdown("# AI Agents 🤖🕵🏻")

    with gr.Tabs():
        with gr.TabItem("Intro"):
            gr.Markdown(md.description)

        with gr.TabItem("SEO Content Generator Agent"):
            with gr.Accordion("How to get GEMINI API KEY ============================================", open=False):
                gr.Markdown(md.gemini_api_key)
            with gr.Accordion("How to get SERPER API KEY ============================================", open=False):
                gr.Markdown(md.serper_api_key)
            with gr.Row():
                with gr.Column(scale=1):
                    model_dropdown = gr.Dropdown(
                        llm_models,
                        label="1. Select AI Model",
                        value=llm_models[0]
                    )
                    gemini_key = gr.Textbox(
                        label="2. Enter Gemini API Key",
                        type="password",
                        placeholder="Paste your Gemini API key here..."
                    )
                    search_toggle = gr.Radio(
                        ["Yes", "No"],
                        label="3. Enable Online Search?",
                        value="No"
                    )
                    serper_key = gr.Textbox(
                        label="4. Enter Serper API Key",
                        type="password",
                        visible=False,
                        placeholder="Paste Serper API key if enabled..."
                    )
                    topic_input = gr.Textbox(
                        label="5. Enter Content Topic",
                        placeholder="Enter your article topic here..."
                    )
                    run_btn = gr.Button("Generate Content", variant="primary")

                with gr.Column(scale=3):
                    output = gr.Markdown(
                        label="Generated Content",
                        value="Your content will appear here..."
                    )
                    with gr.Accordion("Process Logs", open=False):
                        logs = gr.Markdown()

    # Event handlers
    model_dropdown.change(set_model, model_dropdown)
    search_toggle.change(
        toggle_serper_input,
        inputs=search_toggle,
        outputs=serper_key
    )
    run_btn.click(
        run_crew,
        inputs=[gemini_key, search_toggle, serper_key, topic_input],
        outputs=[output, logs],
        show_progress="full"
    )

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