import os
import gradio as gr
from textwrap import dedent
import google.generativeai as genai

# Tool import
from crewai.tools.gemini_tools import GeminiSearchTools
from langchain.tools.yahoo_finance_news import YahooFinanceNewsTool
from crewai.tools.browser_tools import BrowserTools
from crewai.tools.sec_tools import SECTools
from crewai.tools.mixtral_tools import MixtralSearchTools

# Google Langchain
from langchain_google_genai import GoogleGenerativeAI

#Crew imports
from crewai import Agent, Task, Crew, Process

# Retrieve API Key from Environment Variable
GOOGLE_AI_STUDIO = os.environ.get('GOOGLE_API_KEY')

# Ensure the API key is available
if not GOOGLE_AI_STUDIO:
    raise ValueError("API key not found. Please set the GOOGLE_AI_STUDIO2 environment variable.")

# Set gemini_llm
gemini_llm = GoogleGenerativeAI(model="gemini-pro", google_api_key=GOOGLE_AI_STUDIO)

# Base Example with Gemini Search

def crewai_process(research_topic):
    # Define your agents with roles and goals
    crazy1 = Agent(
        role='Homeless Person',
        goal=f'Discuss issue with {research_topic} and come to a deeper understanding',
        backstory="""Not doing well mentally living on the street""",
        verbose=True,
        allow_delegation=False,
        llm = gemini_llm,
        tools=[
                MixtralSearchTools.mixtral_crazy,
                GeminiSearchTools.gemini_search
                
      ]

    )
    crazy2 = Agent(
        role='Distressed Mom',
        goal='Resond to answers',
        backstory="""Struggling with depression and loss of a child""",
        verbose=True,
        allow_delegation=True,
        llm = gemini_llm,
        tools=[
                MixtralSearchTools.mixtral_crazy,
                GeminiSearchTools.gemini_search
      ]
     
        # Add tools and other optional parameters as needed
    )

    # Create tasks for your agents
    task1 = Task(
        description=f"""Discuss the  {research_topic}. Use Mixtral
        do 10 rounds of chat with. react to crazy2 """,
        agent=crazy1
    )

    task2 = Task(
        description=f"""Discuss the  {research_topic}. Use Mixtral
        do 10 rounds of chat with. react to crazy1 """,
        agent=crazy2
    )

    # Instantiate your crew with a sequential process
    crew = Crew(
        agents=[crazy1, crazy2],
        tasks=[task1, task2],
        verbose=2,
        process=Process.sequential
    )

    # Get your crew to work!
    result = crew.kickoff()
    
    return result

# Create a Gradio interface
iface = gr.Interface(
    fn=crewai_process, 
    inputs=gr.Textbox(lines=2, placeholder="Enter Research Topic Here..."), 
    outputs="text",
    title="CrewAI on Gemini (Blog Post Writer)",
    description="Input a research topic to get a comprehensive analysis (in logs) and a blog post draft (in output). To learn more connect with Mike Lively on LinkedIn at https://www.linkedin.com/in/awsmulticloud/ or join his cloud Meetup at https://www.meetup.com/florence-aws-user-group-meetup/"
)

# Launch the interface
iface.launch(debug=True)