"""ADK agent with vector search RAG tool."""

from functools import partial

from google import genai
from google.adk.agents.llm_agent import Agent
from google.adk.runners import Runner
from google.adk.tools.mcp_tool import McpToolset
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
from google.cloud.firestore_v1.async_client import AsyncClient
from google.genai.types import Content, Part

from va_agent.auth import auth_headers_provider
from va_agent.config import settings
from va_agent.dynamic_instruction import provide_dynamic_instruction
from va_agent.governance import GovernancePlugin
from va_agent.notifications import FirestoreNotificationBackend
from va_agent.session import FirestoreSessionService

# MCP Toolset for RAG knowledge search
toolset = McpToolset(
    connection_params=StreamableHTTPConnectionParams(url=settings.mcp_remote_url),
    header_provider=auth_headers_provider,
)


# Shared Firestore client for session service and notifications
firestore_db = AsyncClient(database=settings.firestore_db)

# Session service with compaction
session_service = FirestoreSessionService(
    db=firestore_db,
    compaction_token_threshold=10_000,
    genai_client=genai.Client(),
)

# Notification service
notification_service = FirestoreNotificationBackend(
    db=firestore_db,
    collection_path=settings.notifications_collection_path,
    max_to_notify=settings.notifications_max_to_notify,
)

# Agent with static and dynamic instructions
governance = GovernancePlugin()
agent = Agent(
    model=settings.agent_model,
    name=settings.agent_name,
    instruction=partial(provide_dynamic_instruction, notification_service),
    static_instruction=Content(
        role="user",
        parts=[Part(text=settings.agent_instructions)],
    ),
    tools=[toolset],
    after_model_callback=governance.after_model_callback,
)

# Runner
runner = Runner(
    app_name="va_agent",
    agent=agent,
    session_service=session_service,
    auto_create_session=True,
)