import os
import gradio as gr
from utils.auth import authenticate_google_services
from utils.llm import Google_Gemini_LLM
from agents.gmail_agent import Gmail_Agent
from agents.google_calendar_agent import GoogleCalendar_Agent

# Authenticate and initialize agents
creds = authenticate_google_services()
GMAIL_AGENT = Gmail_Agent(creds)
CALENDAR_Agent = GoogleCalendar_Agent(creds)

# Initialize LLM
# Use os.environ.get instead of userdata.get
GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
# GOOGLE_API_KEY="AIzaSyC3dFyhudUdU63sz7_-fAYOqMHx2T13BqY"
if not GOOGLE_API_KEY:
    raise ValueError("GOOGLE_API_KEY environment variable not set")

MODEL = "gemini-1.5-flash"
YOUR_AI_AGENT = Google_Gemini_LLM(GOOGLE_API_KEY, MODEL)
YOUR_AI_AGENT.append_function_tools(GMAIL_AGENT.get_all_tools())
YOUR_AI_AGENT.append_function_tools(CALENDAR_Agent.get_all_tools())

def generate_response(user_prompt, history):
    response = YOUR_AI_AGENT.generate_content(user_prompt)
    while(response.function_call):
        function_name = response.function_call.name
        function_args = response.function_call.args
        if function_name == "query_gmail_tool":
            messages = GMAIL_AGENT.read_gmail_messages(**function_args)
            if not messages:
                formatted_content = "❌ 找不到對應的 Gmail，請確認您的搜尋條件是否合理，或者請重新執行！"
                return formatted_content
            else:
                formatted_content = GMAIL_AGENT.format_messages(messages)
            user_prompt += f"\n根據使用者的搜尋條件，以下是對應的信件內容（可能有多封信件）：<信件內容開始>{formatted_content}<信件內容結束>. 請仔細閱讀以上的信件內容，並將資料整理統整給使用者！請以簡潔且新處的方式傳達內容！"
        elif function_name == "send_email_tool":
            formatted_content = GMAIL_AGENT.send_email(**function_args)
            return formatted_content
        elif function_name == "list_Calendar_events":
            formatted_content = CALENDAR_Agent.list_events(**function_args)
            user_prompt += f"\n根據使用者的搜尋條件，以下是從使用者的日曆中找到對應的事件內容（可能有多封信件）: <事件內容開始>{formatted_content}<事件內容結束>. 請仔細閱讀以上的事件內容，並將資料統整給使用者！請以簡潔且新處的方式傳達內容！"
        elif function_name == "create_Calendar_event":
            formatted_content = CALENDAR_Agent.create_event(**function_args)
            return formatted_content
        elif function_name == "update_Calendar_event":
            formatted_content = CALENDAR_Agent.update_event(**function_args)
            return formatted_content
        elif function_name == "delete_Calendar_event":
            formatted_content = CALENDAR_Agent.delete_event(**function_args)
            return formatted_content
        elif function_name == "get_event_details":
            formatted_content = CALENDAR_Agent.get_event_details(**function_args)
            user_prompt += f"\n根據使用者的搜尋條件，以下是從使用者的日曆中找到對應的事件內容 <事件內容開始>{formatted_content}<事件內容結束>. 請仔細閱讀以上的事件內容，並且完成使用者的指示（例如：統整事件內容、寄送提醒信件給參與事件的人...等等）!"
        else:
            raise NotImplementedError(f"Function Call {function_name} is not supported yet.")
        response = YOUR_AI_AGENT.generate_content(user_prompt)
    return response.text

# It's better to get the email from the user's input
# than to hardcode it.
demo = gr.ChatInterface(
    fn=generate_response,
    chatbot=gr.Chatbot(
        label="📬 AI 助理",
    ),
    title="Gmail & Calendar 智慧助理",
    description="輸入自然語言，讓 AI 幫你管理 Gmail 或 Google 行事曆",
    examples=[
        "找出2025/6月標註星號的信件",
        "請幫我列出近期的事件",
        "建立2025/8/10早上10點的網球活動",
    ]
)

if __name__ == "__main__":
    demo.launch()