"""
Agno runtime implementation.

This is the main runtime class that provides Agno framework integration
for the Agent Control Plane.
"""

import asyncio
import queue
import threading
import time
import structlog
from typing import Dict, Any, Optional, AsyncIterator, Callable, TYPE_CHECKING

from ..base import (
    RuntimeType,
    RuntimeExecutionResult,
    RuntimeExecutionContext,
    RuntimeCapabilities,
    BaseRuntime,
    RuntimeRegistry,
)
from control_plane_api.worker.services.event_publisher import (
    EventPublisher,
    EventPublisherConfig,
    EventPriority,
)
from control_plane_api.worker.utils.tool_validation import (
    validate_and_sanitize_tools,
    sanitize_tool_name,
)

from .config import build_agno_agent_config
from .hooks import create_tool_hook_for_streaming, create_tool_hook_with_callback
from .mcp_builder import build_agno_mcp_tools
from .utils import (
    build_conversation_messages,
    extract_usage,
    extract_tool_messages,
    extract_response_content,
)

if TYPE_CHECKING:
    from control_plane_client import ControlPlaneClient
    from services.cancellation_manager import CancellationManager

logger = structlog.get_logger(__name__)


@RuntimeRegistry.register(RuntimeType.DEFAULT)
class AgnoRuntime(BaseRuntime):
    """
    Runtime implementation using Agno framework.

    This runtime wraps the Agno-based agent execution logic,
    providing a clean interface that conforms to the AgentRuntime protocol.

    Features:
    - LiteLLM-based model execution
    - Real-time streaming with event batching
    - Tool execution hooks
    - Conversation history support
    - Comprehensive usage tracking
    """

    def __init__(
        self,
        control_plane_client: "ControlPlaneClient",
        cancellation_manager: "CancellationManager",
        **kwargs,
    ):
        """
        Initialize the Agno runtime.

        Args:
            control_plane_client: Client for Control Plane API
            cancellation_manager: Manager for execution cancellation
            **kwargs: Additional configuration options
        """
        super().__init__(control_plane_client, cancellation_manager, **kwargs)
        self._custom_tools: Dict[str, Any] = {}  # tool_id -> tool instance

    def get_runtime_type(self) -> RuntimeType:
        """Return RuntimeType.DEFAULT."""
        return RuntimeType.DEFAULT

    def get_capabilities(self) -> RuntimeCapabilities:
        """Return Agno runtime capabilities."""
        return RuntimeCapabilities(
            streaming=True,
            tools=True,
            mcp=True,  # Agno supports MCP via MCPTools
            hooks=True,
            cancellation=True,
            conversation_history=True,
            custom_tools=True  # Agno supports custom Python tools
        )

    def _validate_mcp_tool_names(self, mcp_tool: Any, execution_id: str) -> Any:
        """
        Validate and sanitize MCP tool function names.

        This ensures MCP tools from external servers meet universal LLM provider requirements.
        Agno's MCPTools.functions property contains the actual tool definitions.

        Args:
            mcp_tool: Connected MCPTools instance
            execution_id: Execution ID for logging

        Returns:
            The MCP tool instance with validated function names
        """
        if not hasattr(mcp_tool, 'functions') or not mcp_tool.functions:
            return mcp_tool

        server_name = getattr(mcp_tool, '_server_name', mcp_tool.name)

        # Convert functions dict to list of Function objects for validation
        # mcp_tool.functions is a Dict[str, Function], we need to validate the Function objects
        functions_list = list(mcp_tool.functions.values())

        # Validate and sanitize the function names
        validated_functions, validation_report = validate_and_sanitize_tools(
            functions_list,
            tool_name_getter=lambda f: getattr(f, 'name', str(f)),
            auto_fix=True,
            provider_context=f"agno_mcp_{server_name}"
        )

        # Log validation results
        sanitized_count = sum(1 for r in validation_report if r['action'] == 'sanitized')
        filtered_count = sum(1 for r in validation_report if r['action'] == 'filtered')

        if sanitized_count > 0:
            self.logger.warning(
                "mcp_tool_names_sanitized",
                server_name=server_name,
                execution_id=execution_id,
                sanitized_count=sanitized_count,
                total_functions=len(mcp_tool.functions),
                details=[r for r in validation_report if r['action'] == 'sanitized'][:5]  # Limit details
            )

        if filtered_count > 0:
            self.logger.error(
                "mcp_tool_names_filtered",
                server_name=server_name,
                execution_id=execution_id,
                filtered_count=filtered_count,
                total_functions=len(mcp_tool.functions),
                details=[r for r in validation_report if r['action'] == 'filtered']
            )

        # Reconstruct the functions dict from validated Function objects
        # Preserve the Dict[str, Function] structure that Agno expects
        from collections import OrderedDict
        validated_functions_dict = OrderedDict()
        for func in validated_functions:
            func_name = getattr(func, 'name', str(func))
            validated_functions_dict[func_name] = func

        # Update the MCP tool with validated functions dict
        mcp_tool.functions = validated_functions_dict

        return mcp_tool

    async def _execute_impl(
        self, context: RuntimeExecutionContext
    ) -> RuntimeExecutionResult:
        """
        Execute agent using Agno framework without streaming.

        Args:
            context: Execution context with prompt, history, config

        Returns:
            RuntimeExecutionResult with response and metadata
        """
        mcp_tools_instances = []

        try:
            # Build MCP tools from context
            mcp_tools_instances = build_agno_mcp_tools(context.mcp_servers)

            # Connect MCP tools
            connected_mcp_tools = []
            for mcp_tool in mcp_tools_instances:
                try:
                    await mcp_tool.connect()

                    # Verify the tool is actually initialized (agno doesn't raise on failure)
                    if not mcp_tool.initialized:
                        server_name = getattr(mcp_tool, '_server_name', mcp_tool.name)
                        error_msg = f"Failed to initialize MCP tool: {server_name}"
                        self.logger.error(
                            "mcp_tool_initialization_failed",
                            server_name=server_name,
                            execution_id=context.execution_id,
                            error=error_msg,
                        )
                        raise RuntimeError(error_msg)

                    # Verify it has tools available
                    if not mcp_tool.functions:
                        server_name = getattr(mcp_tool, '_server_name', mcp_tool.name)
                        error_msg = f"MCP tool {server_name} has no functions available"
                        self.logger.error(
                            "mcp_tool_has_no_functions",
                            server_name=server_name,
                            execution_id=context.execution_id,
                            error=error_msg,
                        )
                        raise RuntimeError(error_msg)

                    self.logger.info(
                        "mcp_tool_connected",
                        execution_id=context.execution_id,
                        server_name=getattr(mcp_tool, '_server_name', mcp_tool.name),
                        function_count=len(mcp_tool.functions),
                    )

                    # UNIVERSAL VALIDATION: Validate MCP tool names
                    validated_mcp_tool = self._validate_mcp_tool_names(mcp_tool, context.execution_id)
                    connected_mcp_tools.append(validated_mcp_tool)

                except Exception as e:
                    server_name = getattr(mcp_tool, '_server_name', mcp_tool.name)
                    self.logger.error(
                        "mcp_tool_connection_failed",
                        error=str(e),
                        error_type=type(e).__name__,
                        server_name=server_name,
                        execution_id=context.execution_id,
                    )
                    import traceback
                    self.logger.debug(
                        "mcp_tool_connection_error_traceback",
                        traceback=traceback.format_exc(),
                        execution_id=context.execution_id,
                    )

                    # Publish MCP connection error event
                    try:
                        from control_plane_api.worker.utils.error_publisher import (
                            ErrorEventPublisher, ErrorSeverity, ErrorCategory
                        )
                        error_publisher = ErrorEventPublisher(self.control_plane)
                        await error_publisher.publish_error(
                            execution_id=context.execution_id,
                            exception=e,
                            severity=ErrorSeverity.WARNING,
                            category=ErrorCategory.MCP_CONNECTION,
                            stage="initialization",
                            component="mcp_server",
                            operation=f"connect_{server_name}",
                            metadata={"server_name": server_name},
                            recovery_actions=[
                                "Verify MCP server is running and accessible",
                                "Check MCP server configuration",
                                "Review network connectivity",
                            ],
                        )
                    except Exception as publish_error:
                        # Log warning but don't block execution
                        self.logger.warning(
                            f"Failed to publish MCP connection error: {publish_error}",
                            execution_id=context.execution_id,
                        )

                    # Continue with other MCP tools even if one fails

            # Use only successfully connected tools
            mcp_tools_instances = connected_mcp_tools

            # Merge regular skills with custom tools
            # IMPORTANT: Deep copy skills to isolate Function objects between executions
            # This prevents schema corruption from shared mutable state in Function.parameters
            from copy import deepcopy

            all_skills = []
            if context.skills:
                for skill in context.skills:
                    try:
                        # Deep copy the skill to ensure Function objects are isolated
                        # This prevents process_entrypoint() from modifying shared state
                        if hasattr(skill, 'functions') and hasattr(skill.functions, 'items'):
                            copied_skill = deepcopy(skill)
                            all_skills.append(copied_skill)
                            self.logger.debug(
                                "skill_deep_copied",
                                skill_name=getattr(skill, 'name', 'unknown'),
                                function_count=len(skill.functions) if hasattr(skill, 'functions') else 0,
                                execution_id=context.execution_id,
                            )
                        else:
                            # For non-Toolkit skills, use as-is
                            all_skills.append(skill)
                    except Exception as e:
                        # If deep copy fails, fall back to original skill and log warning
                        self.logger.warning(
                            "skill_deep_copy_failed_using_original",
                            skill_name=getattr(skill, 'name', 'unknown'),
                            error=str(e),
                            error_type=type(e).__name__,
                            execution_id=context.execution_id,
                        )
                        all_skills.append(skill)

            # Add custom tools
            if self._custom_tools:
                for tool_id, custom_tool in self._custom_tools.items():
                    try:
                        # Get toolkit from custom tool
                        toolkit = custom_tool.get_tools()

                        # Extract tools - handle both Toolkit objects and iterables
                        if hasattr(toolkit, 'tools'):
                            all_skills.extend(toolkit.tools)
                        elif hasattr(toolkit, '__iter__'):
                            all_skills.extend(toolkit)
                        else:
                            all_skills.append(toolkit)

                        self.logger.debug(
                            "custom_tool_loaded",
                            tool_id=tool_id,
                            execution_id=context.execution_id
                        )
                    except Exception as e:
                        self.logger.error(
                            "custom_tool_load_failed",
                            tool_id=tool_id,
                            error=str(e),
                            execution_id=context.execution_id
                        )

            # Extract metadata for Langfuse tracking
            user_id = None
            session_id = None
            agent_name = None

            if context.user_metadata:
                user_id = context.user_metadata.get("user_email") or context.user_metadata.get("user_id")
                session_id = context.user_metadata.get("session_id") or context.execution_id
                agent_name = context.user_metadata.get("agent_name") or context.agent_id

            # DEBUG: Log metadata extraction
            self.logger.warning(
                "🔍 DEBUG: AGNO RUNTIME (_execute_impl) - METADATA EXTRACTION",
                context_user_metadata=context.user_metadata,
                extracted_user_id=user_id,
                extracted_session_id=session_id,
                extracted_agent_name=agent_name,
                organization_id=context.organization_id,
                # Change 9 (P1): Log which field was chosen for user_id
                user_id_source="user_email" if (context.user_metadata and context.user_metadata.get("user_email")) else "user_id",
                raw_user_email=context.user_metadata.get("user_email") if context.user_metadata else None,
                raw_user_id=context.user_metadata.get("user_id") if context.user_metadata else None,
            )

            # Create Agno agent with all tools (skills + MCP tools) and metadata
            agent = build_agno_agent_config(
                agent_id=context.agent_id,
                system_prompt=context.system_prompt,
                model_id=context.model_id,
                skills=all_skills,
                mcp_tools=mcp_tools_instances,
                tool_hooks=None,
                user_id=user_id,
                session_id=session_id,
                organization_id=context.organization_id,
                agent_name=agent_name,
                skill_configs=context.skill_configs,
                user_metadata=context.user_metadata,
                task_execution_id=context.execution_id,
            )

            # Register for cancellation
            self.cancellation_manager.register(
                execution_id=context.execution_id,
                instance=agent,
                instance_type="agent",
            )

            # Log tool schema snapshots for debugging
            # This helps detect schema inconsistencies and parameter mismatches
            if context.execution_id:
                for skill in all_skills:
                    if hasattr(skill, 'functions') and hasattr(skill.functions, 'items'):
                        skill_name = getattr(skill, 'name', 'unknown')
                        for func_name, func_obj in skill.functions.items():
                            if hasattr(func_obj, 'parameters') and isinstance(func_obj.parameters, dict):
                                param_names = list(func_obj.parameters.get('properties', {}).keys())
                                self.logger.debug(
                                    "tool_schema_snapshot",
                                    execution_id=context.execution_id,
                                    skill_name=skill_name,
                                    tool_name=func_name,
                                    parameters=param_names,
                                )

            # Build conversation context
            messages = build_conversation_messages(context.conversation_history)

            # Determine if we need async execution (when MCP tools are present)
            has_async_tools = len(mcp_tools_instances) > 0

            # Execute without streaming
            if has_async_tools:
                # Use async agent.arun() for MCP tools
                if messages:
                    result = await agent.arun(context.prompt, stream=False, messages=messages)
                else:
                    result = await agent.arun(context.prompt, stream=False)
            else:
                # Use sync agent.run() for non-MCP tools
                def run_agent():
                    if messages:
                        return agent.run(context.prompt, stream=False, messages=messages)
                    else:
                        return agent.run(context.prompt, stream=False)

                # Run in thread pool to avoid blocking
                result = await asyncio.to_thread(run_agent)

            # Cleanup
            self.cancellation_manager.unregister(context.execution_id)

            # Extract response and metadata
            response_content = extract_response_content(result)
            usage = extract_usage(result)
            tool_messages = extract_tool_messages(result)

            return RuntimeExecutionResult(
                response=response_content,
                usage=usage,
                success=True,
                finish_reason="stop",
                run_id=getattr(result, "run_id", None),
                model=context.model_id,
                tool_messages=tool_messages,
            )

        except asyncio.CancelledError:
            # Handle cancellation
            self.cancellation_manager.cancel(context.execution_id)
            self.cancellation_manager.unregister(context.execution_id)
            raise

        except Exception as e:
            self.logger.error(
                "agno_execution_failed",
                execution_id=context.execution_id,
                error=str(e),
            )
            self.cancellation_manager.unregister(context.execution_id)

            # Publish error event
            try:
                from control_plane_api.worker.utils.error_publisher import (
                    ErrorEventPublisher, ErrorSeverity, ErrorCategory
                )
                error_publisher = ErrorEventPublisher(self.control_plane)
                await error_publisher.publish_error(
                    execution_id=context.execution_id,
                    exception=e,
                    severity=ErrorSeverity.CRITICAL,
                    category=ErrorCategory.UNKNOWN,
                    stage="execution",
                    component="agno_runtime",
                    operation="agent_execution",
                    include_stack_trace=True,
                )
            except Exception:
                pass  # Never break execution flow

            return RuntimeExecutionResult(
                response="",
                usage={},
                success=False,
                error=str(e),
            )

        finally:
            # Close MCP tool connections
            for mcp_tool in mcp_tools_instances:
                try:
                    await mcp_tool.close()
                    self.logger.debug(
                        "mcp_tool_closed",
                        execution_id=context.execution_id,
                    )
                except Exception as e:
                    self.logger.error(
                        "mcp_tool_close_failed",
                        error=str(e),
                        execution_id=context.execution_id,
                    )

    async def _stream_execute_impl(
        self,
        context: RuntimeExecutionContext,
        event_callback: Optional[Callable[[Dict], None]] = None,
    ) -> AsyncIterator[RuntimeExecutionResult]:
        """
        Execute agent with streaming using Agno framework with efficient event batching.

        This implementation uses the EventPublisher service to batch message chunks,
        reducing HTTP requests by 90-96% while keeping tool events immediate.

        Args:
            context: Execution context
            event_callback: Optional callback for real-time events

        Yields:
            RuntimeExecutionResult chunks as they arrive in real-time
        """
        # Create event publisher with batching
        event_publisher = EventPublisher(
            control_plane=self.control_plane,
            execution_id=context.execution_id,
            config=EventPublisherConfig.from_env(),
        )

        mcp_tools_instances = []

        try:
            # Build MCP tools from context
            mcp_tools_instances = build_agno_mcp_tools(context.mcp_servers)

            # Connect MCP tools
            connected_mcp_tools = []
            for mcp_tool in mcp_tools_instances:
                try:
                    await mcp_tool.connect()

                    # Verify the tool is actually initialized (agno doesn't raise on failure)
                    if not mcp_tool.initialized:
                        server_name = getattr(mcp_tool, '_server_name', mcp_tool.name)
                        error_msg = f"Failed to initialize MCP tool: {server_name}"
                        self.logger.error(
                            "mcp_tool_initialization_failed",
                            server_name=server_name,
                            execution_id=context.execution_id,
                            error=error_msg,
                        )
                        raise RuntimeError(error_msg)

                    # Verify it has tools available
                    if not mcp_tool.functions:
                        server_name = getattr(mcp_tool, '_server_name', mcp_tool.name)
                        error_msg = f"MCP tool {server_name} has no functions available"
                        self.logger.error(
                            "mcp_tool_has_no_functions",
                            server_name=server_name,
                            execution_id=context.execution_id,
                            error=error_msg,
                        )
                        raise RuntimeError(error_msg)

                    self.logger.info(
                        "mcp_tool_connected_streaming",
                        execution_id=context.execution_id,
                        server_name=getattr(mcp_tool, '_server_name', mcp_tool.name),
                        function_count=len(mcp_tool.functions),
                    )

                    # UNIVERSAL VALIDATION: Validate MCP tool names
                    validated_mcp_tool = self._validate_mcp_tool_names(mcp_tool, context.execution_id)
                    connected_mcp_tools.append(validated_mcp_tool)

                except Exception as e:
                    self.logger.error(
                        "mcp_tool_connection_failed_streaming",
                        error=str(e),
                        error_type=type(e).__name__,
                        server_name=getattr(mcp_tool, '_server_name', mcp_tool.name),
                        execution_id=context.execution_id,
                    )
                    import traceback
                    self.logger.debug(
                        "mcp_tool_connection_error_traceback",
                        traceback=traceback.format_exc(),
                        execution_id=context.execution_id,
                    )
                    # Continue with other MCP tools even if one fails

            # Use only successfully connected tools
            mcp_tools_instances = connected_mcp_tools

            # Build conversation context
            messages = build_conversation_messages(context.conversation_history)

            # Determine if we need async execution (when MCP tools are present)
            has_async_tools = len(mcp_tools_instances) > 0

            # Stream execution - publish events INSIDE the thread (like old code)
            accumulated_response = ""
            run_result = None

            # Create queue for streaming chunks from thread to async
            chunk_queue = queue.Queue()

            # Generate unique message ID
            message_id = f"{context.execution_id}_msg_{int(time.time() * 1000000)}"

            # Merge regular skills with custom tools
            all_skills = list(context.skills) if context.skills else []

            # Add custom tools
            if self._custom_tools:
                for tool_id, custom_tool in self._custom_tools.items():
                    try:
                        # Get toolkit from custom tool
                        toolkit = custom_tool.get_tools()

                        # Extract tools - handle both Toolkit objects and iterables
                        if hasattr(toolkit, 'tools'):
                            all_skills.extend(toolkit.tools)
                        elif hasattr(toolkit, '__iter__'):
                            all_skills.extend(toolkit)
                        else:
                            all_skills.append(toolkit)

                        self.logger.debug(
                            "custom_tool_loaded_streaming",
                            tool_id=tool_id,
                            execution_id=context.execution_id
                        )
                    except Exception as e:
                        self.logger.error(
                            "custom_tool_load_failed_streaming",
                            tool_id=tool_id,
                            error=str(e),
                            execution_id=context.execution_id
                        )

            # Initialize enforcement service
            enforcement_context = {
                "organization_id": context.organization_id,
                "user_email": context.user_email,
                "user_id": context.user_id,
                "user_roles": context.user_roles or [],
                "team_id": context.team_id,
                "team_name": context.team_name,
                "agent_id": context.agent_id,
                "environment": context.environment,
                "model_id": context.model_id,
            }

            # Import enforcement dependencies
            from control_plane_api.app.lib.policy_enforcer_client import create_policy_enforcer_client
            from control_plane_api.worker.services.tool_enforcement import ToolEnforcementService
            import os

            # Get enforcer client (using the same token as the control plane)
            enforcer_client = None
            enforcement_service = None

            # Check if enforcement is enabled (opt-in via environment variable)
            enforcement_enabled = os.environ.get("KUBIYA_ENFORCE_ENABLED", "").lower() in ("true", "1", "yes")

            if not enforcement_enabled:
                self.logger.info(
                    "policy_enforcement_disabled",
                    reason="KUBIYA_ENFORCE_ENABLED not set",
                    execution_id=context.execution_id[:8],
                    note="Set KUBIYA_ENFORCE_ENABLED=true to enable policy enforcement"
                )
            else:
                try:
                    # Get enforcer URL - default to control plane enforcer proxy
                    enforcer_url = os.environ.get("ENFORCER_SERVICE_URL")
                    if not enforcer_url:
                        # Use control plane's enforcer proxy as default
                        control_plane_url = os.environ.get("CONTROL_PLANE_URL", "http://localhost:8000")
                        enforcer_url = f"{control_plane_url.rstrip('/')}/api/v1/enforcer"
                        self.logger.debug(
                            "using_control_plane_enforcer_proxy",
                            enforcer_url=enforcer_url,
                            execution_id=context.execution_id[:8],
                        )

                    # Use async context manager properly (we're in an async function)
                    enforcer_client_context = create_policy_enforcer_client(
                        enforcer_url=enforcer_url,
                        api_key=self.control_plane.api_key,
                        auth_type="UserKey"
                    )
                    enforcer_client = await enforcer_client_context.__aenter__()
                    if enforcer_client:
                        enforcement_service = ToolEnforcementService(enforcer_client)
                        self.logger.info(
                            "policy_enforcement_enabled",
                            enforcer_url=enforcer_url,
                            execution_id=context.execution_id[:8],
                        )
                except Exception as e:
                    self.logger.warning(
                        "enforcement_service_init_failed",
                        error=str(e),
                        execution_id=context.execution_id[:8],
                    )

            # Create tool hook that publishes directly to Control Plane with enforcement
            tool_hook = create_tool_hook_for_streaming(
                control_plane=self.control_plane,
                execution_id=context.execution_id,
                message_id=message_id,  # Link tools to this assistant message turn
                enforcement_context=enforcement_context,
                enforcement_service=enforcement_service,
            )

            # Extract metadata for Langfuse tracking
            user_id = None
            session_id = None
            agent_name = None

            if context.user_metadata:
                user_id = context.user_metadata.get("user_email") or context.user_metadata.get("user_id")
                session_id = context.user_metadata.get("session_id") or context.execution_id
                agent_name = context.user_metadata.get("agent_name") or context.agent_id

            # DEBUG: Log metadata extraction
            self.logger.warning(
                "🔍 DEBUG: AGNO RUNTIME (_stream_execute_impl) - METADATA EXTRACTION",
                context_user_metadata=context.user_metadata,
                extracted_user_id=user_id,
                extracted_session_id=session_id,
                extracted_agent_name=agent_name,
                organization_id=context.organization_id,
                # Change 9 (P1): Log which field was chosen for user_id
                user_id_source="user_email" if (context.user_metadata and context.user_metadata.get("user_email")) else "user_id",
                raw_user_email=context.user_metadata.get("user_email") if context.user_metadata else None,
                raw_user_id=context.user_metadata.get("user_id") if context.user_metadata else None,
            )

            # Create Agno agent with all tools (skills + MCP tools), tool hooks, and metadata
            agent = build_agno_agent_config(
                agent_id=context.agent_id,
                system_prompt=context.system_prompt,
                model_id=context.model_id,
                skills=all_skills,
                mcp_tools=mcp_tools_instances,
                tool_hooks=[tool_hook],
                user_id=user_id,
                session_id=session_id,
                organization_id=context.organization_id,
                agent_name=agent_name,
                skill_configs=context.skill_configs,
                user_metadata=context.user_metadata,
                task_execution_id=context.execution_id,
            )

            # Register for cancellation
            self.cancellation_manager.register(
                execution_id=context.execution_id,
                instance=agent,
                instance_type="agent",
            )

            # Cache execution metadata
            self.control_plane.cache_metadata(context.execution_id, "AGENT")

            def stream_agent_run():
                """
                Run agent with streaming and publish events directly to Control Plane.
                This runs in a thread pool, so blocking HTTP calls are OK here.
                Put chunks in queue for async iterator to yield in real-time.
                """
                nonlocal accumulated_response, run_result
                run_id_published = False

                # Use thread-local event loop from control_plane client
                # This ensures all async operations (tool hooks, event publishing) share the same loop
                # and it persists until explicitly cleaned up at the end of execution
                thread_loop = self.control_plane._get_thread_event_loop()

                try:
                    # Use async streaming for MCP tools, sync for others
                    if has_async_tools:
                        # For async tools (MCP), we need to use agent.arun() in an async context
                        # Use the thread-local event loop instead of creating a new one
                        if messages:
                            stream_response = thread_loop.run_until_complete(
                                agent.arun(context.prompt, stream=True, messages=messages)
                            )
                        else:
                            stream_response = thread_loop.run_until_complete(
                                agent.arun(context.prompt, stream=True)
                            )
                    else:
                        # Use sync agent.run() for non-MCP tools
                        if messages:
                            stream_response = agent.run(
                                context.prompt,
                                stream=True,
                                messages=messages,
                            )
                        else:
                            stream_response = agent.run(context.prompt, stream=True)

                    # Iterate over streaming chunks and publish IMMEDIATELY
                    for chunk in stream_response:
                        # Capture run_id for cancellation (first chunk)
                        if not run_id_published and hasattr(chunk, "run_id") and chunk.run_id:
                            self.cancellation_manager.set_run_id(
                                context.execution_id, chunk.run_id
                            )

                            # Publish run_id event
                            self.control_plane.publish_event(
                                execution_id=context.execution_id,
                                event_type="run_started",
                                data={
                                    "run_id": chunk.run_id,
                                    "execution_id": context.execution_id,
                                    "cancellable": True,
                                }
                            )
                            run_id_published = True

                        # Extract content
                        chunk_content = ""
                        if hasattr(chunk, "content") and chunk.content:
                            if isinstance(chunk.content, str):
                                chunk_content = chunk.content
                            elif hasattr(chunk.content, "text"):
                                chunk_content = chunk.content.text

                        # Filter out whitespace-only chunks to prevent "(no content)" in UI
                        if chunk_content and chunk_content.strip():
                            accumulated_response += chunk_content

                            # Queue chunk for batched publishing (via EventPublisher in async context)
                            # This reduces 300 HTTP requests → 12 requests (96% reduction)
                            chunk_queue.put(("chunk", chunk_content, message_id))

                    # Store final result
                    run_result = stream_response

                    # Signal completion
                    chunk_queue.put(("done", run_result))

                except Exception as e:
                    self.logger.error("streaming_error", error=str(e))
                    chunk_queue.put(("error", e))
                    raise

            # Start streaming in background thread
            stream_thread = threading.Thread(target=stream_agent_run, daemon=True)
            stream_thread.start()

            # Yield chunks as they arrive in the queue and publish via EventPublisher
            while True:
                try:
                    # Non-blocking get with short timeout for responsiveness
                    queue_item = await asyncio.to_thread(chunk_queue.get, timeout=0.1)

                    if queue_item[0] == "chunk":
                        # Unpack chunk data
                        _, chunk_content, msg_id = queue_item

                        # Publish chunk via EventPublisher (batched, non-blocking)
                        await event_publisher.publish(
                            event_type="message_chunk",
                            data={
                                "role": "assistant",
                                "content": chunk_content,
                                "is_chunk": True,
                                "message_id": msg_id,
                            },
                            priority=EventPriority.NORMAL,  # Batched
                        )

                        # Yield chunk immediately to iterator
                        yield RuntimeExecutionResult(
                            response=chunk_content,
                            usage={},
                            success=True,
                        )
                    elif queue_item[0] == "done":
                        # Final result - extract metadata and break
                        run_result = queue_item[1]
                        break
                    elif queue_item[0] == "error":
                        # Error occurred in thread
                        raise queue_item[1]

                except queue.Empty:
                    # Queue empty, check if thread is still alive
                    if not stream_thread.is_alive():
                        # Thread died without putting "done" - something went wrong
                        break
                    # Thread still running, continue waiting
                    continue

            # Wait for thread to complete
            await asyncio.to_thread(stream_thread.join, timeout=5.0)

            # Yield final result with complete metadata
            usage = extract_usage(run_result) if run_result else {}
            tool_messages = extract_tool_messages(run_result) if run_result else []

            yield RuntimeExecutionResult(
                response=accumulated_response,  # Full accumulated response
                usage=usage,
                success=True,
                finish_reason="stop",
                run_id=getattr(run_result, "run_id", None) if run_result else None,
                model=context.model_id,
                tool_messages=tool_messages,
                metadata={"accumulated_response": accumulated_response},
            )

        finally:
            # Flush and close event publisher to ensure all batched events are sent
            await event_publisher.flush()
            await event_publisher.close()

            # Clean up thread-local event loop used by tool hooks
            # This prevents resource leaks and "await wasn't used with future" errors
            self.control_plane.close_thread_event_loop()

            # Close MCP tool connections
            for mcp_tool in mcp_tools_instances:
                try:
                    await mcp_tool.close()
                    self.logger.debug(
                        "mcp_tool_closed_streaming",
                        execution_id=context.execution_id,
                    )
                except Exception as e:
                    self.logger.error(
                        "mcp_tool_close_failed_streaming",
                        error=str(e),
                        execution_id=context.execution_id,
                    )

            # Close enforcer client context manager (fix resource leak)
            if 'enforcer_client_context' in locals() and enforcer_client_context is not None:
                try:
                    await enforcer_client_context.__aexit__(None, None, None)
                    self.logger.debug(
                        "enforcer_client_closed",
                        execution_id=context.execution_id[:8],
                    )
                except Exception as e:
                    self.logger.warning(
                        "enforcer_client_close_failed",
                        error=str(e),
                        execution_id=context.execution_id[:8],
                    )

            # Cleanup
            self.cancellation_manager.unregister(context.execution_id)

    # ==================== Custom Tool Extension API ====================

    def get_custom_tool_requirements(self) -> Dict[str, Any]:
        """
        Get requirements for creating custom tools for Agno runtime.

        Returns:
            Dictionary with format, examples, and documentation for Agno custom tools
        """
        return {
            "format": "python_class",
            "description": "Python class with get_tools() method returning Agno Toolkit",
            "example_code": '''
from agno.tools import Toolkit

class MyCustomTool:
    """Custom tool for Agno runtime."""

    def get_tools(self) -> Toolkit:
        """Return Agno toolkit with custom functions."""
        return Toolkit(
            name="my_tool",
            tools=[self.my_function]
        )

    def my_function(self, arg: str) -> str:
        """Tool function description."""
        return f"Result: {arg}"
            ''',
            "documentation_url": "https://docs.agno.ai/custom-tools",
            "required_methods": ["get_tools"],
            "schema": {
                "type": "object",
                "required": ["get_tools"],
                "properties": {
                    "get_tools": {
                        "type": "method",
                        "returns": "Toolkit"
                    }
                }
            }
        }

    def validate_custom_tool(self, tool: Any) -> tuple[bool, Optional[str]]:
        """
        Validate a custom tool for Agno runtime.

        Args:
            tool: Tool instance to validate

        Returns:
            Tuple of (is_valid, error_message)
        """
        # Check for get_tools method
        if not hasattr(tool, 'get_tools'):
            return False, "Tool must have get_tools() method"

        if not callable(getattr(tool, 'get_tools')):
            return False, "get_tools must be callable"

        # Try calling to validate return type
        try:
            toolkit = tool.get_tools()

            # Check if it's a Toolkit-like object (has tools attribute or is iterable)
            if not (hasattr(toolkit, 'tools') or hasattr(toolkit, '__iter__')):
                return False, f"get_tools() must return Toolkit or iterable, got {type(toolkit)}"

        except Exception as e:
            return False, f"get_tools() failed: {str(e)}"

        return True, None

    def register_custom_tool(self, tool: Any, metadata: Optional[Dict] = None) -> str:
        """
        Register a custom tool with Agno runtime.

        Args:
            tool: Tool instance with get_tools() method
            metadata: Optional metadata (name, description, etc.)

        Returns:
            Tool identifier for this registered tool

        Raises:
            ValueError: If tool validation fails
        """
        # Validate first
        is_valid, error = self.validate_custom_tool(tool)
        if not is_valid:
            raise ValueError(f"Invalid custom tool: {error}")

        # Generate tool ID
        tool_name = metadata.get("name") if metadata else tool.__class__.__name__
        tool_id = f"custom_{tool_name}_{id(tool)}"

        # Store tool instance
        self._custom_tools[tool_id] = tool

        self.logger.info(
            "custom_tool_registered",
            tool_id=tool_id,
            tool_class=tool.__class__.__name__,
            tool_name=tool_name
        )

        return tool_id

    def get_registered_custom_tools(self) -> list[str]:
        """
        Get list of registered custom tool identifiers.

        Returns:
            List of tool IDs
        """
        return list(self._custom_tools.keys())
