# 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 inspect
import logging
from typing import Awaitable
from typing import Callable
from typing import Optional
import uuid

from a2a.server.agent_execution import AgentExecutor
from a2a.server.agent_execution import RequestContext
from a2a.server.events.event_queue import EventQueue
from a2a.types import Message
from a2a.types import Task
from typing_extensions import override

from .. import _compat
from ...runners import Runner
from ...sessions import base_session_service
from ...utils.context_utils import Aclosing
from ..agent.interceptors.new_integration_extension import _NEW_A2A_ADK_INTEGRATION_EXTENSION
from ..converters.from_adk_event import create_error_status_event
from ..converters.long_running_functions import handle_user_input
from ..converters.long_running_functions import LongRunningFunctions
from ..converters.request_converter import AgentRunRequest
from ..converters.utils import _get_adk_metadata_key
from ..experimental import a2a_experimental
from .config import A2aAgentExecutorConfig
from .executor_context import ExecutorContext
from .utils import execute_after_agent_interceptors
from .utils import execute_after_event_interceptors
from .utils import execute_before_agent_interceptors

logger = logging.getLogger('google_adk.' + __name__)


@a2a_experimental
class _A2aAgentExecutor(AgentExecutor):
  """An AgentExecutor that runs an ADK Agent against an A2A request and

  publishes updates to an event queue.
  """

  def __init__(
      self,
      *,
      runner: Runner | Callable[..., Runner | Awaitable[Runner]],
      config: Optional[A2aAgentExecutorConfig] = None,
  ):
    super().__init__()
    self._runner = runner
    self._config = config or A2aAgentExecutorConfig()

  @override
  async def cancel(self, context: RequestContext, event_queue: EventQueue):
    """Cancel the execution."""
    # TODO: Implement proper cancellation logic if needed
    raise NotImplementedError('Cancellation is not supported')

  @override
  async def execute(
      self,
      context: RequestContext,
      event_queue: EventQueue,
  ):
    """Executes an A2A request and publishes updates to the event queue

    specified. It runs as following:
    * Takes the input from the A2A request
    * Convert the input to ADK input content, and runs the ADK agent
    * Collects output events of the underlying ADK Agent
    * Converts the ADK output events into A2A task updates
    * Publishes the updates back to A2A server via event queue
    """
    if not context.message:
      raise ValueError('A2A request must have a message')

    context = await execute_before_agent_interceptors(
        context, self._config.execute_interceptors
    )

    runner = await self._resolve_runner()
    try:
      run_request = self._config.request_converter(
          context,
          self._config.a2a_part_converter,
      )
      await self._resolve_session(run_request, runner)

      executor_context = ExecutorContext(
          app_name=runner.app_name,
          user_id=run_request.user_id,
          session_id=run_request.session_id,
          runner=runner,
      )

      # for new task, create a task submitted event
      if not context.current_task:
        await event_queue.enqueue_event(
            Task(
                id=context.task_id,
                status=_compat.make_task_status(_compat.TS_SUBMITTED),
                context_id=context.context_id,
                history=[context.message],
                metadata=self._get_invocation_metadata(executor_context),
            )
        )
      else:
        # Check if the user input is responding to the agent's
        # request for input.
        missing_user_input_event = handle_user_input(context)
        if missing_user_input_event:
          _compat.set_event_metadata(
              missing_user_input_event,
              self._get_invocation_metadata(executor_context),
          )
          await event_queue.enqueue_event(missing_user_input_event)
          return

      await event_queue.enqueue_event(
          _compat.make_task_status_update_event(
              task_id=context.task_id,
              context_id=context.context_id,
              status=_compat.make_task_status(_compat.TS_WORKING),
              final=False,
              metadata=self._get_invocation_metadata(executor_context),
          )
      )

      # Handle the request and publish updates to the event queue
      await self._handle_request(
          context,
          executor_context,
          event_queue,
          runner,
          run_request,
      )
    except Exception as e:
      logger.error('Error handling A2A request: %s', e, exc_info=True)
      # Publish failure event
      try:
        await event_queue.enqueue_event(
            _compat.make_task_status_update_event(
                task_id=context.task_id,
                context_id=context.context_id,
                status=_compat.make_task_status(
                    _compat.TS_FAILED,
                    message=Message(
                        message_id=str(uuid.uuid4()),
                        role=_compat.ROLE_AGENT,
                        parts=[_compat.make_text_part(str(e))],
                    ),
                ),
                final=True,
            )
        )
      except Exception as enqueue_error:
        logger.error(
            'Failed to publish failure event: %s', enqueue_error, exc_info=True
        )

  async def _handle_request(
      self,
      context: RequestContext,
      executor_context: ExecutorContext,
      event_queue: EventQueue,
      runner: Runner,
      run_request: AgentRunRequest,
  ):
    agents_artifact: dict[str, str] = {}
    error_event = None
    long_running_functions = LongRunningFunctions(
        self._config.gen_ai_part_converter
    )
    async with Aclosing(runner.run_async(**vars(run_request))) as agen:
      async for adk_event in agen:
        # Handle error scenarios
        if adk_event and (adk_event.error_code or adk_event.error_message):
          error_event = create_error_status_event(
              adk_event,
              context.task_id,
              context.context_id,
          )

        # Handle long running function calls
        adk_event = long_running_functions.process_event(adk_event)

        for a2a_event in self._config.adk_event_converter(
            adk_event,
            agents_artifact,
            context.task_id,
            context.context_id,
            self._config.gen_ai_part_converter,
        ):
          _compat.set_event_metadata(
              a2a_event, self._get_invocation_metadata(executor_context)
          )
          a2a_events = await execute_after_event_interceptors(
              a2a_event,
              executor_context,
              adk_event,
              self._config.execute_interceptors,
          )
          for e in a2a_events:
            await event_queue.enqueue_event(e)

    if error_event:
      final_event = error_event
    elif long_running_functions.has_long_running_function_calls():
      final_event = (
          long_running_functions.create_long_running_function_call_event(
              context.task_id, context.context_id
          )
      )
    else:
      final_event = _compat.make_task_status_update_event(
          task_id=context.task_id,
          context_id=context.context_id,
          status=_compat.make_task_status(_compat.TS_COMPLETED),
          final=True,
      )

    _compat.set_event_metadata(
        final_event, self._get_invocation_metadata(executor_context)
    )
    final_event = await execute_after_agent_interceptors(
        executor_context, final_event, self._config.execute_interceptors
    )
    await event_queue.enqueue_event(final_event)

  async def _resolve_runner(self) -> Runner:
    """Resolve the runner, handling cases where it's a callable that returns a Runner."""
    if isinstance(self._runner, Runner):
      return self._runner
    if callable(self._runner):
      result = self._runner()

      if inspect.iscoroutine(result):
        resolved_runner = await result
      else:
        resolved_runner = result

      self._runner = resolved_runner
      return resolved_runner

    raise TypeError(
        'Runner must be a Runner instance or a callable that returns a'
        f' Runner, got {type(self._runner)}'
    )

  async def _resolve_session(
      self,
      run_request: AgentRunRequest,
      runner: Runner,
  ):
    session_id = run_request.session_id
    # create a new session if not exists
    user_id = run_request.user_id
    session = await runner.session_service.get_session(
        app_name=runner.app_name,
        user_id=user_id,
        session_id=session_id,
        # Checking existence doesn't require event history.
        config=base_session_service.GetSessionConfig(num_recent_events=0),
    )
    if session is None:
      session = await runner.session_service.create_session(
          app_name=runner.app_name,
          user_id=user_id,
          state={},
          session_id=session_id,
      )
      # Update run_request with the new session_id
      run_request.session_id = session.id

  def _get_invocation_metadata(
      self, executor_context: ExecutorContext
  ) -> dict[str, str]:
    return {
        _get_adk_metadata_key('app_name'): executor_context.app_name,
        _get_adk_metadata_key('user_id'): executor_context.user_id,
        _get_adk_metadata_key('session_id'): executor_context.session_id,
        # TODO: Remove this metadata once the new agent executor
        # is fully adopted.
        _NEW_A2A_ADK_INTEGRATION_EXTENSION: {'adk_agent_executor_v2': True},
    }
