# Copyright 2025 Wei Sun (Jack)
#
# 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.

"""Base agent class for adkx extensions."""

from __future__ import annotations

from typing import AsyncGenerator

from google.adk.agents import BaseAgent as ADKBaseAgent
from google.adk.events.event import Event
from pydantic import ConfigDict
from pydantic import Field
from typing_extensions import override


class BaseAgent(ADKBaseAgent):
  """Extended base class for adkx agents.

  Inherits all functionality from google.adk.agents.BaseAgent but disables
  agent hierarchies (parent_agent and sub_agents).

  Inherited from ADKBaseAgent:
  - name: str - The agent's name
  - description: str - Description of the agent's capability
  - before_agent_callback: Optional callback before agent run
  - after_agent_callback: Optional callback after agent run

  Disabled from ADKBaseAgent:
  - parent_agent (forbidden)
  - sub_agents (forbidden)

  All fields are immutable after creation.
  """

  model_config = ConfigDict(
      frozen=True,
      arbitrary_types_allowed=True,
      extra='forbid',
  )

  # Forbid parent_agent and sub_agents
  parent_agent: None = Field(default=None, init=False, exclude=True, repr=False)
  sub_agents: list = Field(
      default_factory=list, init=False, exclude=True, repr=False
  )

  @override
  async def run_async(
      self,
      parent_context,
  ) -> AsyncGenerator[Event, None]:
    """Entry method to run an agent without tracing.

    This override removes the tracing span from the parent implementation,
    providing a lighter-weight execution path.

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

    Yields:
      Event: the events generated by the agent.
    """
    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
