"""
Substrate Agent - The foundational agent for Daita Agents.

This agent provides a blank slate that users can build upon to create
custom agents for any task, with simplified error handling and retry capabilities.
All operations are automatically traced without any configuration required.
"""
import asyncio
import logging
import os
from datetime import datetime
from typing import Dict, Any, Optional, List, Union, Callable

from ..config.base import AgentConfig, AgentType
from ..core.interfaces import LLMProvider
from ..core.exceptions import (
    DaitaError, AgentError, LLMError, PluginError,
    ValidationError, InvalidDataError, NotFoundError
)
from .base import BaseAgent

logger = logging.getLogger(__name__)

# Import unified plugin access
from ..plugins import PluginAccess
from ..llm.factory import create_llm_provider
from ..config.settings import settings
from ..core.tools import AgentTool, ToolRegistry

class SubstrateAgent(BaseAgent):
    """
    The foundational agent with unified prompt system and handler-based task execution.

    SubstrateAgent uses a two-tier prompt system for composable AI behavior:
    - Agent-level prompt (WHO the agent is) - set at creation
    - Call-level instructions (WHAT to do) - set per process() call

    All operations are automatically traced without any configuration required.

    Args:
        name: Agent name
        prompt: Agent identity and default behavior. Can be:
            - str: Single prompt for all tasks
            - dict: Task-specific prompts {"analyze": "...", "transform": "..."}
        focus: Column/field filtering for data operations
        handlers: Custom task handlers {task_name: handler_function}
        tools: Plugin instances or AgentTool objects
        llm_provider: LLM provider name or instance
        model: Model name to use
        api_key: API key for LLM provider (auto-detected if not provided)

    Examples:
        Basic agent:
            >>> agent = SubstrateAgent(name="processor")
            >>> result = await agent.process("analyze", data)

        Agent with identity:
            >>> agent = SubstrateAgent(
            ...     name="validator",
            ...     prompt="You are a data quality expert. Always return JSON."
            ... )
            >>> result = await agent.process("analyze", data)

        Agent with call-specific instructions:
            >>> agent = SubstrateAgent(prompt="You validate data quality.")
            >>> result = await agent.process(
            ...     "analyze",
            ...     data,
            ...     context={"instructions": "Focus on email validation"}
            ... )

        Task-specific prompts:
            >>> agent = SubstrateAgent(
            ...     prompt={
            ...         "analyze": "You are a data analyst.",
            ...         "transform": "You are a data transformer."
            ...     }
            ... )

        Composition (agent identity + specific instructions):
            >>> agent = SubstrateAgent(
            ...     prompt="You are a data quality expert."
            ... )
            >>> result = await agent.process(
            ...     "analyze",
            ...     batch_data,
            ...     context={"instructions": "This is legacy data - be careful"}
            ... )
            >>> # LLM receives both the agent prompt AND the specific instructions

        Data operations example:
            >>> validator = SubstrateAgent(
            ...     name="ETL Validator",
            ...     prompt="You validate data quality for ETL pipelines."
            ... )
            >>> result = await validator.process(
            ...     "analyze",
            ...     staging_data,
            ...     context={"instructions": "Validate against production schema"}
            ... )
    """
    
    # Class-level defaults for smart constructor
    _default_llm_provider = "openai"
    _default_model = "gpt-4"

    @classmethod
    def configure_defaults(cls, **kwargs):
        """Set global defaults for all SubstrateAgent instances."""
        for key, value in kwargs.items():
            setattr(cls, f'_default_{key}', value)

    def __new__(cls, name=None, **kwargs):
        """Smart constructor with auto-configuration."""
        # Auto-configuration from environment and defaults
        if not kwargs.get('llm_provider'):
            kwargs['llm_provider'] = getattr(cls, '_default_llm_provider', 'openai')
        if not kwargs.get('model'):
            kwargs['model'] = getattr(cls, '_default_model', 'gpt-4')
        if not kwargs.get('api_key'):
            provider = kwargs.get('llm_provider', 'openai')
            # Only try to get API key if provider is a string (not an object)
            if isinstance(provider, str):
                kwargs['api_key'] = settings.get_llm_api_key(provider)

        return super().__new__(cls)
    
    def __init__(
        self,
        name: Optional[str] = None,
        llm_provider: Optional[Union[str, LLMProvider]] = None,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        config: Optional[AgentConfig] = None,
        agent_id: Optional[str] = None,
        prompt: Optional[Union[str, Dict[str, str]]] = None,
        focus: Optional[Union[List[str], str, Dict[str, Any]]] = None,
        handlers: Optional[Dict[str, Callable]] = None,
        relay: Optional[str] = None,
        mcp: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None,
        display_reasoning: bool = False,
        **kwargs
    ):
        """
        Initialize the Substrate Agent with smart constructor pattern.

        This constructor auto-creates LLM providers and provides sensible
        defaults while preserving all functionality.

        Args:
            name: Agent name (required for direct instantiation)
            llm_provider: LLM provider name ("openai", "anthropic") or instance
            model: Model name ("gpt-4", "claude-3-sonnet-20240229")
            api_key: API key for LLM provider (auto-detected from env if not provided)
            config: Agent configuration (auto-generated if not provided)
            agent_id: Unique identifier for the agent
            prompt: Custom prompt or prompt templates
            focus: Default focus configuration for data processing
            handlers: Dictionary of task handlers {task_name: handler_function}
            relay: Name of relay channel for publishing results
            mcp: MCP server(s) for tool integration - single dict or list of dicts
            display_reasoning: Enable minimal decision display in console
            **kwargs: Additional configuration options (can include 'tools' parameter)
        """
        # Auto-create LLM provider if needed
        if isinstance(llm_provider, str) or llm_provider is None:
            provider_name = llm_provider or self._default_llm_provider
            model_name = model or self._default_model
            api_key_to_use = api_key or settings.get_llm_api_key(provider_name)
            
            if api_key_to_use:
                llm_provider = create_llm_provider(
                    provider=provider_name,
                    model=model_name,
                    api_key=api_key_to_use,
                    agent_id=agent_id
                )
            else:
                logger.warning(f"No API key found for {provider_name}. LLM functionality will be disabled.")
                llm_provider = None
        # Create default config if none provided
        if config is None:
            config = AgentConfig(
                name=name or "Substrate Agent",
                type=AgentType.SUBSTRATE,
                **kwargs
            )
        
        # Initialize base agent (which handles automatic tracing)
        super().__init__(config, llm_provider, agent_id, name)
        
        # Store customization options
        self.prompt = prompt
        self.default_focus = focus
        self.relay = relay
        
        # Decision display setup
        self.display_reasoning = display_reasoning
        self._decision_display = None
        
        if display_reasoning:
            self._setup_decision_display()

        # Tool management (unified system)
        self.tool_registry = ToolRegistry()
        self.tool_sources = kwargs.get('tools', [])  # Plugins, AgentTool instances, or callables
        self._tools_setup = False

        # MCP server integration
        self.mcp_registry = None
        self.mcp_tools = []
        if mcp is not None:
            # Normalize to list
            mcp_servers = [mcp] if isinstance(mcp, dict) else mcp
            self._mcp_server_configs = mcp_servers
            # MCP setup happens lazily on first use to avoid blocking init
        else:
            self._mcp_server_configs = []

        # Task handlers - users can define custom behavior
        self.handlers: Dict[str, Callable] = handlers or {}

        # Built-in handlers for common operations
        self._register_builtin_handlers()

        # Plugin access for direct plugin usage
        self.plugins = PluginAccess()

        logger.debug(f"Substrate Agent {self.name} initialized with {len(self.handlers)} handlers")
    
    
    def _setup_decision_display(self):
        """Setup minimal decision display for local development."""
        try:
            from ..display.console import create_console_decision_display
            from ..core.decision_tracing import register_agent_decision_stream

            # Create display
            self._decision_display = create_console_decision_display(
                agent_name=self.name,
                agent_id=self.agent_id
            )

            # Register with decision streaming system
            register_agent_decision_stream(
                agent_id=self.agent_id,
                callback=self._decision_display.handle_event
            )

            logger.debug(f"Decision display enabled for agent {self.name}")

        except Exception as e:
            logger.warning(f"Failed to setup decision display: {e}")
            self.display_reasoning = False
            self._decision_display = None

    async def _setup_mcp_tools(self):
        """
        Setup MCP servers and discover available tools.

        This is called lazily on first agent.process() to avoid blocking
        agent initialization with MCP server connections.
        """
        if self.mcp_registry is not None:
            # Already setup
            return

        if not self._mcp_server_configs:
            # No MCP servers configured
            return

        try:
            from ..plugins.mcp import MCPServer, MCPToolRegistry

            logger.info(f"Setting up {len(self._mcp_server_configs)} MCP server(s) for {self.name}")

            # Create registry
            self.mcp_registry = MCPToolRegistry()

            # Connect to each server and register tools
            for server_config in self._mcp_server_configs:
                server = MCPServer(
                    command=server_config.get("command"),
                    args=server_config.get("args", []),
                    env=server_config.get("env", {}),
                    server_name=server_config.get("name")
                )

                # Add to registry (automatically connects and discovers tools)
                await self.mcp_registry.add_server(server)

            # Get all tools from registry
            self.mcp_tools = self.mcp_registry.get_all_tools()

            logger.info(f"MCP setup complete: {self.mcp_registry.tool_count} tools from {self.mcp_registry.server_count} server(s)")

        except ImportError:
            logger.error(
                "MCP SDK not installed. Install with: pip install mcp\n"
                "See: https://github.com/modelcontextprotocol/python-sdk"
            )
            raise

        except Exception as e:
            logger.error(f"Failed to setup MCP servers: {str(e)}")
            raise

    async def _setup_tools(self):
        """
        Discover and register tools from all sources.

        Called lazily on first process() call to avoid blocking initialization.
        Sources can be:
        - Plugin instances with get_tools() method
        - AgentTool instances directly
        - MCP server configurations
        """
        if self._tools_setup:
            return  # Already setup

        # 1. Setup MCP tools first
        if self._mcp_server_configs and self.mcp_registry is None:
            await self._setup_mcp_tools()
            # Convert MCP tools to AgentTool format
            for mcp_tool in self.mcp_tools:
                agent_tool = AgentTool.from_mcp_tool(mcp_tool, self.mcp_registry)
                self.tool_registry.register(agent_tool)

        # 2. Register plugin tools
        for source in self.tool_sources:
            if isinstance(source, AgentTool):
                # Direct AgentTool registration
                self.tool_registry.register(source)
                logger.debug(f"Registered tool: {source.name}")

            elif hasattr(source, 'get_tools'):
                # Plugin with get_tools() method
                plugin_tools = source.get_tools()
                if plugin_tools:
                    self.tool_registry.register_many(plugin_tools)
                    logger.info(
                        f"Registered {len(plugin_tools)} tools from "
                        f"{source.__class__.__name__}"
                    )
            else:
                logger.warning(
                    f"Invalid tool source: {source}. "
                    f"Expected AgentTool or plugin with get_tools() method."
                )

        self._tools_setup = True
        logger.info(
            f"Agent {self.name} initialized with {self.tool_registry.tool_count} tools"
        )

    async def _call_llm(
        self,
        task: str,
        data: Any,
        context: Dict[str, Any],
        default_system_prompt: Optional[str] = None
    ) -> str:
        """
        Centralized LLM gateway - ALL LLM calls go through here.

        Composes: agent prompt + context instructions + data

        Args:
            task: Task name (for task-specific prompts)
            data: Input data to process
            context: Call-level context with optional 'instructions'
            default_system_prompt: Fallback if no agent prompt configured

        Returns:
            LLM response string

        Raises:
            AgentError: If no LLM configured
        """
        if not self.llm:
            raise AgentError(f"No LLM provider configured for task: {task}")

        # 1. Build system message (agent identity/behavior)
        system_prompt = None
        if self.prompt:
            if isinstance(self.prompt, dict):
                # Task-specific prompt
                system_prompt = self.prompt.get(task, self.prompt.get('default'))
            else:
                # Single prompt for all tasks
                system_prompt = str(self.prompt)

        # Use default if no agent prompt
        if not system_prompt and default_system_prompt:
            system_prompt = default_system_prompt

        # 2. Build user message (instructions + data)
        user_parts = []

        # Add call-specific instructions
        instructions = context.get('instructions')
        if instructions:
            user_parts.append(str(instructions))

        # Add data
        if data is not None:
            # Smart data formatting
            if isinstance(data, str):
                data_str = data
            elif isinstance(data, (list, dict)):
                import json
                try:
                    data_str = json.dumps(data, indent=2, default=str)
                except:
                    data_str = str(data)
            else:
                data_str = str(data)

            # Truncate very long data
            if len(data_str) > 10000:
                data_str = data_str[:10000] + "\n... (truncated)"

            user_parts.append(f"Data:\n{data_str}")

        user_message = "\n\n".join(user_parts) if user_parts else ""

        # 3. Call LLM with proper structure
        # Check if provider supports system messages (most do)
        if hasattr(self.llm, 'generate_with_system'):
            response = await self.llm.generate_with_system(
                system=system_prompt,
                user=user_message
            )
        else:
            # Fallback: concatenate into single prompt
            if system_prompt:
                full_prompt = f"{system_prompt}\n\n{user_message}"
            else:
                full_prompt = user_message
            response = await self.llm.generate(full_prompt)

        return response

    def _register_builtin_handlers(self):
        """Register built-in handlers for common tasks (without overwriting custom handlers)."""
        builtin_handlers = {
            'analyze': self._handle_analyze,
            'transform': self._handle_transform,
            'process_data': self._handle_process_data,
            'custom': self._handle_custom,
            'relay_message': self._handle_relay_message,
            'llm_query': self._handle_llm_query
        }

        # Only add built-in handlers if custom handler doesn't already exist
        for task, handler in builtin_handlers.items():
            if task not in self.handlers:
                self.handlers[task] = handler
    
    async def _process_once(
        self,
        task: str,
        data: Any,
        context: Dict[str, Any],
        attempt: int,
        max_attempts: int
    ) -> Dict[str, Any]:
        """
        Execute the task once using the SubstrateAgent's handler system.

        This override provides the SubstrateAgent's specific task routing logic
        while letting the parent handle retry and tracing automatically.
        """
        # Setup all tools (lazy initialization - includes MCP and plugin tools)
        await self._setup_tools()

        # Apply focus if specified
        focus_to_use = context.get('focus') or self.default_focus
        if focus_to_use and data is not None:
            try:
                from ..core.focus import apply_focus
                data = apply_focus(data, focus_to_use)
            except Exception as e:
                logger.warning(f"Focus application failed: {str(e)}")
                # Continue with original data if focus fails
        
        # Find the appropriate handler
        handler = None
        if task in self.handlers:
            handler = self.handlers[task]
        elif self.llm:
            # Fall back to LLM query if no specific handler found
            handler = self._handle_llm_query
        else:
            # No handler and no LLM available - use default behavior
            return self._handle_default(data, context, self)

        # Execute the handler
        try:
            if asyncio.iscoroutinefunction(handler):
                result = await handler(data, context, self)
            else:
                result = handler(data, context, self)
            
            # Ensure result is a dictionary
            if not isinstance(result, dict):
                result = {"handler_result": result}
            
            # Add metadata
            result.update({
                "task": task,
                "agent_id": self.agent_id,
                "agent_name": self.name,
                "attempt": attempt,
                "handler_used": handler.__name__ if hasattr(handler, '__name__') else str(handler),
                "timestamp": datetime.utcnow().isoformat()
            })
            
            # Publish to relay if configured and result is successful
            if self.relay:
                await self._publish_to_relay(result, context)
            
            return result
            
        except Exception as e:
            logger.error(f"Handler execution failed in {self.name}: {str(e)}")
            raise
    
    # Built-in Handlers
    
    async def _handle_llm_query(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
        """Handle direct LLM queries using unified gateway."""
        if not self.llm:
            raise AgentError("No LLM provider configured for LLM query")

        # Use gateway for consistent behavior
        response = await self._call_llm(
            task=context.get('task', 'llm_query'),
            data=data,
            context=context,
            default_system_prompt="You are a helpful AI assistant. Process the user's request."
        )

        return {
            "agent_response": response,
            "llm_provider": self.llm.provider_name if hasattr(self.llm, 'provider_name') else 'unknown'
        }

    async def _llm_with_tools(self, prompt: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Handle LLM query with MCP tool calling support.

        This implements the tool calling flow:
        1. Send prompt with available tools to LLM
        2. If LLM requests tool call, execute it via MCP
        3. Return tool result to LLM for final response
        4. Return final response to user
        """
        # Convert MCP tools to LLM function format
        tool_definitions = [tool.to_llm_function() for tool in self.mcp_tools]

        logger.debug(f"LLM query with {len(tool_definitions)} MCP tools available")

        # For now, we'll do a simple implementation that appends tool info to the prompt
        # Full function calling integration requires LLM provider updates
        # This is MVP: tools are described, LLM can reference them in response

        # Build tool descriptions for prompt
        tool_descriptions = "\n\nAvailable tools:\n"
        for tool in self.mcp_tools:
            tool_descriptions += f"- {tool.name}: {tool.description}\n"

        enhanced_prompt = f"{prompt}{tool_descriptions}"

        response = await self.llm.generate(enhanced_prompt)

        return {
            "agent_response": response,
            "prompt_used": enhanced_prompt,
            "mcp_tools_available": [t.name for t in self.mcp_tools],
            "llm_provider": self.llm.provider_name if hasattr(self.llm, 'provider_name') else 'unknown'
        }

    async def call_mcp_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
        """
        Manually call an MCP tool.

        This allows users to explicitly call MCP tools from their handlers
        or for testing purposes.

        Args:
            tool_name: Name of the MCP tool to call
            arguments: Arguments for the tool

        Returns:
            Tool execution result

        Example:
            ```python
            result = await agent.call_mcp_tool("read_file", {"path": "/data/file.txt"})
            ```
        """
        if not self.mcp_registry:
            raise RuntimeError("No MCP servers configured. Add mcp parameter to SubstrateAgent.")

        return await self.mcp_registry.call_tool(tool_name, arguments)
    
    async def _handle_analyze(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
        """Handle analysis tasks with unified prompt system."""
        try:
            # Basic data type analysis (non-LLM)
            analysis_result = {
                "data_type": type(data).__name__,
                "data_size": len(data) if hasattr(data, '__len__') else None,
                "keys": list(data.keys()) if isinstance(data, dict) else None,
                "sample": str(data)[:200] if data else None
            }

            # LLM analysis using gateway
            if self.llm:
                llm_analysis = await self._call_llm(
                    task="analyze",
                    data=data,
                    context=context,
                    default_system_prompt="You are a data analyst. Analyze the provided data and provide detailed insights."
                )
                analysis_result["llm_analysis"] = llm_analysis

            return analysis_result

        except Exception as e:
            logger.error(f"Analysis failed: {str(e)}")
            return {
                "analysis_error": str(e),
                "data_type": type(data).__name__,
                "message": "Analysis completed with errors"
            }
    
    async def _handle_transform(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
        """Handle transformation tasks with LLM support."""
        try:
            # Use LLM if available
            if self.llm:
                transformed = await self._call_llm(
                    task="transform",
                    data=data,
                    context=context,
                    default_system_prompt="You are a data transformer. Transform the data according to instructions."
                )

                return {
                    "transformed": transformed,
                    "transformation_type": "llm_guided",
                    "original_type": type(data).__name__
                }

            # Fallback: basic transformations if no LLM
            if isinstance(data, dict):
                transformed = {k: str(v).upper() if isinstance(v, str) else v for k, v in data.items()}
                transformation_type = "dict_string_uppercase"
            elif isinstance(data, list):
                transformed = [str(item).title() if isinstance(item, str) else item for item in data]
                transformation_type = "list_string_title"
            elif isinstance(data, str):
                transformed = data.title()
                transformation_type = "string_title"
            else:
                transformed = str(data)
                transformation_type = "to_string"

            return {
                "transformed": transformed,
                "transformation_type": transformation_type,
                "original_type": type(data).__name__
            }

        except Exception as e:
            logger.error(f"Transformation failed: {str(e)}")
            return {
                "transformation_error": str(e),
                "original_data": data,
                "message": "Transformation completed with errors"
            }
    
    async def _handle_process_data(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
        """Handle generic data processing with optional LLM."""
        result = {
            'status': 'success',
            'message': 'Data processed',
            'input_type': type(data).__name__,
            'processed_at': datetime.utcnow().isoformat(),
            'task': 'process_data'
        }

        # Add LLM processing if available
        if self.llm:
            processed = await self._call_llm(
                task="process_data",
                data=data,
                context=context,
                default_system_prompt="Process this data as requested."
            )
            result['llm_processed'] = processed

        return result
    
    async def _handle_custom(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
        """Handle custom tasks."""
        return {
            'status': 'success',
            'message': 'Custom task executed',
            'data_received': data is not None,
            'context_keys': list(context.keys()),
            'task': 'custom'
        }
    
    async def _handle_relay_message(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
        """Handle incoming relay messages from other agents."""
        return {
            'status': 'success',
            'message': 'Received relay message',
            'source_agent': context.get('source_agent', 'unknown'),
            'data_preview': str(data)[:200] if data else None,
            'task': 'relay_message'
        }
    
    def _handle_default(self, data: Any, context: Dict[str, Any], agent) -> Dict[str, Any]:
        """Default handler when no specific handler is found."""
        return {
            'status': 'success',
            'message': f'Default handler executed for task: {context.get("task", "unknown")}',
            'data_type': type(data).__name__,
            'data_received': data is not None,
            'task': context.get('task', 'unknown')
        }
    
    # User customization methods
    
    def add_handler(self, task: str, handler: Callable):
        """Add a custom task handler."""
        if not task:
            raise ValidationError("Task name cannot be empty")
        if not callable(handler):
            raise ValidationError("Handler must be callable")
        
        self.handlers[task] = handler
        logger.debug(f"Added handler for task: {task}")
    
    def remove_handler(self, task: str):
        """Remove a task handler."""
        if task in self.handlers:
            del self.handlers[task]
            logger.debug(f"Removed handler for task: {task}")
    
    def add_plugin(self, plugin: Any):
        """
        Add a plugin to the agent's tool sources.

        The plugin's tools will be registered on next tool setup.
        """
        self.tool_sources.append(plugin)
        logger.debug(f"Added plugin: {plugin.__class__.__name__}")

    def register_tool(self, tool: AgentTool) -> None:
        """
        Register a single tool manually.

        Useful for adding custom tools after agent initialization.

        Args:
            tool: AgentTool instance to register

        Example:
            ```python
            agent = SubstrateAgent(name="my_agent")

            custom_tool = AgentTool.from_function(my_custom_function)
            agent.register_tool(custom_tool)
            ```
        """
        self.tool_registry.register(tool)

    def register_tools(self, tools: List[AgentTool]) -> None:
        """
        Register multiple tools manually.

        Args:
            tools: List of AgentTool instances
        """
        self.tool_registry.register_many(tools)

    async def call_tool(self, name: str, arguments: Dict[str, Any]) -> Any:
        """
        Execute a tool by name.

        Provides manual tool execution for testing or custom handlers.

        Args:
            name: Tool name
            arguments: Tool arguments dict

        Returns:
            Tool execution result

        Example:
            ```python
            result = await agent.call_tool("query_database", {"sql": "SELECT 1"})
            ```
        """
        await self._setup_tools()
        return await self.tool_registry.execute(name, arguments)

    @property
    def available_tools(self) -> List[AgentTool]:
        """
        Get list of all available tools.

        Returns:
            List of AgentTool instances
        """
        return self.tool_registry.tools.copy()

    @property
    def tool_names(self) -> List[str]:
        """Get list of all tool names"""
        return self.tool_registry.tool_names

    async def stop(self) -> None:
        """Stop agent and clean up all resources including MCP connections."""
        # Clean up MCP connections first
        if self.mcp_registry:
            try:
                await self.mcp_registry.disconnect_all()
                logger.info(f"Cleaned up MCP connections for agent {self.name}")
            except Exception as e:
                logger.warning(f"Error cleaning up MCP connections: {e}")

        # Call parent stop for standard cleanup
        await super().stop()

    def get_token_usage(self) -> Dict[str, int]:
        """
        Get token usage for this agent using automatic tracing.

        Returns comprehensive token statistics from the unified tracing system.
        """
        if not self.llm or not hasattr(self.llm, 'get_token_stats'):
            # Fallback for agents without LLM or tracing
            return {
                'total_tokens': 0,
                'prompt_tokens': 0,
                'completion_tokens': 0,
                'requests': 0
            }

        return self.llm.get_token_stats()
    
    async def _publish_to_relay(self, result: Dict[str, Any], context: Dict[str, Any]):
        """Publish result to relay channel."""
        try:
            from ..core.relay import publish
            
            await publish(
                channel=self.relay,
                agent_response=result,
                publisher=self.name
            )
            logger.debug(f"Published result to relay channel: {self.relay}")
        except Exception as e:
            logger.warning(f"Failed to publish to relay channel {self.relay}: {str(e)}")
            # Don't re-raise - relay failures shouldn't break main processing
    
    @property
    def health(self) -> Dict[str, Any]:
        """Enhanced health information for SubstrateAgent."""
        base_health = super().health

        # Add SubstrateAgent-specific health info
        base_health.update({
            'handlers': {
                'count': len(self.handlers),
                'builtin': len([h for h in self.handlers.keys() if not h.startswith('_')]),
                'custom': len([h for h in self.handlers.keys() if h.startswith('_')])
            },
            'tools': {
                'count': self.tool_registry.tool_count,
                'setup': self._tools_setup,
                'names': self.tool_registry.tool_names if self._tools_setup else []
            },
            'relay': {
                'enabled': self.relay is not None,
                'channel': self.relay
            },
            'llm': {
                'available': self.llm is not None,
                'provider': self.llm.provider_name if self.llm and hasattr(self.llm, 'provider_name') else None
            }
        })

        return base_health