# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Sample agent demonstrating Agent Registry discovery."""

import os

from google.adk.agents.llm_agent import LlmAgent
from google.adk.integrations.agent_registry import AgentRegistry
from google.adk.models.google_llm import Gemini

# Project and location can be set via environment variables:
# GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION
project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
location = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")

# Initialize Agent Registry client
registry = AgentRegistry(project_id=project_id, location=location)

# List agents, MCP servers, and endpoints resource names from the registry.
# They can be used to initialize the agent, toolset, and model below.
print(f"Listing agents in {project_id}/{location}...")
agents = registry.list_agents()
for agent in agents.get("agents", []):
  print(f"- Agent: {agent.get('displayName')} ({agent.get('name')})")

print(f"\nListing MCP servers in {project_id}/{location}...")
mcp_servers = registry.list_mcp_servers()
for server in mcp_servers.get("mcpServers", []):
  print(f"- MCP Server: {server.get('displayName')} ({server.get('name')})")

print(f"\nListing endpoints in {project_id}/{location}...")
endpoints = registry.list_endpoints()
for endpoint in endpoints.get("endpoints", []):
  print(f"- Endpoint: {endpoint.get('displayName')} ({endpoint.get('name')})")

# Search agents and MCP servers matching a query
print(f"\nSearching agents matching 'Workspace' in {project_id}/{location}...")
matching_agents = registry.search_agents(search_string="Workspace")
for agent in matching_agents.get("agents", []):
  print(f"- Found Agent: {agent.get('displayName')} ({agent.get('name')})")

print(
    "\nSearching MCP servers matching 'agentregistry' in"
    f" {project_id}/{location}..."
)
matching_servers = registry.search_mcp_servers(search_string="agentregistry")
for server in matching_servers.get("mcpServers", []):
  print(
      f"- Found MCP Server: {server.get('displayName')} ({server.get('name')})"
  )

# Example of using a specific agent or MCP server from the registry:
# (Note: These names should be full resource names as returned by list methods)

# 1. Using a Remote A2A Agent as a sub-agent
# TODO: Replace AGENT_NAME with your agent name
remote_agent = registry.get_remote_a2a_agent(
    f"projects/{project_id}/locations/{location}/agents/AGENT_NAME"
)

# 2. Using an MCP Server in a toolset
# TODO: Replace MCP_SERVER_NAME with your MCP server name
mcp_toolset = registry.get_mcp_toolset(
    f"projects/{project_id}/locations/{location}/mcpServers/MCP_SERVER_NAME"
)

# 3. Getting a specific model endpoint configuration
# This returns a string like:
# "projects/adk12345/locations/us-central1/publishers/google/models/gemini-2.5-flash"
# TODO: Replace ENDPOINT_NAME with your endpoint name
model_name = registry.get_model_name(
    f"projects/{project_id}/locations/{location}/endpoints/ENDPOINT_NAME"
)

# Initialize the model using the resolved model name from registry.
gemini_model = Gemini(model=model_name)

root_agent = LlmAgent(
    model=gemini_model,
    name="discovery_agent",
    instruction=(
        "You have access to tools and sub-agents discovered via Registry."
    ),
    tools=[mcp_toolset],
    sub_agents=[remote_agent],
)