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
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize ComposioToolSet and OpenAI LLM
toolset = ComposioToolSet(api_key=os.getenv('COMPOSIO_API_KEY')) 
tools = toolset.get_tools(apps=[App.TWITTER])

llm = OpenAI(model="gpt-4o", api_key=os.getenv('OPENAI_API_KEY'))

# Set up prefix messages for the agent
prefix_messages = [
    ChatMessage(
        role="system",
        content=(
            f"""
            You are a Twitter wrapped generator. Based on the Twitter username provided, analyze the user's profile, recent tweets, and engagement data.
            Create a personalized "Twitter Wrapped" summary highlighting their top tweets, most engaging content, follower growth, and other key insights.
            Generate the output in a structured JSON format that can be easily parsed programmatically. Include fields like "top_tweets", "engagement_stats", "follower_growth", and "summary_sheet_link".
            """
        ),
    )
]

# Initialize the agent
agent = FunctionCallingAgentWorker(
    tools=tools,
    llm=llm,
    prefix_messages=prefix_messages,
    max_function_calls=10,
    allow_parallel_tool_calls=False,
    verbose=True,
).as_agent()

def generate_wrapped(username):
    """
    Function to generate a "Twitter Wrapped" summary based on the Twitter username provided by the user.
    """
    user_input = f"Create a Twitter Wrapped summary for the username: {username}"
    response = agent.chat(user_input)
    return response

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("""### Twitter Wrapped Generator
    Enter a Twitter username below to generate your personalized Twitter Wrapped summary.
    """)

    username_input = gr.Textbox(label="Twitter Username", placeholder="e.g., @elonmusk")
    output = gr.Textbox(label="Output", placeholder="Your Twitter Wrapped summary and Google Sheet link will appear here.", lines=10)

    generate_button = gr.Button("Generate Wrapped")

    generate_button.click(fn=generate_wrapped, inputs=username_input, outputs=output)

# Launch the Gradio app
demo.launch()