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

from __future__ import annotations

import json
from typing import Any
from typing import Optional
from typing import TYPE_CHECKING

from google.genai import types
from pydantic import BaseModel
from pydantic import Field
from pydantic import model_validator
from typing_extensions import override

from . import _automatic_function_calling_util
from ..agents.common_configs import AgentRefConfig
from ..events._branch_path import _BranchPath
from ..features import FeatureName
from ..features import is_feature_enabled
from ..memory.in_memory_memory_service import InMemoryMemoryService
from ..utils._schema_utils import SchemaType
from ..utils._schema_utils import validate_schema
from ..utils.context_utils import Aclosing
from ._forwarding_artifact_service import ForwardingArtifactService
from .base_tool import BaseTool
from .tool_configs import BaseToolConfig
from .tool_configs import ToolArgsConfig
from .tool_context import ToolContext

if TYPE_CHECKING:
  from ..agents.base_agent import BaseAgent


def _part_to_text(part: types.Part) -> str:
  """Returns user-visible text from a Part, including code execution output."""
  if part.text:
    return part.text
  if part.code_execution_result and part.code_execution_result.output:
    return part.code_execution_result.output.rstrip('\n')
  if part.executable_code and part.executable_code.code:
    return part.executable_code.code
  return ''


def _get_input_schema(agent: BaseAgent) -> Optional[type[BaseModel]]:
  """Extracts the input_schema from an agent.

  For LlmAgent, returns its input_schema directly.
  For agents with sub_agents, recursively searches the first sub-agent for an
  input_schema.

  Args:
    agent: The agent to extract input_schema from.

  Returns:
    The input_schema if found, None otherwise.
  """
  from ..agents.llm_agent import LlmAgent

  if isinstance(agent, LlmAgent):
    return agent.input_schema

  # For composite agents, check the first sub-agent
  if agent.sub_agents:
    return _get_input_schema(agent.sub_agents[0])

  return None


def _get_output_schema(agent: BaseAgent) -> Optional[SchemaType]:
  """Extracts the output_schema from an agent.

  For LlmAgent, returns its output_schema directly.
  For agents with sub_agents, recursively searches the last sub-agent for an
  output_schema.

  Args:
    agent: The agent to extract output_schema from.

  Returns:
    The output_schema if found, None otherwise.
  """
  from ..agents.llm_agent import LlmAgent

  if isinstance(agent, LlmAgent):
    return agent.output_schema

  # For composite agents, check the last sub-agent
  if agent.sub_agents:
    return _get_output_schema(agent.sub_agents[-1])

  return None


class AgentTool(BaseTool):
  """A tool that wraps an agent.

  This tool allows an agent to be called as a tool within a larger application.
  The agent's input schema is used to define the tool's input parameters, and
  the agent's output is returned as the tool's result.

  Note:
    To expose an agent as an inline tool of a parent ``LlmAgent``, prefer
    setting ``mode='single_turn'`` on the sub-agent and attaching it via
    ``sub_agents=[...]`` instead of wrapping it with ``AgentTool``. The
    framework then exposes the sub-agent as a tool automatically and runs it
    inline in the parent's session. See the single-turn mode guide for details.

  Attributes:
    agent: The agent to wrap.
    skip_summarization: Whether to skip summarization of the agent output.
    include_plugins: Whether to propagate plugins from the parent runner context
      to the agent's runner. When True (default), the agent will inherit all
      plugins from its parent. Set to False to run the agent with an isolated
      plugin environment.
  """

  def __init__(
      self,
      agent: BaseAgent,
      skip_summarization: bool = False,
      *,
      include_plugins: bool = True,
      propagate_grounding_metadata: bool = False,
  ):
    self.agent = agent
    self.skip_summarization: bool = skip_summarization
    self.include_plugins = include_plugins
    self.propagate_grounding_metadata = propagate_grounding_metadata

    super().__init__(name=agent.name, description=agent.description)

  @model_validator(mode='before')
  @classmethod
  def populate_name(cls, data: Any) -> Any:
    data['name'] = data['agent'].name
    return data

  @override
  def _get_declaration(self) -> types.FunctionDeclaration:
    from ..utils.variant_utils import GoogleLLMVariant

    input_schema = _get_input_schema(self.agent)
    output_schema = _get_output_schema(self.agent)

    if input_schema:
      result = _automatic_function_calling_util.build_function_declaration(
          func=input_schema, variant=self._api_variant
      )
      # Override the description with the agent's description
      result.description = self.agent.description
    else:
      if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
        result = types.FunctionDeclaration(
            name=self.name,
            description=self.agent.description,
            parameters_json_schema={
                'type': 'object',
                'properties': {
                    'request': {'type': 'string'},
                },
                'required': ['request'],
            },
        )
      else:
        result = types.FunctionDeclaration(
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    'request': types.Schema(
                        type=types.Type.STRING,
                    ),
                },
                required=['request'],
            ),
            description=self.agent.description,
            name=self.name,
        )

    # Set response schema for non-GEMINI_API variants
    if self._api_variant != GoogleLLMVariant.GEMINI_API:
      # Determine response type based on agent's output schema
      if output_schema:
        # Agent has structured output schema - response is an object
        if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
          result.response_json_schema = {'type': 'object'}
        else:
          result.response = types.Schema(type=types.Type.OBJECT)
      else:
        # Agent returns text - response is a string
        if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
          result.response_json_schema = {'type': 'string'}
        else:
          result.response = types.Schema(type=types.Type.STRING)

    result.name = self.name
    return result

  @override
  async def run_async(
      self,
      *,
      args: dict[str, Any],
      tool_context: ToolContext,
  ) -> Any:
    from ..runners import Runner
    from ..sessions.in_memory_session_service import InMemorySessionService

    if self.skip_summarization:
      tool_context.actions.skip_summarization = True

    input_schema = _get_input_schema(self.agent)
    if input_schema:
      input_value = input_schema.model_validate(args)
      content = types.Content(
          role='user',
          parts=[
              types.Part.from_text(
                  text=input_value.model_dump_json(exclude_none=True)
              )
          ],
      )
    else:
      if 'request' in args:
        request_text = args['request']
      else:
        request_text = json.dumps(args, ensure_ascii=False, sort_keys=True)
      content = types.Content(
          role='user',
          parts=[types.Part.from_text(text=request_text)],
      )
    invocation_context = tool_context._invocation_context
    parent_app_name = (
        invocation_context.app_name if invocation_context else None
    )
    child_app_name = parent_app_name or self.agent.name
    plugins = (
        tool_context._invocation_context.plugin_manager.plugins
        if self.include_plugins
        else None
    )
    runner = Runner(
        app_name=child_app_name,
        agent=self.agent,
        artifact_service=ForwardingArtifactService(tool_context),
        session_service=InMemorySessionService(),
        memory_service=InMemoryMemoryService(),
        credential_service=tool_context._invocation_context.credential_service,
        plugins=plugins,
    )
    # When plugins are inherited from the parent runner, the parent still owns
    # them; tell the sub-Runner's plugin manager to skip closing them on exit
    # so shared plugins (e.g. observability exporters) are not torn down while
    # the parent is still using them.
    if self.include_plugins:
      runner.plugin_manager.set_skip_closing_plugins(True)

    state_dict = {
        k: v
        for k, v in tool_context.state.to_dict().items()
        if not k.startswith('_adk')  # Filter out adk internal states
    }
    session = await runner.session_service.create_session(
        app_name=child_app_name,
        user_id=tool_context._invocation_context.user_id,
        state=state_dict,
    )

    last_content = None
    last_grounding_metadata = None
    async with Aclosing(
        runner.run_async(
            user_id=session.user_id, session_id=session.id, new_message=content
        )
    ) as agen:
      async for event in agen:
        # Forward state delta to parent session.
        if event.actions.state_delta:
          tool_context.state.update(event.actions.state_delta)
        if event.content:
          last_content = event.content
          last_grounding_metadata = event.grounding_metadata

    # Clean up runner resources (especially MCP sessions)
    # to avoid "Attempted to exit cancel scope in a different task" errors
    await runner.close()

    if last_content is None or last_content.parts is None:
      return ''
    parts_text = (_part_to_text(p) for p in last_content.parts if not p.thought)
    merged_text = '\n'.join(t for t in parts_text if t)
    output_schema = _get_output_schema(self.agent)
    if output_schema:
      tool_result = validate_schema(output_schema, merged_text)
    else:
      tool_result = merged_text

    if self.propagate_grounding_metadata and last_grounding_metadata:
      tool_context.state['temp:_adk_grounding_metadata'] = (
          last_grounding_metadata
      )

    return tool_result

  @override
  @classmethod
  def from_config(
      cls, config: ToolArgsConfig, config_abs_path: str
  ) -> AgentTool:
    from ..agents import config_agent_utils

    agent_tool_config = AgentToolConfig.model_validate(config.model_dump())

    agent = config_agent_utils.resolve_agent_reference(
        agent_tool_config.agent, config_abs_path
    )
    return cls(
        agent=agent,
        skip_summarization=agent_tool_config.skip_summarization,
        include_plugins=agent_tool_config.include_plugins,
    )


class AgentToolConfig(BaseToolConfig):
  """The config for the AgentTool."""

  agent: AgentRefConfig
  """The reference to the agent instance."""

  skip_summarization: bool = False
  """Whether to skip summarization of the agent output."""

  include_plugins: bool = True
  """Whether to include plugins from parent runner context."""


class _SingleTurnAgentTool(AgentTool):
  """A tool that wraps a single-turn agent and runs it via ctx.run_node.

  This is only used in mode='chat' LlmAgent.
  """

  @override
  async def run_async(
      self,
      *,
      args: dict[str, Any],
      tool_context: ToolContext,
  ) -> Any:
    input_schema = _get_input_schema(self.agent)
    if input_schema:
      try:
        node_input = input_schema.model_validate(args)
      except Exception as e:
        return f'Error validating input: {e}'
    else:
      node_input = args.get('request')

    # Align subagent branch scoping with node execution (Node as Tool) using function_call_id.
    fc_id = tool_context.function_call_id
    base_branch = tool_context.get_invocation_context().branch
    tool_branch = _BranchPath.create_sub_branch(
        base_branch, name=self.agent.name, run_id=fc_id
    )

    try:
      return await tool_context.run_node(
          self.agent,
          node_input=node_input,
          override_branch=tool_branch,
          use_sub_branch=False,
      )
    except Exception as e:
      return f'Error running sub-agent: {e}'


class _DefaultTaskInput(BaseModel):
  request: str = Field(
      description='Detailed instructions or context for the task sub-agent.'
  )


class _TaskAgentTool(AgentTool):
  """A tool that wraps a task-mode agent and acts as a framework delegation marker.

  This is only used in mode='chat' LlmAgent. The wrapper intercepts calls
  to this tool to drive task sub-agent execution via ctx.run_node.
  """

  def __init__(
      self,
      agent: BaseAgent,
      skip_summarization: bool = False,
      *,
      include_plugins: bool = True,
      propagate_grounding_metadata: bool = False,
  ):
    super().__init__(
        agent,
        skip_summarization,
        include_plugins=include_plugins,
        propagate_grounding_metadata=propagate_grounding_metadata,
    )
    self._defers_response = True

  @override
  def _get_declaration(self) -> types.FunctionDeclaration:
    from ..utils.variant_utils import GoogleLLMVariant

    input_schema = _get_input_schema(self.agent) or _DefaultTaskInput

    from . import _function_tool_declarations

    result = (
        _function_tool_declarations.build_function_declaration_with_json_schema(
            func=input_schema
        )
    )
    base_desc = self.agent.description or ''
    suffix = (
        '\nIMPORTANT: This tool delegates execution to a specialized agent.'
        ' Do NOT call this tool in parallel with any other tools.'
    )
    result.description = f'{base_desc}{suffix}'.strip()
    result.name = self.name

    if self._api_variant != GoogleLLMVariant.GEMINI_API:
      output_schema = _get_output_schema(self.agent)
      if output_schema:
        result.response_json_schema = {'type': 'object'}
      else:
        result.response_json_schema = {'type': 'string'}

    return result

  @override
  async def run_async(
      self,
      *,
      args: dict[str, Any],
      tool_context: ToolContext,
  ) -> Any:
    # Framework handles task delegation dispatch directly via the wrapper.
    return None
