# 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

from abc import ABC
import os
from typing import Any
from typing import Optional
from typing import TYPE_CHECKING

from deprecated import deprecated
from google.genai import types

from .tool_context import ToolContext

if TYPE_CHECKING:
  from ..models.llm_request import LlmRequest


class BaseTool(ABC):
  """The base class for all tools."""

  name: str
  """The name of the tool."""
  description: str
  """The description of the tool."""

  is_long_running: bool = False
  """Whether the tool is a long running operation, which typically returns a
  resource id first and finishes the operation later."""

  def __init__(self, *, name, description, is_long_running: bool = False):
    self.name = name
    self.description = description
    self.is_long_running = is_long_running

  def _get_declaration(self) -> Optional[types.FunctionDeclaration]:
    """Gets the OpenAPI specification of this tool in the form of a FunctionDeclaration.

    NOTE
    - Required if subclass uses the default implementation of
      `process_llm_request` to add function declaration to LLM request.
    - Otherwise, can be skipped, e.g. for a built-in GoogleSearch tool for
      Gemini.

    Returns:
      The FunctionDeclaration of this tool, or None if it doesn't need to be
      added to LlmRequest.config.
    """
    return None

  async def run_async(
      self, *, args: dict[str, Any], tool_context: ToolContext
  ) -> Any:
    """Runs the tool with the given arguments and context.

    NOTE
    - Required if this tool needs to run at the client side.
    - Otherwise, can be skipped, e.g. for a built-in GoogleSearch tool for
      Gemini.

    Args:
      args: The LLM-filled arguments.
      tool_context: The context of the tool.

    Returns:
      The result of running the tool.
    """
    raise NotImplementedError(f'{type(self)} is not implemented')

  async def process_llm_request(
      self, *, tool_context: ToolContext, llm_request: LlmRequest
  ) -> None:
    """Processes the outgoing LLM request for this tool.

    Use cases:
    - Most common use case is adding this tool to the LLM request.
    - Some tools may just preprocess the LLM request before it's sent out.

    Args:
      tool_context: The context of the tool.
      llm_request: The outgoing LLM request, mutable this method.
    """
    if (function_declaration := self._get_declaration()) is None:
      return

    llm_request.tools_dict[self.name] = self
    if tool_with_function_declarations := _find_tool_with_function_declarations(
        llm_request
    ):
      if tool_with_function_declarations.function_declarations is None:
        tool_with_function_declarations.function_declarations = []
      tool_with_function_declarations.function_declarations.append(
          function_declaration
      )
    else:
      llm_request.config = (
          types.GenerateContentConfig()
          if not llm_request.config
          else llm_request.config
      )
      llm_request.config.tools = (
          [] if not llm_request.config.tools else llm_request.config.tools
      )
      llm_request.config.tools.append(
          types.Tool(function_declarations=[function_declaration])
      )

  @property
  def _api_variant(self) -> str:
    use_vertexai = os.environ.get('GOOGLE_GENAI_USE_VERTEXAI', '0').lower() in [
        'true',
        '1',
    ]
    return 'VERTEX_AI' if use_vertexai else 'GOOGLE_AI'


def _find_tool_with_function_declarations(
    llm_request: LlmRequest,
) -> Optional[types.Tool]:
  # TODO: add individual tool with declaration and merge in google_llm.py
  if not llm_request.config or not llm_request.config.tools:
    return None

  return next(
      (
          tool
          for tool in llm_request.config.tools
          if isinstance(tool, types.Tool) and tool.function_declarations
      ),
      None,
  )