import gradio as gr
from composio_llamaindex import ComposioToolSet, App, Action
from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.llms import ChatMessage
from llama_index.llms.openai import OpenAI
import os

def initialize_tools(composio_api_key, openai_api_key):
    # Initialize tools and language model
    toolset = ComposioToolSet(api_key=composio_api_key)
    tools = toolset.get_tools(apps=[App.WEBTOOL, App.EXA])
    llm = OpenAI(model="gpt-4", api_key=openai_api_key)
    
    prefix_messages = [
        ChatMessage(
            role="system",
            content=(
                "You are a Lead Outreach Agent that is equipped with great tools for research "
                "and is an expert writer. Your job is to first research some info about the lead "
                "given to you and then draft a perfect ideal email for whatever input task is given to you. "
                "Always write the subject, content of the email and nothing else."
                "Deduce the lead's name by the email you are sending it to"
            )
        )
    ]
    
    agent = FunctionCallingAgentWorker(
        tools=tools,
        llm=llm,
        prefix_messages=prefix_messages,
        max_function_calls=10,
        allow_parallel_tool_calls=False,
        verbose=True
    ).as_agent()
    
    return agent, toolset

def generate_email(lead, email_id, purpose, composio_api_key, openai_api_key):
    try:
        agent, _ = initialize_tools(composio_api_key, openai_api_key)
        response = agent.chat(
            f"These are the lead details that we know {lead} and their email is{email_id}. This is the purpose to write the email: {purpose}. "
            "Write a well written email for the purpose to the lead."
        )
        return response.response
    except Exception as e:
        return f"Error generating email: {str(e)}"

def send_email(email_content, email_id, composio_api_key):
    try:
        t = email_content + "This is the email id:" + email_id
        composio_toolset = ComposioToolSet(api_key=composio_api_key)
        composio_toolset.execute_action(
            "gmail_send_email",
            params={},
            text=str(t),
        )
        return "Email sent successfully"
    except Exception as e:
        return f"Error sending email: {str(e)}"

# Setting up Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("### Lead Outreach Email Generator")
    
    with gr.Row():
        composio_api_key_input = gr.Textbox(
            label="Composio API Key",
            placeholder="Enter your Composio API key...",
            type="password"
        )
        openai_api_key_input = gr.Textbox(
            label="OpenAI API Key",
            placeholder="Enter your OpenAI API key...",
            type="password"
        )
    
    lead_input = gr.Textbox(
        label="Lead Details",
        placeholder="Enter lead details here..."
    )
    email_input = gr.Textbox(
        label="Email ID",
        placeholder="Enter recipient's email here..."
    )
    purpose_input = gr.Textbox(
        label="Purpose and your details",
        placeholder="Enter purpose of the email, dont forget to mention your name in it..."
    )
    
    submit_button = gr.Button("Generate Email")
    send_button = gr.Button("Send Email")
    
    output_email = gr.Textbox(
        label="Generated Email",
        interactive=False
    )
    output_message = gr.Textbox(
        label="Email Status",
        placeholder="Not sent yet"
    )
    
    submit_button.click(
        generate_email,
        inputs=[
            lead_input,
            email_input,
            purpose_input,
            composio_api_key_input,
            openai_api_key_input
        ],
        outputs=output_email
    )
    
    send_button.click(
        send_email,
        inputs=[
            output_email,
            email_input,
            composio_api_key_input
        ],
        outputs=output_message
    )

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()