"""
Handles ADK Agent and Runner initialization, including tool loading and callback assignment.
"""

from typing import Dict, List, Optional, Union, Callable, Tuple, Set, Any, TYPE_CHECKING, Type
import functools
import inspect
from solace_ai_connector.common.log import log
from solace_ai_connector.common.utils import import_module

from .app_llm_agent import AppLlmAgent
from .tool_wrapper import ADKToolWrapper
from .embed_resolving_mcp_toolset import EmbedResolvingMCPToolset
from google.adk.runners import Runner
from google.adk.models import BaseLlm
from google.adk.tools import BaseTool, ToolContext
from google.adk import tools as adk_tools_module
from google.adk.agents.callback_context import CallbackContext
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.tools.mcp_tool.mcp_session_manager import (
    SseServerParams,
    StdioConnectionParams,
)

from mcp import StdioServerParameters

if TYPE_CHECKING:
    from ..sac.component import SamAgentComponent

from ..tools.registry import tool_registry
from ..tools.tool_definition import BuiltinTool
from ..tools.dynamic_tool import DynamicTool, DynamicToolProvider
from ..tools.tool_config_types import AnyToolConfig


from ...agent.adk import callbacks as adk_callbacks
from ...agent.adk.models.lite_llm import LiteLlm


def _find_dynamic_tool_class(module) -> Optional[type]:
    """Finds a single non-abstract DynamicTool subclass in a module."""
    found_classes = []
    for name, obj in inspect.getmembers(module, inspect.isclass):
        if (
            issubclass(obj, DynamicTool)
            and obj is not DynamicTool
            and not inspect.isabstract(obj)
        ):
            found_classes.append(obj)
    if len(found_classes) > 1:
        raise TypeError(
            f"Module '{module.__name__}' contains multiple DynamicTool subclasses. "
            "Please specify which one to use with 'class_name' in the config."
        )
    return found_classes[0] if found_classes else None


async def _execute_lifecycle_hook(
    component: "SamAgentComponent",
    func_name: Optional[str],
    module_name: str,
    base_path: Optional[str],
    tool_config_model: AnyToolConfig,
):
    """Dynamically loads and executes a lifecycle hook function."""
    if not func_name:
        return

    log.info(
        "%s Executing lifecycle hook: %s.%s",
        component.log_identifier,
        module_name,
        func_name,
    )

    try:
        module = import_module(module_name, base_path=base_path)
        func = getattr(module, func_name)

        if not inspect.iscoroutinefunction(func):
            raise TypeError(
                f"Lifecycle hook '{func_name}' in module '{module_name}' must be an async function."
            )

        await func(component, tool_config_model)
        log.info(
            "%s Successfully executed lifecycle hook: %s.%s",
            component.log_identifier,
            module_name,
            func_name,
        )
    except Exception as e:
        log.exception(
            "%s Fatal error during lifecycle hook execution for '%s.%s': %s",
            component.log_identifier,
            module_name,
            func_name,
            e,
        )
        raise RuntimeError(f"Tool lifecycle initialization failed: {e}") from e


def _create_cleanup_partial(
    component: "SamAgentComponent",
    func_name: Optional[str],
    module_name: str,
    base_path: Optional[str],
    tool_config_model: AnyToolConfig,
) -> Optional[Callable]:
    """Creates a functools.partial for a cleanup hook function."""
    if not func_name:
        return None

    try:
        module = import_module(module_name, base_path=base_path)
        func = getattr(module, func_name)

        if not inspect.iscoroutinefunction(func):
            raise TypeError(
                f"Lifecycle hook '{func_name}' in module '{module_name}' must be an async function."
            )

        return functools.partial(func, component, tool_config_model)
    except Exception as e:
        log.exception(
            "%s Fatal error creating partial for cleanup hook '%s.%s': %s",
            component.log_identifier,
            module_name,
            func_name,
            e,
        )
        raise RuntimeError(f"Tool lifecycle setup failed: {e}") from e


def _find_dynamic_tool_provider_class(module) -> Optional[type]:
    """Finds a single non-abstract DynamicToolProvider subclass in a module."""
    found_classes = []
    for name, obj in inspect.getmembers(module, inspect.isclass):
        if (
            issubclass(obj, DynamicToolProvider)
            and obj is not DynamicToolProvider
            and not inspect.isabstract(obj)
        ):
            found_classes.append(obj)
    if len(found_classes) > 1:
        raise TypeError(
            f"Module '{module.__name__}' contains multiple DynamicToolProvider subclasses. "
            "Only one is permitted per module."
        )
    return found_classes[0] if found_classes else None


async def load_adk_tools(
    component,
) -> Tuple[List[Union[BaseTool, Callable]], List[BuiltinTool], List[Callable]]:
    """
    Loads all configured tools for the agent.
    - Explicitly configured tools (Python, MCP, ADK Built-ins) from YAML.
    - SAM Built-in tools (Artifact, Data, etc.) from the tool registry,
      filtered by agent configuration.

    Args:
        component: The SamAgentComponent instance.

    Returns:
        A tuple containing:
        - A list of loaded tool callables/instances for the ADK agent.
        - A list of enabled BuiltinTool definition objects for prompt generation.
        - A list of awaitable cleanup functions for the tools.

    Raises:
        ImportError: If a configured tool or its dependencies cannot be loaded.
    """
    loaded_tools: List[Union[BaseTool, Callable]] = []
    enabled_builtin_tools: List[BuiltinTool] = []
    loaded_tool_names: Set[str] = set()
    cleanup_hooks: List[Callable] = []
    tools_config = component.get_config("tools", [])

    from pydantic import TypeAdapter, BaseModel, ValidationError

    any_tool_adapter = TypeAdapter(AnyToolConfig)

    def _check_and_register_tool_name(name: str, source: str):
        """Checks for duplicate tool names and raises ValueError if found."""
        if name in loaded_tool_names:
            raise ValueError(
                f"Configuration Error: Duplicate tool name '{name}' found from source '{source}'. "
                "This name is already in use. Please resolve the conflict by renaming or "
                "disabling one of the tools in your agent's configuration."
            )
        loaded_tool_names.add(name)

    if not tools_config:
        log.info(
            "%s No explicit tools configured in 'tools' list.", component.log_identifier
        )
    else:
        log.debug(
            "%s Processing %d tool configurations: %s",
            component.log_identifier,
            len(tools_config),
            [tc.get("tool_type") for tc in tools_config],
        )
        log.info(
            "%s Loading tools from 'tools' list configuration...",
            component.log_identifier,
        )
        for tool_config in tools_config:
            newly_loaded_tools = []
            try:
                tool_config_model = any_tool_adapter.validate_python(tool_config)
                tool_type = tool_config_model.tool_type.lower()

                if tool_type == "python":
                    module_name = tool_config.get("component_module")
                    base_path = tool_config.get("component_base_path")
                    if not module_name:
                        raise ValueError(
                            "'component_module' is required for python tools."
                        )
                    module = import_module(module_name, base_path=base_path)

                    # Case 1: Simple function-based tool
                    if "function_name" in tool_config:
                        function_name = tool_config.get("function_name")
                        tool_name = tool_config.get("tool_name")
                        tool_description = tool_config.get("tool_description")

                        func = getattr(module, function_name)
                        if not callable(func):
                            raise TypeError(
                                f"'{function_name}' in module '{module_name}' is not callable."
                            )

                        specific_tool_config = tool_config.get("tool_config")
                        tool_callable = ADKToolWrapper(
                            func,
                            specific_tool_config,
                            function_name,
                            origin="python",
                            raw_string_args=tool_config.get("raw_string_args", []),
                        )

                        if tool_name:
                            function_name = tool_name
                            tool_callable.__name__ = tool_name

                        if tool_description:
                            tool_callable.__doc__ = tool_description

                        _check_and_register_tool_name(
                            function_name, f"python:{module_name}"
                        )
                        loaded_tools.append(tool_callable)
                        newly_loaded_tools.append(tool_callable)
                        log.info(
                            "%s Loaded Python tool: %s from %s.",
                            component.log_identifier,
                            function_name,
                            module_name,
                        )
                    # Case 2: Advanced class-based dynamic tool or provider
                    else:
                        specific_tool_config = tool_config.get("tool_config")
                        dynamic_tools = []

                        # Determine the class to load
                        tool_class = None
                        class_name = tool_config.get("class_name")
                        if class_name:
                            tool_class = getattr(module, class_name)
                        else:
                            # Auto-discover: provider first, then single tool
                            tool_class = _find_dynamic_tool_provider_class(module)
                            if not tool_class:
                                tool_class = _find_dynamic_tool_class(module)

                        if not tool_class:
                            raise TypeError(
                                f"Module '{module_name}' does not contain a 'function_name' or 'class_name' to load, "
                                "and no DynamicTool or DynamicToolProvider subclass could be auto-discovered."
                            )

                        # Check for a Pydantic model declaration on the tool class
                        config_model: Optional[Type["BaseModel"]] = getattr(
                            tool_class, "config_model", None
                        )
                        validated_config: Union[dict, "BaseModel"] = specific_tool_config

                        if config_model:
                            log.debug(
                                "%s Found config_model '%s' for tool class '%s'. Validating...",
                                component.log_identifier,
                                config_model.__name__,
                                tool_class.__name__,
                            )
                            try:
                                # Validate the raw dict and get a Pydantic model instance
                                validated_config = config_model.model_validate(
                                    specific_tool_config or {}
                                )
                                log.debug(
                                    "%s Successfully validated tool_config for '%s'.",
                                    component.log_identifier,
                                    tool_class.__name__,
                                )
                            except ValidationError as e:
                                # Provide a clear error message and raise
                                error_msg = (
                                    f"Configuration error for tool '{tool_class.__name__}' from module '{module_name}'. "
                                    f"The provided 'tool_config' in your YAML is invalid:\n{e}"
                                )
                                log.error("%s %s", component.log_identifier, error_msg)
                                raise ValueError(error_msg) from e

                        # Instantiate tools from the class
                        if issubclass(tool_class, DynamicToolProvider):
                            provider_instance = tool_class()
                            dynamic_tools = (
                                provider_instance.get_all_tools_for_framework(
                                    tool_config=validated_config
                                )
                            )
                            log.info(
                                "%s Loaded %d tools from DynamicToolProvider '%s' in %s",
                                component.log_identifier,
                                len(dynamic_tools),
                                tool_class.__name__,
                                module_name,
                            )
                        elif issubclass(tool_class, DynamicTool):
                            tool_instance = tool_class(tool_config=validated_config)
                            dynamic_tools = [tool_instance]
                        else:
                            raise TypeError(
                                f"Class '{tool_class.__name__}' in module '{module_name}' is not a valid "
                                "DynamicTool or DynamicToolProvider subclass."
                            )

                        # Process all generated tools
                        for tool in dynamic_tools:
                            tool.origin = "dynamic"
                            declaration = tool._get_declaration()
                            if not declaration:
                                log.warning(
                                    "Dynamic tool '%s' from module '%s' did not generate a valid declaration. Skipping.",
                                    tool.__class__.__name__,
                                    module_name,
                                )
                                continue

                            _check_and_register_tool_name(
                                declaration.name, f"dynamic:{module_name}"
                            )
                            loaded_tools.append(tool)
                            newly_loaded_tools.append(tool)
                            log.info(
                                "%s Loaded dynamic tool: %s from %s",
                                component.log_identifier,
                                declaration.name,
                                module_name,
                            )

                    # --- Lifecycle Hook Execution for Python Tools ---
                    module_name = tool_config_model.component_module
                    base_path = tool_config_model.component_base_path

                    # 1. YAML Init
                    await _execute_lifecycle_hook(
                        component,
                        tool_config_model.init_function,
                        module_name,
                        base_path,
                        tool_config_model,
                    )

                    # 2. DynamicTool Init
                    for tool_instance in newly_loaded_tools:
                        if isinstance(tool_instance, DynamicTool):
                            log.info(
                                "%s Executing .init() method for DynamicTool '%s'.",
                                component.log_identifier,
                                tool_instance.tool_name,
                            )
                            await tool_instance.init(component, tool_config_model)

                    # 3. Collect Cleanup Hooks (LIFO order)
                    yaml_cleanup_partial = _create_cleanup_partial(
                        component,
                        tool_config_model.cleanup_function,
                        module_name,
                        base_path,
                        tool_config_model,
                    )

                    dynamic_cleanup_partials = []
                    for tool_instance in newly_loaded_tools:
                        if isinstance(tool_instance, DynamicTool):
                            dynamic_cleanup_partials.append(
                                functools.partial(
                                    tool_instance.cleanup, component, tool_config_model
                                )
                            )

                    # Add to main list, prepending to get LIFO execution
                    if yaml_cleanup_partial:
                        cleanup_hooks.insert(0, yaml_cleanup_partial)

                    for partial_func in reversed(dynamic_cleanup_partials):
                        cleanup_hooks.insert(0, partial_func)

                elif tool_type == "builtin":
                    tool_name = tool_config.get("tool_name")
                    if not tool_name:
                        raise ValueError("'tool_name' required for builtin tool.")

                    _check_and_register_tool_name(tool_name, "builtin")

                    sam_tool_def = tool_registry.get_tool_by_name(tool_name)
                    if sam_tool_def:
                        specific_tool_config = tool_config.get("tool_config")
                        tool_callable = ADKToolWrapper(
                            sam_tool_def.implementation,
                            specific_tool_config,
                            sam_tool_def.name,
                            origin="builtin",
                            raw_string_args=sam_tool_def.raw_string_args,
                        )
                        loaded_tools.append(tool_callable)
                        newly_loaded_tools.append(tool_callable)
                        enabled_builtin_tools.append(sam_tool_def)
                        log.info(
                            "%s Loaded SAM built-in tool: %s",
                            component.log_identifier,
                            sam_tool_def.name,
                        )
                        continue

                    adk_tool = getattr(adk_tools_module, tool_name, None)
                    if adk_tool and isinstance(adk_tool, (BaseTool, Callable)):
                        adk_tool.origin = "adk_builtin"
                        loaded_tools.append(adk_tool)
                        log.info(
                            "%s Loaded ADK built-in tool: %s",
                            component.log_identifier,
                            tool_name,
                        )
                        continue

                    raise ValueError(
                        f"Built-in tool '{tool_name}' not found in SAM or ADK registry."
                    )

                elif tool_type == "builtin-group":
                    group_name = tool_config.get("group_name")
                    if not group_name:
                        raise ValueError("'group_name' required for builtin-group.")

                    tools_in_group = tool_registry.get_tools_by_category(group_name)
                    if not tools_in_group:
                        log.warning("No tools found for built-in group: %s", group_name)
                        continue

                    initializers_to_run: Dict[Callable, Dict] = {}
                    for tool_def in tools_in_group:
                        if (
                            tool_def.initializer
                            and tool_def.initializer not in initializers_to_run
                        ):
                            initializers_to_run[tool_def.initializer] = tool_config.get(
                                "config", {}
                            )

                    for init_func, init_config in initializers_to_run.items():
                        try:
                            log.info(
                                "%s Running initializer '%s' for tool group '%s'.",
                                component.log_identifier,
                                init_func.__name__,
                                group_name,
                            )
                            init_func(component, init_config)
                            log.info(
                                "%s Successfully executed initializer '%s' for tool group '%s'.",
                                component.log_identifier,
                                init_func.__name__,
                                group_name,
                            )
                        except Exception as e:
                            log.exception(
                                "%s Failed to run initializer '%s' for tool group '%s': %s",
                                component.log_identifier,
                                init_func.__name__,
                                group_name,
                                e,
                            )
                            raise e

                    group_tool_count = 0
                    for tool_def in tools_in_group:
                        _check_and_register_tool_name(
                            tool_def.name, f"builtin-group:{group_name}"
                        )
                        specific_tool_config = tool_config.get("tool_configs", {}).get(
                            tool_def.name
                        )
                        tool_callable = ADKToolWrapper(
                            tool_def.implementation,
                            specific_tool_config,
                            tool_def.name,
                            origin="builtin",
                            raw_string_args=tool_def.raw_string_args,
                        )
                        loaded_tools.append(tool_callable)
                        newly_loaded_tools.append(tool_callable)
                        enabled_builtin_tools.append(tool_def)
                        group_tool_count += 1
                    log.info(
                        "Loaded %d tools from built-in group: %s",
                        group_tool_count,
                        group_name,
                    )

                elif tool_type == "mcp":
                    tool_name = tool_config.get("tool_name")
                    if not tool_name:
                        log.info(
                            "%s No specific 'tool_name' for MCP tool, will load all tools from server unless tool_filter is specified in MCPToolset itself.",
                            component.log_identifier,
                        )

                    connection_params_config = tool_config.get("connection_params")
                    if not connection_params_config:
                        raise ValueError("'connection_params' required for mcp tool.")

                    connection_type = connection_params_config.get("type", "").lower()
                    connection_args = {
                        k: v for k, v in connection_params_config.items() if k != "type"
                    }
                    connection_args["timeout"] = connection_args.get("timeout", 30)

                    environment_variables = tool_config.get("environment_variables")
                    env_param = {}
                    if connection_type == "stdio" and environment_variables:
                        if isinstance(environment_variables, dict):
                            env_param = environment_variables
                            log.debug(
                                "%s Found environment_variables for stdio MCP tool.",
                                component.log_identifier,
                            )
                        else:
                            log.warning(
                                "%s 'environment_variables' provided for stdio MCP tool but it is not a dictionary. Ignoring.",
                                component.log_identifier,
                            )

                    if connection_type == "stdio":
                        cmd_arg = connection_args.get("command")
                        args_list = connection_args.get("args", [])
                        if isinstance(cmd_arg, list):
                            command_str = " ".join(cmd_arg)
                        elif isinstance(cmd_arg, str):
                            command_str = cmd_arg
                        else:
                            raise ValueError(
                                f"MCP tool 'command' parameter must be a string or a list of strings, got {type(cmd_arg)}"
                            )
                        if not isinstance(args_list, list):
                            raise ValueError(
                                f"MCP tool 'args' parameter must be a list, got {type(args_list)}"
                            )
                        final_connection_args = {
                            k: v
                            for k, v in connection_args.items()
                            if k not in ["command", "args", "timeout"]
                        }
                        connection_params = StdioConnectionParams(
                            server_params=StdioServerParameters(
                                command=command_str,
                                args=args_list,
                                **final_connection_args,
                                env=env_param if env_param else None,
                            ),
                            timeout=connection_args.get("timeout"),
                        )

                    elif connection_type == "sse":
                        connection_params = SseServerParams(**connection_args)
                    else:
                        raise ValueError(
                            f"Unsupported MCP connection type: {connection_type}"
                        )

                    tool_filter_list = [tool_name] if tool_name else None
                    if tool_filter_list:
                        log.info(
                            "%s MCP tool config specifies tool_name: '%s'. Applying as tool_filter.",
                            component.log_identifier,
                            tool_name,
                        )

                    mcp_toolset_instance = EmbedResolvingMCPToolset(
                        connection_params=connection_params,
                        tool_filter=tool_filter_list,
                        tool_config=tool_config,
                    )
                    mcp_toolset_instance.origin = "mcp"

                    # Check for duplicates from the MCP server
                    try:
                        mcp_tools = await mcp_toolset_instance.get_tools()
                        log.debug(
                            "%s Successfully discovered %d tools from MCP server: %s",
                            component.log_identifier,
                            len(mcp_tools),
                            [tool.name for tool in mcp_tools],
                        )
                        for mcp_tool in mcp_tools:
                            log.debug(
                                "%s Registering MCP tool: %s",
                                component.log_identifier,
                                mcp_tool.name,
                            )
                            _check_and_register_tool_name(mcp_tool.name, "mcp")
                    except Exception as e:
                        log.error(
                            "%s Failed to discover tools from MCP server: %s",
                            component.log_identifier,
                            str(e),
                        )
                        raise

                    loaded_tools.append(mcp_toolset_instance)
                    newly_loaded_tools.append(mcp_toolset_instance)
                    log.info(
                        "%s Initialized MCPToolset (filter: %s) for server: %s",
                        component.log_identifier,
                        (tool_filter_list if tool_filter_list else "none (all tools)"),
                        connection_params,
                    )

                else:
                    log.warning(
                        "%s Unknown tool type '%s' in config: %s",
                        component.log_identifier,
                        tool_type,
                        tool_config,
                    )

            except Exception as e:
                log.error(
                    "%s Failed to load tool config %s: %s",
                    component.log_identifier,
                    tool_config,
                    e,
                )
                raise e

    internal_tool_names = ["_notify_artifact_save"]
    if component.get_config("enable_auto_continuation", True):
        internal_tool_names.append("_continue_generation")

    for tool_name in internal_tool_names:
        try:
            _check_and_register_tool_name(tool_name, "internal")
        except ValueError:
            log.debug(
                "%s Internal tool '%s' was already loaded explicitly. Skipping implicit load.",
                component.log_identifier,
                tool_name,
            )
            continue

        tool_def = tool_registry.get_tool_by_name(tool_name)
        if tool_def:
            # Wrap the implementation to ensure its description is passed to the LLM
            tool_callable = ADKToolWrapper(
                tool_def.implementation,
                None,  # No specific config for internal tools
                tool_def.name,
                origin="internal",
            )

            tool_callable.__doc__ = tool_def.description

            loaded_tools.append(tool_callable)
            enabled_builtin_tools.append(tool_def)
            log.info(
                "%s Implicitly loaded internal framework tool: %s",
                component.log_identifier,
                tool_def.name,
            )
        else:
            log.warning(
                "%s Could not find internal framework tool '%s' in registry. Related features may not work.",
                component.log_identifier,
                tool_name,
            )

    log.info(
        "%s Finished loading tools. Total tools for ADK: %d. Total SAM built-ins for prompt: %d. Total cleanup hooks: %d. Peer tools added dynamically.",
        component.log_identifier,
        len(loaded_tools),
        len(enabled_builtin_tools),
        len(cleanup_hooks),
    )
    return loaded_tools, enabled_builtin_tools, cleanup_hooks


def initialize_adk_agent(
    component,
    loaded_explicit_tools: List[Union[BaseTool, Callable]],
    enabled_builtin_tools: List[BuiltinTool],
) -> AppLlmAgent:
    """
    Initializes the ADK LlmAgent based on component configuration.
    Assigns callbacks for peer tool injection, dynamic instruction injection,
    artifact metadata injection, embed resolution, and logging.

    Args:
        component: The A2A_ADK_HostComponent instance.
        loaded_explicit_tools: The list of pre-loaded non-peer tools.

    Returns:
        An initialized LlmAgent instance.

    Raises:
        ValueError: If configuration is invalid.
        ImportError: If required dependencies are missing.
        Exception: For other initialization errors.
    """
    agent_name = component.get_config("agent_name")
    log.info(
        "%s Initializing ADK Agent '%s' (Peer tools & instructions added via callback)...",
        component.log_identifier,
        agent_name,
    )

    model_config = component.get_config("model")
    adk_model_instance: Union[str, BaseLlm]
    if isinstance(model_config, str):
        adk_model_instance = model_config
    elif isinstance(model_config, dict):
        if model_config.get("type") is None:
            # Use setdefault to add keys only if they are not already present in the YAML
            model_config.setdefault("num_retries", 3)
            model_config.setdefault("timeout", 120)
            log.info(
                "%s Applying default resilience settings for LiteLlm model (num_retries=%s, timeout=%s). These can be overridden in YAML.",
                component.log_identifier,
                model_config["num_retries"],
                model_config["timeout"],
            )

        try:

            adk_model_instance = LiteLlm(**model_config)
            log.info(
                "%s Initialized LiteLlm model: %s",
                component.log_identifier,
                model_config.get("model"),
            )
        except ImportError:
            log.error(
                "%s LiteLlm dependency not found. Cannot use dictionary model config.",
                component.log_identifier,
            )
            raise
        except Exception as e:
            log.error(
                "%s Failed to initialize model from dictionary config: %s",
                component.log_identifier,
                e,
            )
            raise
    else:
        raise ValueError(
            f"{component.log_identifier} Invalid 'model' configuration type: {type(model_config)}"
        )

    instruction = component._resolve_instruction_provider(
        component.get_config("instruction", "")
    )
    global_instruction = component._resolve_instruction_provider(
        component.get_config("global_instruction", "")
    )
    planner = component.get_config("planner")
    code_executor = component.get_config("code_executor")

    try:
        agent = AppLlmAgent(
            name=agent_name,
            model=adk_model_instance,
            instruction=instruction,
            global_instruction=global_instruction,
            tools=loaded_explicit_tools,
            planner=planner,
            code_executor=code_executor,
        )

        agent.host_component = component
        log.debug(
            "%s Attached host_component reference to AppLlmAgent.",
            component.log_identifier,
        )
        callbacks_in_order_for_before_model = []

        callbacks_in_order_for_before_model.append(
            adk_callbacks.repair_history_callback
        )
        log.info(
            "%s Added repair_history_callback to before_model chain.",
            component.log_identifier,
        )

        if hasattr(component, "_inject_peer_tools_callback"):
            callbacks_in_order_for_before_model.append(
                component._inject_peer_tools_callback
            )
            log.info(
                "%s Added _inject_peer_tools_callback to before_model chain.",
                component.log_identifier,
            )

        if hasattr(component, "_filter_tools_by_capability_callback"):
            callbacks_in_order_for_before_model.append(
                component._filter_tools_by_capability_callback
            )
            log.info(
                "%s Added _filter_tools_by_capability_callback to before_model chain.",
                component.log_identifier,
            )
        if hasattr(component, "_inject_gateway_instructions_callback"):
            callbacks_in_order_for_before_model.append(
                component._inject_gateway_instructions_callback
            )
            log.info(
                "%s Added _inject_gateway_instructions_callback to before_model chain.",
                component.log_identifier,
            )

        dynamic_instruction_callback_with_component = functools.partial(
            adk_callbacks.inject_dynamic_instructions_callback,
            host_component=component,
            active_builtin_tools=enabled_builtin_tools,
        )
        callbacks_in_order_for_before_model.append(
            dynamic_instruction_callback_with_component
        )
        log.info(
            "%s Added inject_dynamic_instructions_callback to before_model chain.",
            component.log_identifier,
        )

        solace_llm_trigger_callback_with_component = functools.partial(
            adk_callbacks.solace_llm_invocation_callback, host_component=component
        )

        def final_before_model_wrapper(
            callback_context: CallbackContext, llm_request: LlmRequest
        ) -> Optional[LlmResponse]:
            early_response: Optional[LlmResponse] = None
            for cb_func in callbacks_in_order_for_before_model:
                response = cb_func(callback_context, llm_request)
                if response:
                    early_response = response
                    break

            solace_llm_trigger_callback_with_component(callback_context, llm_request)

            if early_response:
                return early_response

            return None

        agent.before_model_callback = final_before_model_wrapper
        log.info(
            "%s Final before_model_callback chain (Solace logging now occurs last) assigned to agent.",
            component.log_identifier,
        )

        tool_invocation_start_cb_with_component = functools.partial(
            adk_callbacks.notify_tool_invocation_start_callback,
            host_component=component,
        )
        agent.before_tool_callback = tool_invocation_start_cb_with_component
        log.info(
            "%s Assigned notify_tool_invocation_start_callback as before_tool_callback.",
            component.log_identifier,
        )

        large_response_cb_with_component = functools.partial(
            adk_callbacks.manage_large_mcp_tool_responses_callback,
            host_component=component,
        )
        metadata_injection_cb_with_component = functools.partial(
            adk_callbacks.after_tool_callback_inject_metadata, host_component=component
        )
        track_artifacts_cb_with_component = functools.partial(
            adk_callbacks.track_produced_artifacts_callback, host_component=component
        )
        notify_tool_result_cb_with_component = functools.partial(
            adk_callbacks.notify_tool_execution_result_callback,
            host_component=component,
        )

        async def chained_after_tool_callback(
            tool: BaseTool,
            args: Dict,
            tool_context: ToolContext,
            tool_response: Dict,
        ) -> Optional[Dict]:
            log.debug(
                "%s Tool callback chain started for tool: %s, response type: %s",
                component.log_identifier,
                tool.name,
                type(tool_response).__name__,
            )

            try:
                # First, notify the UI about the raw result.
                # This is a fire-and-forget notification that does not modify the response.
                notify_tool_result_cb_with_component(
                    tool, args, tool_context, tool_response
                )

                # Now, proceed with the existing chain that modifies the response for the LLM.
                processed_by_large_handler = await large_response_cb_with_component(
                    tool, args, tool_context, tool_response
                )
                response_for_metadata_injector = (
                    processed_by_large_handler
                    if processed_by_large_handler is not None
                    else tool_response
                )

                final_response_after_metadata = (
                    await metadata_injection_cb_with_component(
                        tool, args, tool_context, response_for_metadata_injector
                    )
                )

                final_result = (
                    final_response_after_metadata
                    if final_response_after_metadata is not None
                    else response_for_metadata_injector
                )

                # Track produced artifacts. This callback does not modify the response.
                await track_artifacts_cb_with_component(
                    tool, args, tool_context, final_result
                )

                log.debug(
                    "%s Tool callback chain completed for tool: %s, final response type: %s",
                    component.log_identifier,
                    tool.name,
                    type(final_result).__name__,
                )

                return final_result

            except Exception as e:
                log.exception(
                    "%s Error in tool callback chain for tool %s: %s",
                    component.log_identifier,
                    tool.name,
                    e,
                )
                return tool_response

        agent.after_tool_callback = chained_after_tool_callback
        log.info(
            "%s Chained 'manage_large_mcp_tool_responses_callback' and 'after_tool_callback_inject_metadata' as after_tool_callback.",
            component.log_identifier,
        )

        # --- After Model Callbacks Chain ---
        # The callbacks are executed in the order they are added to this list.
        callbacks_in_order_for_after_model = []

        # 1. Fenced Artifact Block Processing (must run before auto-continue)
        artifact_block_cb = functools.partial(
            adk_callbacks.process_artifact_blocks_callback, host_component=component
        )
        callbacks_in_order_for_after_model.append(artifact_block_cb)
        log.info(
            "%s Added process_artifact_blocks_callback to after_model chain.",
            component.log_identifier,
        )

        # 2. Auto-Continuation (may short-circuit the chain)
        auto_continue_cb = functools.partial(
            adk_callbacks.auto_continue_on_max_tokens_callback, host_component=component
        )
        callbacks_in_order_for_after_model.append(auto_continue_cb)
        log.info(
            "%s Added auto_continue_on_max_tokens_callback to after_model chain.",
            component.log_identifier,
        )

        # 3. Solace LLM Response Logging
        solace_llm_response_cb = functools.partial(
            adk_callbacks.solace_llm_response_callback, host_component=component
        )
        callbacks_in_order_for_after_model.append(solace_llm_response_cb)

        # 4. Chunk Logging
        log_chunk_cb = functools.partial(
            adk_callbacks.log_streaming_chunk_callback, host_component=component
        )
        callbacks_in_order_for_after_model.append(log_chunk_cb)

        async def final_after_model_wrapper(
            callback_context: CallbackContext, llm_response: LlmResponse
        ) -> Optional[LlmResponse]:
            for cb_func in callbacks_in_order_for_after_model:
                # Await async callbacks, call sync callbacks
                if inspect.iscoroutinefunction(cb_func):
                    response = await cb_func(callback_context, llm_response)
                else:
                    response = cb_func(callback_context, llm_response)

                # If a callback returns a response, it hijacks the flow.
                if response:
                    return response
            return None

        agent.after_model_callback = final_after_model_wrapper
        log.info(
            "%s Chained all after_model_callbacks and assigned to agent.",
            component.log_identifier,
        )

        log.info(
            "%s ADK Agent '%s' created. Callbacks assigned.",
            component.log_identifier,
            agent_name,
        )
        return agent
    except Exception as e:
        log.error(
            "%s Failed to create ADK Agent '%s': %s",
            component.log_identifier,
            agent_name,
            e,
        )
        raise


def initialize_adk_runner(component) -> Runner:
    """
    Initializes the ADK Runner.

    Args:
        component: The A2A_ADK_HostComponent instance.

    Returns:
        An initialized Runner instance.

    Raises:
        Exception: For runner initialization errors.
    """
    agent_name = component.get_config("agent_name")
    log.info(
        "%s Initializing ADK Runner for agent '%s'...",
        component.log_identifier,
        agent_name,
    )
    try:
        runner = Runner(
            app_name=agent_name,
            agent=component.adk_agent,
            session_service=component.session_service,
            artifact_service=component.artifact_service,
            memory_service=component.memory_service,
        )
        log.info("%s ADK Runner created successfully.", component.log_identifier)
        return runner
    except Exception as e:
        log.error("%s Failed to create ADK Runner: %s", component.log_identifier, e)
        raise
