# 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.

"""Simplified Agent implementation."""

from __future__ import annotations

import asyncio
from typing import AsyncGenerator
import uuid

from google.adk.agents.invocation_context import InvocationContext
from google.adk.agents.run_config import StreamingMode
from google.adk.events.event import Event
from google.adk.events.event_actions import EventActions
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.tools.tool_context import ToolContext
from google.genai import types
from pydantic import Field
from typing_extensions import override

from ..models import Gemini
from ..tools import BaseTool
from ..tools import ToolResult
from .base_agent import BaseAgent


class Agent(BaseAgent):
  """LLM-based agent.

  This is a simplified agent with just the essential fields:
  - model: The LLM to use
  - instruction: What the agent should do
  - tools: Tools available to the agent

  All fields are immutable after creation.

  Extension Points:
    Subclasses can override these protected methods to customize behavior:
    - _build_identity_instruction(): Customize agent identity presentation
    - _build_system_instruction(): Customize system instructions
    - _build_tools(): Customize tool selection and registration
    - _build_conversation(): Customize conversation history handling
  """

  model: str | BaseLlm
  """The model to use for this agent.

  Either a model name (e.g., "gemini-2.5-flash", "gemini-2.0-flash") or an
  instance of a BaseLlm subclass for custom LLM implementations.
  """

  instruction: str = ""
  """Instructions for what this agent should do.

  Describes the agent's role, responsibilities, and behavior.
  """

  tools: list[BaseTool] = Field(default_factory=list)
  """Tools available to the agent.

  List of tools that the agent can invoke during execution.
  """

  # ============================================================================
  # Core Implementation
  # ============================================================================

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

    Implements LLM loop with tool execution:
    1. Build LlmRequest (system instruction, conversation history, tools)
    2. Loop until no more function calls:
       - Call LLM and yield model response event
       - Extract function calls from response
       - If function calls exist, execute tools concurrently
       - Yield function response event with merged EventActions
    3. Exit when LLM returns no function calls

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

    Yields:
      Event: the events generated by the agent.
    """
    # Tool execution loop
    # TODO: Optimize by reusing LLM client across invocations once we have
    # proper lifecycle management (e.g., server shutdown hooks to close clients)
    llm_client = (
        Gemini(model=self.model) if isinstance(self.model, str) else self.model
    )

    while True:
      # Build LLM request from session
      llm_request = self._build_llm_request(ctx)

      # Call LLM (with streaming if SSE mode, keep last response)
      llm_response = None
      async for response in llm_client.generate_content_async(
          llm_request,
          stream=bool(
              ctx.run_config
              and ctx.run_config.streaming_mode == StreamingMode.SSE
          ),
      ):
        if response.partial:
          yield self._create_model_event(response, ctx)
        else:
          llm_response = response

      if llm_response is None:
        break

      # Ensure all function calls have IDs before yielding the event
      self._ensure_function_call_ids(llm_response.content)

      # Yield model response event
      yield self._create_model_event(llm_response, ctx)

      # Check for function calls
      function_calls = self._extract_function_calls(llm_response.content)
      if not function_calls:
        break

      # Execute tools and yield response
      function_response_parts, event_actions = await self._execute_all_tools(
          function_calls, ctx
      )
      yield self._create_function_response_event(
          function_response_parts, event_actions, ctx
      )

  # ============================================================================
  # LLM Request Building
  # ============================================================================

  def _build_llm_request(self, ctx: InvocationContext) -> LlmRequest:
    """Build LlmRequest from session events and agent configuration."""
    # Only pass model string to LlmRequest if self.model is a string
    # If it's a BaseLlm instance, the model is handled by the instance itself
    model_str = self.model if isinstance(self.model, str) else None
    llm_request = LlmRequest(model=model_str)
    self._build_system_instruction(llm_request, ctx)
    self._build_conversation(llm_request, ctx)
    self._build_tools(llm_request, ctx)
    return llm_request

  def _build_identity_instruction(self) -> str:
    """Build identity instruction for the agent.

    Extension point: Override to customize how the agent introduces itself
    (e.g., custom name format, additional identity attributes, role clarification).

    Returns:
      Identity instruction describing the agent's name and description.
    """
    parts = [f'You are known as "{self.name}".']

    if self.description:
      parts.append(f'The description about you is "{self.description}".')

    return " ".join(parts)

  def _build_system_instruction(
      self, llm_request: LlmRequest, ctx: InvocationContext
  ) -> None:
    """Add system instruction to LlmRequest.

    Extension point: Override to customize system instruction behavior
    (e.g., add context-specific instructions, dynamic instructions).

    Args:
      llm_request: The LLM request to add system instruction to.
      ctx: The invocation context.
    """
    instructions = []

    # Add identity instruction
    identity_instruction = self._build_identity_instruction()
    if identity_instruction:
      instructions.append(identity_instruction)

    # Add custom instruction
    if self.instruction:
      instructions.append(self.instruction)

    if instructions:
      llm_request.append_instructions(instructions)

  def _build_tools(
      self, llm_request: LlmRequest, ctx: InvocationContext
  ) -> None:
    """Add tools to LlmRequest.

    Extension point: Override to customize tool selection and registration
    (e.g., dynamic tool filtering, context-dependent tool availability).

    Args:
      llm_request: The LLM request to add tools to.
      ctx: The invocation context.
    """
    llm_request.append_tools(self.tools)

  def _build_conversation(
      self, llm_request: LlmRequest, ctx: InvocationContext
  ) -> None:
    """Add session event contents to LlmRequest.

    Extension point: Override to customize conversation history handling
    (e.g., event filtering, summarization, context window management).

    Args:
      llm_request: The LLM request to add contents to.
      ctx: The invocation context containing session events.
    """
    for event in ctx.session.events:
      if event.content:
        llm_request.contents.append(event.content)

  # ============================================================================
  # Event Creation
  # ============================================================================

  def _create_model_event(
      self, llm_response: LlmResponse, ctx: InvocationContext
  ) -> Event:
    """Create event for model response."""
    return Event(
        **llm_response.model_dump(exclude_none=True),
        author=self.name,
        branch=ctx.branch,
        invocation_id=ctx.invocation_id,
    )

  def _create_function_response_event(
      self,
      function_response_parts: list[types.Part],
      event_actions: EventActions,
      ctx: InvocationContext,
  ) -> Event:
    """Create event for function responses."""
    function_response_content = types.UserContent(parts=function_response_parts)
    return Event(
        author=self.name,
        content=function_response_content,
        actions=event_actions,
        branch=ctx.branch,
        invocation_id=ctx.invocation_id,
    )

  # ============================================================================
  # Tool Execution
  # ============================================================================

  async def _execute_all_tools(
      self, function_calls: list[types.FunctionCall], ctx: InvocationContext
  ) -> tuple[list[types.Part], EventActions]:
    """Execute all function calls concurrently and collect response parts and actions.

    Returns:
      Tuple of (function response parts, merged EventActions from all tools).
    """
    # Create contexts for all tool calls
    tool_contexts = [
        self._create_tool_context(function_call, ctx)
        for function_call in function_calls
    ]

    # Execute all tools concurrently
    tool_results = await asyncio.gather(*[
        self._execute_tool(function_call, tool_context)
        for function_call, tool_context in zip(function_calls, tool_contexts)
    ])

    # Collect response parts and merge actions
    function_response_parts: list[types.Part] = []
    merged_actions = EventActions()

    for function_call, tool_result, tool_context in zip(
        function_calls, tool_results, tool_contexts
    ):
      parts = self._tool_result_to_parts(function_call, tool_result)
      function_response_parts.extend(parts)
      self._merge_event_actions(merged_actions, tool_context.actions)

    return function_response_parts, merged_actions

  def _ensure_function_call_ids(self, content: types.Content | None) -> None:
    """Ensure all function calls have IDs, generating adk-{uuid} if missing.

    Args:
      content: LLM response content to process.
    """
    if not content or not content.parts:
      return

    for part in content.parts:
      function_call = part.function_call
      if not function_call:
        continue

      if not function_call.id:
        function_call.id = f"adk-{uuid.uuid4()}"

  def _extract_function_calls(
      self, content: types.Content | None
  ) -> list[types.FunctionCall]:
    """Extract function calls from LLM response content.

    Args:
      content: The content from LLM response.

    Returns:
      List of function calls found in the content.
    """
    if not content or not content.parts:
      return []

    function_calls: list[types.FunctionCall] = []
    for part in content.parts:
      if part.function_call:
        function_calls.append(part.function_call)

    return function_calls

  async def _execute_tool(
      self, function_call: types.FunctionCall, tool_context: ToolContext
  ) -> ToolResult:
    """Execute a tool based on function call.

    Args:
      function_call: The function call from LLM.
      tool_context: The tool execution context.

    Returns:
      The tool execution result as a ToolResult.
    """
    # Validate function call has a name
    if not function_call.name:
      return ToolResult(status="error", details=["Function call missing name"])

    # Find the tool
    tool = self._find_tool_by_name(function_call.name)
    if not tool:
      return ToolResult(
          status="error",
          details=[f"Tool '{function_call.name}' not found"],
      )

    # Execute the tool
    try:
      return await tool.run_async(
          args=function_call.args or {}, tool_context=tool_context
      )
    except Exception as e:
      return ToolResult(status="error", details=[str(e)])

  def _find_tool_by_name(self, tool_name: str) -> BaseTool | None:
    """Find a tool by name, returns None if not found."""
    return next((t for t in self.tools if t.name == tool_name), None)

  def _create_tool_context(
      self, function_call: types.FunctionCall, ctx: InvocationContext
  ) -> ToolContext:
    """Create ToolContext for tool execution."""
    return ToolContext(
        invocation_context=ctx,
        function_call_id=function_call.id,
        event_actions=EventActions(),
    )

  def _tool_result_to_parts(
      self, function_call: types.FunctionCall, tool_result: ToolResult
  ) -> list[types.Part]:
    """Convert ToolResult to Parts for LLM response.

    Args:
      function_call: The original function call from the LLM.
      tool_result: The result from tool execution.

    Returns:
      List of Parts with FunctionResponse (with matching ID) and any media parts.
    """
    return tool_result.to_parts(
        name=function_call.name or "", id=function_call.id or ""
    )

  def _merge_event_actions(
      self, target: EventActions, source: EventActions
  ) -> None:
    """Merge event actions from source into target.

    Args:
      target: The target EventActions to merge into.
      source: The source EventActions to merge from.
    """
    # Merge scalar fields (take source if set)
    if source.skip_summarization is not None:
      target.skip_summarization = source.skip_summarization
    if source.transfer_to_agent is not None:
      target.transfer_to_agent = source.transfer_to_agent
    if source.escalate is not None:
      target.escalate = source.escalate
    if source.compaction is not None:
      target.compaction = source.compaction
    if source.end_of_agent is not None:
      target.end_of_agent = source.end_of_agent
    if source.agent_state is not None:
      target.agent_state = source.agent_state
    if source.rewind_before_invocation_id is not None:
      target.rewind_before_invocation_id = source.rewind_before_invocation_id

    # Merge dict fields (extend)
    target.state_delta.update(source.state_delta)
    target.artifact_delta.update(source.artifact_delta)
    target.requested_auth_configs.update(source.requested_auth_configs)
    target.requested_tool_confirmations.update(
        source.requested_tool_confirmations
    )
