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
from crewai.tools.zephyr_tools import ZephyrSearchTools


# 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

TITLE1 = """<h1 align="center">SmartMix - Your Safe Place</h1>"""
TITLE2 = """<h3 align="left">"This is an agent simulated group therapy session providing a safe, judgment-free environment, allowing for open exploration of sensitive topics. Please input the topic you would like to discuss.  Active phrases produce realistic interactions."</h3>"""
TITLE3 = """<h3 align="center">"To see active group discussion click on logs during run."</h3>"""

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_normal,
                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

with gr.Blocks() as iface:
    gr.HTML(TITLE1)
    gr.HTML(TITLE2)
    gr.HTML(TITLE3)
    with gr.Row():
        with gr.Column(scale=1):
            gr.Image(value="crewai/resources/smartmix.jpg")
        with gr.Column(scale=5):
            run_button_crewai = gr.Button(value="Run", variant="primary", scale=1)
            run_button_crewai.click(
                fn=crewai_process, 
                inputs=gr.Textbox(lines=2,label="Topic Input", placeholder="Enter Discussion Topic..."), 
                outputs=gr.Textbox(label="Group Synopsis"),
            )
                    

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