# Copyright 2025 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
from typing import Any
from typing import AsyncGenerator
from typing import Awaitable
from typing import Callable
from typing import final
from typing import Optional
from typing import TYPE_CHECKING
from typing import Union

from google.genai import types
from opentelemetry import trace
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
from pydantic import field_validator
from typing_extensions import override
from typing_extensions import TypeAlias

from ..events.event import Event
from .callback_context import CallbackContext

if TYPE_CHECKING:
  from .invocation_context import InvocationContext

tracer = trace.get_tracer('gcp.vertex.agent')

_SingleAgentCallback: TypeAlias = Callable[
    [CallbackContext],
    Union[Awaitable[Optional[types.Content]], Optional[types.Content]],
]

BeforeAgentCallback: TypeAlias = Union[
    _SingleAgentCallback,
    list[_SingleAgentCallback],
]

AfterAgentCallback: TypeAlias = Union[
    _SingleAgentCallback,
    list[_SingleAgentCallback],
]


class BaseAgent(BaseModel):
  """Base class for all agents in Agent Development Kit."""

  model_config = ConfigDict(
      arbitrary_types_allowed=True,
      extra='forbid',
  )
  """The pydantic model config."""

  name: str
  """The agent's name.

  Agent name must be a Python identifier and unique within the agent tree.
  Agent name cannot be "user", since it's reserved for end-user's input.
  """

  description: str = ''
  """Description about the agent's capability.

  The model uses this to determine whether to delegate control to the agent.
  One-line description is enough and preferred.
  """

  parent_agent: Optional[BaseAgent] = Field(default=None, init=False)
  """The parent agent of this agent.

  Note that an agent can ONLY be added as sub-agent once.

  If you want to add one agent twice as sub-agent, consider to create two agent
  instances with identical config, but with different name and add them to the
  agent tree.
  """
  sub_agents: list[BaseAgent] = Field(default_factory=list)
  """The sub-agents of this agent."""

  before_agent_callback: Optional[BeforeAgentCallback] = None
  """Callback or list of callbacks to be invoked before the agent run.

  When a list of callbacks is provided, the callbacks will be called in the
  order they are listed until a callback does not return None.

  Args:
    callback_context: MUST be named 'callback_context' (enforced).

  Returns:
    Optional[types.Content]: The content to return to the user.
      When the content is present, the agent run will be skipped and the
      provided content will be returned to user.
  """
  after_agent_callback: Optional[AfterAgentCallback] = None
  """Callback or list of callbacks to be invoked after the agent run.

  When a list of callbacks is provided, the callbacks will be called in the
  order they are listed until a callback does not return None.

  Args:
    callback_context: MUST be named 'callback_context' (enforced).

  Returns:
    Optional[types.Content]: The content to return to the user.
      When the content is present, the provided content will be used as agent
      response and appended to event history as agent response.
  """

  @final
  async def run_async(
      self,
      parent_context: InvocationContext,
  ) -> AsyncGenerator[Event, None]:
    """Entry method to run an agent via text-based conversation.

    Args:
      parent_context: InvocationContext, the invocation context of the parent
        agent.

    Yields:
      Event: the events generated by the agent.
    """

    with tracer.start_as_current_span(f'agent_run [{self.name}]'):
      ctx = self._create_invocation_context(parent_context)

      if event := await self.__handle_before_agent_callback(ctx):
        yield event
      if ctx.end_invocation:
        return

      async for event in self._run_async_impl(ctx):
        yield event

      if ctx.end_invocation:
        return

      if event := await self.__handle_after_agent_callback(ctx):
        yield event

  @final
  async def run_live(
      self,
      parent_context: InvocationContext,
  ) -> AsyncGenerator[Event, None]:
    """Entry method to run an agent via video/audio-based conversation.

    Args:
      parent_context: InvocationContext, the invocation context of the parent
        agent.

    Yields:
      Event: the events generated by the agent.
    """
    with tracer.start_as_current_span(f'agent_run [{self.name}]'):
      ctx = self._create_invocation_context(parent_context)
      # TODO(hangfei): support before/after_agent_callback

      async for event in self._run_live_impl(ctx):
        yield event

  async def _run_async_impl(
      self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    """Core logic to run this agent via text-based conversation.

    Args:
      ctx: InvocationContext, the invocation context for this agent.

    Yields:
      Event: the events generated by the agent.
    """
    raise NotImplementedError(
        f'_run_async_impl for {type(self)} is not implemented.'
    )
    yield  # AsyncGenerator requires having at least one yield statement

  async def _run_live_impl(
      self, ctx: InvocationContext
  ) -> AsyncGenerator[Event, None]:
    """Core logic to run this agent via video/audio-based conversation.

    Args:
      ctx: InvocationContext, the invocation context for this agent.

    Yields:
      Event: the events generated by the agent.
    """
    raise NotImplementedError(
        f'_run_live_impl for {type(self)} is not implemented.'
    )
    yield  # AsyncGenerator requires having at least one yield statement

  @property
  def root_agent(self) -> BaseAgent:
    """Gets the root agent of this agent."""
    root_agent = self
    while root_agent.parent_agent is not None:
      root_agent = root_agent.parent_agent
    return root_agent

  def find_agent(self, name: str) -> Optional[BaseAgent]:
    """Finds the agent with the given name in this agent and its descendants.

    Args:
      name: The name of the agent to find.

    Returns:
      The agent with the matching name, or None if no such agent is found.
    """
    if self.name == name:
      return self
    return self.find_sub_agent(name)

  def find_sub_agent(self, name: str) -> Optional[BaseAgent]:
    """Finds the agent with the given name in this agent's descendants.

    Args:
      name: The name of the agent to find.

    Returns:
      The agent with the matching name, or None if no such agent is found.
    """
    for sub_agent in self.sub_agents:
      if result := sub_agent.find_agent(name):
        return result
    return None

  def _create_invocation_context(
      self, parent_context: InvocationContext
  ) -> InvocationContext:
    """Creates a new invocation context for this agent."""
    invocation_context = parent_context.model_copy(update={'agent': self})
    return invocation_context

  @property
  def canonical_before_agent_callbacks(self) -> list[_SingleAgentCallback]:
    """The resolved self.before_agent_callback field as a list of _SingleAgentCallback.

    This method is only for use by Agent Development Kit.
    """
    if not self.before_agent_callback:
      return []
    if isinstance(self.before_agent_callback, list):
      return self.before_agent_callback
    return [self.before_agent_callback]

  @property
  def canonical_after_agent_callbacks(self) -> list[_SingleAgentCallback]:
    """The resolved self.after_agent_callback field as a list of _SingleAgentCallback.

    This method is only for use by Agent Development Kit.
    """
    if not self.after_agent_callback:
      return []
    if isinstance(self.after_agent_callback, list):
      return self.after_agent_callback
    return [self.after_agent_callback]

  async def __handle_before_agent_callback(
      self, ctx: InvocationContext
  ) -> Optional[Event]:
    """Runs the before_agent_callback if it exists.

    Returns:
      Optional[Event]: an event if callback provides content or changed state.
    """
    ret_event = None

    if not self.canonical_before_agent_callbacks:
      return ret_event

    callback_context = CallbackContext(ctx)

    for callback in self.canonical_before_agent_callbacks:
      before_agent_callback_content = callback(
          callback_context=callback_context
      )
      if inspect.isawaitable(before_agent_callback_content):
        before_agent_callback_content = await before_agent_callback_content
      if before_agent_callback_content:
        ret_event = Event(
            invocation_id=ctx.invocation_id,
            author=self.name,
            branch=ctx.branch,
            content=before_agent_callback_content,
            actions=callback_context._event_actions,
        )
        ctx.end_invocation = True
        return ret_event

    if callback_context.state.has_delta():
      ret_event = Event(
          invocation_id=ctx.invocation_id,
          author=self.name,
          branch=ctx.branch,
          actions=callback_context._event_actions,
      )

    return ret_event

  async def __handle_after_agent_callback(
      self, invocation_context: InvocationContext
  ) -> Optional[Event]:
    """Runs the after_agent_callback if it exists.

    Returns:
      Optional[Event]: an event if callback provides content or changed state.
    """
    ret_event = None

    if not self.canonical_after_agent_callbacks:
      return ret_event

    callback_context = CallbackContext(invocation_context)

    for callback in self.canonical_after_agent_callbacks:
      after_agent_callback_content = callback(callback_context=callback_context)
      if inspect.isawaitable(after_agent_callback_content):
        after_agent_callback_content = await after_agent_callback_content
      if after_agent_callback_content:
        ret_event = Event(
            invocation_id=invocation_context.invocation_id,
            author=self.name,
            branch=invocation_context.branch,
            content=after_agent_callback_content,
            actions=callback_context._event_actions,
        )
        return ret_event

    if callback_context.state.has_delta():
      ret_event = Event(
          invocation_id=invocation_context.invocation_id,
          author=self.name,
          branch=invocation_context.branch,
          content=after_agent_callback_content,
          actions=callback_context._event_actions,
      )

    return ret_event

  @override
  def model_post_init(self, __context: Any) -> None:
    self.__set_parent_agent_for_sub_agents()

  @field_validator('name', mode='after')
  @classmethod
  def __validate_name(cls, value: str):
    if not value.isidentifier():
      raise ValueError(
          f'Found invalid agent name: `{value}`.'
          ' Agent name must be a valid identifier. It should start with a'
          ' letter (a-z, A-Z) or an underscore (_), and can only contain'
          ' letters, digits (0-9), and underscores.'
      )
    if value == 'user':
      raise ValueError(
          "Agent name cannot be `user`. `user` is reserved for end-user's"
          ' input.'
      )
    return value

  def __set_parent_agent_for_sub_agents(self) -> BaseAgent:
    for sub_agent in self.sub_agents:
      if sub_agent.parent_agent is not None:
        raise ValueError(
            f'Agent `{sub_agent.name}` already has a parent agent, current'
            f' parent: `{sub_agent.parent_agent.name}`, trying to add:'
            f' `{self.name}`'
        )
      sub_agent.parent_agent = self
    return self
