# app/agent.py
from llama_index.core import Settings
from llama_index.llms.gemini import Gemini
from llama_index.core.agent import FunctionCallingAgentWorker
from composio_llamaindex import ComposioToolSet, App
from llama_index.core.llms import ChatMessage

class GmailAgent:
    def __init__(self, google_api_key: str, composio_api_key: str):
        # Set up Gemini LLM
        Settings.llm = Gemini(
            model="models/gemini-1.5-pro",
            api_key=google_api_key
        )

        # Initialize Composio ToolSet
        self.composio_toolset = ComposioToolSet(api_key=composio_api_key)
        self.gmail_tools = self.composio_toolset.get_tools(apps=[App.GMAIL])

        # Set up agent
        prefix_messages = [
            ChatMessage(
                role="system",
                content="""You are a Gmail Assistant that can:
                - Read and analyze email content
                - Draft professional replies
                - Manage email threads
                Use the available Gmail tools to perform these actions."""
            )
        ]

        agent_worker = FunctionCallingAgentWorker.from_tools(
            tools=self.gmail_tools,
            prefix_messages=prefix_messages,
            verbose=True
        )
        self.agent = agent_worker.as_agent()

    def authenticate_gmail(self):
        request = self.composio_toolset.initiate_connection(app=App.GMAIL)
        return request.redirectUrl

    async def process_email(self, thread_id: str, email_content: str, sender_email: str):
        prompt = f"""Analyze this email and create an appropriate reply:
        Thread ID: {thread_id}
        Sender: {sender_email}
        Content: {email_content}
        
        Please draft a professional response and use the appropriate Gmail tool to send it."""
        
        response = await self.agent.achat(prompt)
        return str(response)