# Install the required packages
%pip install -qU llama-index llama-index-tools-mcp klavis

import os
from klavis import Klavis
from klavis.types import McpServerName, ConnectionType
from llama_index.llms.openai import OpenAI
from llama_index.tools.mcp import BasicMCPClient

# Set environment variables
os.environ[
    "OPENAI_API_KEY"
] = "YOUR_OPENAI_API_KEY"  # Replace with your actual OpenAI API key
os.environ[
    "KLAVIS_API_KEY"
] = "YOUR_KLAVIS_API_KEY"  # Replace with your actual Klavis API key

klavis_client = Klavis(api_key=os.getenv("KLAVIS_API_KEY"))

# Create a YouTube MCP server and get the server URL
youtube_mcp_instance = klavis_client.mcp_server.create_server_instance(
    server_name=McpServerName.YOUTUBE,
    user_id="1234",
    platform_name="Klavis",
    connection_type=ConnectionType.STREAMABLE_HTTP,
)

youtube_mcp_server_url = youtube_mcp_instance.server_url
# print(f"🔗 YouTube MCP server created at: {youtube_mcp_server_url}")

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.tools.mcp import (
    get_tools_from_mcp_url,
    aget_tools_from_mcp_url,
)

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

youtube_tools = await aget_tools_from_mcp_url(
    youtube_mcp_server_url, client=BasicMCPClient(youtube_mcp_server_url)
)

youtube_agent = FunctionAgent(
    name="youtube_agent",
    description="Agent using MCP-based tools",
    tools=youtube_tools,
    llm=llm,
    system_prompt="You are an AI assistant that uses MCP tools.",
)

YOUTUBE_VIDEO_URL = "https://www.youtube.com/watch?v=MmiveeGxfX0&t=528s"  # pick a video you like!

response = await youtube_agent.run(
    f"Summarize this video: {YOUTUBE_VIDEO_URL}"
)
print(response)

import webbrowser

klavis_client = Klavis(api_key=os.getenv("KLAVIS_API_KEY"))

# Create YouTube MCP server
youtube_mcp_instance = klavis_client.mcp_server.create_server_instance(
    server_name=McpServerName.YOUTUBE,
    user_id="1234",
    platform_name="Klavis",
    connection_type=ConnectionType.STREAMABLE_HTTP,
)

# Create Gmail MCP server with OAuth authorization
gmail_mcp_instance = klavis_client.mcp_server.create_server_instance(
    server_name=McpServerName.GMAIL,
    user_id="1234",
    platform_name="Klavis",
    connection_type=ConnectionType.STREAMABLE_HTTP,
)

print("✅ Created YouTube and Gmail MCP instances")

# Open Gmail OAuth authorization
webbrowser.open(gmail_mcp_instance.oauth_url)
print(
    f"🔐 Opening OAuth authorization for Gmail, if you are not redirected, please open the following URL in your browser: {gmail_mcp_instance.oauth_url}"
)

from llama_index.llms.openai import OpenAI
from llama_index.core.agent.workflow import FunctionAgent, AgentWorkflow
from llama_index.tools.mcp import (
    BasicMCPClient,
    get_tools_from_mcp_url,
    aget_tools_from_mcp_url,
)

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

# Get MCP server URLs
youtube_mcp_server_url = youtube_mcp_instance.server_url
gmail_mcp_server_url = gmail_mcp_instance.server_url

# Get tools from both MCP servers
youtube_tools = await aget_tools_from_mcp_url(
    youtube_mcp_server_url, client=BasicMCPClient(youtube_mcp_server_url)
)
gmail_tools = await aget_tools_from_mcp_url(
    gmail_mcp_server_url, client=BasicMCPClient(gmail_mcp_server_url)
)

# Create specialized agents
youtube_agent = FunctionAgent(
    name="youtube_agent",
    description="Agent that can summarize YouTube videos",
    tools=youtube_tools,
    llm=llm,
    system_prompt="You are a YouTube video summarization expert. Use MCP tools to analyze and summarize videos.",
    can_handoff_to=["gmail_agent"],
)

gmail_agent = FunctionAgent(
    name="gmail_agent",
    description="Agent that can send emails via Gmail",
    tools=gmail_tools,
    llm=llm,
    system_prompt="You are an email assistant. Use MCP tools to send emails via Gmail.",
)

# Create multi-agent workflow
workflow = AgentWorkflow(
    agents=[youtube_agent, gmail_agent],
    root_agent="youtube_agent",
)

print("🤖 Multi-agent workflow created with YouTube and Gmail agents!")

YOUTUBE_VIDEO_URL = "https://www.youtube.com/watch?v=MmiveeGxfX0&t=528s"  # pick a video you like!
EMAIL_RECIPIENT = "[email protected]"  # Replace with your email

resp = await workflow.run(
    user_msg=f"Summarize this video {YOUTUBE_VIDEO_URL} and send it to {EMAIL_RECIPIENT}"
)
print("\n✅ Report:\n", resp.response.content)