from dataclasses import dataclass
from enum import Enum
from typing import Optional, Dict, Any
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 llama_index.llms.gemini import Gemini
from dotenv import load_dotenv
import gradio as gr
import os
import json
from datetime import datetime

# Load environment variables
load_dotenv()

class ConnectionStatus(Enum):
    PENDING = "pending"
    ACTIVE = "active"
    FAILED = "failed"
    NOT_FOUND = "not_found"

@dataclass
class APIResponse:
    success: bool
    data: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    
    def to_json(self) -> str:
        return json.dumps({
            "success": self.success,
            "data": self.data,
            "error": self.error
        })

class CalendarService:
    def __init__(self):
        self.toolset = ComposioToolSet(api_key=os.getenv('COMPOSIO_API_KEY'))
        self.llm = Gemini(model="models/gemini-2.0-flash", api_key=os.getenv('GOOGLE_API_KEY'))
        self.connections: Dict[str, Dict[str, Any]] = {}
        
    def initiate_connection(self, entity_id: str, redirect_url: Optional[str] = None) -> APIResponse:
        try:
            if not redirect_url:
                redirect_url = "https://yourwebsite.com/connection/success"
                
            connection_request = self.toolset.initiate_connection(
                entity_id=entity_id,
                app=App.GOOGLECALENDAR
            )
            
            self.connections[entity_id] = {
                'status': ConnectionStatus.PENDING.value,
                'redirect_url': connection_request.redirectUrl,
                'created_at': datetime.now().isoformat()
            }
            
            return APIResponse(
                success=True,
                data={
                    'status': ConnectionStatus.PENDING.value,
                    'redirect_url': connection_request.redirectUrl,
                    'wait_time': 60,
                    'message': "Please authenticate using the provided link."
                }
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                error=f"Failed to initiate connection: {str(e)}"
            )
    
    def check_status(self, entity_id: str) -> APIResponse:
        try:
            if entity_id not in self.connections:
                return APIResponse(
                    success=False,
                    error="No connection found for this entity ID"
                )
            
            # In a real implementation, you would check the actual status
            # For now, we'll simulate the status check
            connection = self.connections[entity_id]
            
            return APIResponse(
                success=True,
                data={
                    'status': connection['status'],
                    'message': f"Connection status: {connection['status']}"
                }
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                error=f"Failed to check status: {str(e)}"
            )
    
    def generate_wrapped(self, entity_id: str) -> APIResponse:
        try:
            tools = self.toolset.get_tools([
                Action.GOOGLECALENDAR_GET_CALENDAR,
                Action.GOOGLECALENDAR_FIND_EVENT,
                Action.GOOGLECALENDAR_GET_CURRENT_DATE_TIME,
            ])
            
            agent = FunctionCallingAgentWorker(
                tools=tools,
                llm=self.llm,
                prefix_messages=[
                    ChatMessage(
                        role="system",
                        content="""
                        Use the tools available to you get factual data then
                        Generate a creative Google Calendar Wrapped summary including:
                        1. Total meetings this year
                        2. Overall time spent in meetings
                        3. Busiest month
                        4. Busiest day
                        5. Most frequent meeting participant
                        6. Average meeting duration
                        7. Most common meeting time
                        
                        Be entertaining and witty with the analysis. Conclude by assigning 
                        a Greek god persona based on their meeting patterns.
                        """
                    )
                ],
                max_function_calls=10,
                allow_parallel_tool_calls=False,
                verbose=True
            ).as_agent()
            
            response = agent.chat(
                f"Create a Calendar Wrapped summary for user with entity_id: {entity_id}. "
                "Use the available Google Calendar actions to gather real data about the user's calendar usage."
                "calendar id is primary"
            )
            
            return APIResponse(
                success=True,
                data={
                    'wrapped_summary': str(response)
                }
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                error=f"Failed to generate wrapped: {str(e)}"
            )

def create_gradio_api():
    service = CalendarService()
    
    def handle_connection(entity_id: str, redirect_url: Optional[str] = None) -> str:
        response = service.initiate_connection(entity_id, redirect_url)
        return response.to_json()
    
    def check_status(entity_id: str) -> str:
        response = service.check_status(entity_id)
        return response.to_json()
    
    def generate_wrapped(entity_id: str) -> str:
        response = service.generate_wrapped(entity_id)
        return response.to_json()
    
    # Create API endpoints
    connection_api = gr.Interface(
        fn=handle_connection,
        inputs=[
            gr.Textbox(label="Entity ID"),
            gr.Textbox(label="Redirect URL", placeholder="https://yourwebsite.com/connection/success")
        ],
        outputs=gr.JSON(),
        title="Initialize Calendar Connection",
        description="Start a new calendar connection for an entity",
        examples=[["user123", "https://example.com/callback"]]
    )
    
    status_api = gr.Interface(
        fn=check_status,
        inputs=gr.Textbox(label="Entity ID"),
        outputs=gr.JSON(),
        title="Check Connection Status",
        description="Check the status of an existing connection",
        examples=[["user123"]]
    )
    
    wrapped_api = gr.Interface(
        fn=generate_wrapped,
        inputs=gr.Textbox(label="Entity ID"),
        outputs=gr.JSON(),
        title="Generate Calendar Wrapped",
        description="Generate a calendar wrapped summary for an entity",
        examples=[["user123"]]
    )
    
    # Combine all interfaces
    api = gr.TabbedInterface(
        [connection_api, status_api, wrapped_api],
        ["Connect", "Check Status", "Generate Wrapped"],
        title="Calendar Wrapped API",
    )
    
    return api

if __name__ == "__main__":
    api = create_gradio_api()
    api.launch(server_name="0.0.0.0", server_port=7860)