# 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 enum import Enum

from google.adk.features import FeatureName
from google.adk.features._feature_registry import temporary_feature_override
from google.adk.tools import _automatic_function_calling_util
from google.adk.tools.tool_context import ToolContext
from google.adk.utils.variant_utils import GoogleLLMVariant
from google.genai import types
# TODO: crewai requires python 3.10 as minimum
# from crewai_tools import FileReadTool
from pydantic import BaseModel
import pytest


class TestBuildFunctionDeclarationLegacy:

  @pytest.fixture(autouse=True)
  def disable_feature_flag(self):
    """Disable the JSON_SCHEMA_FOR_FUNC_DECL feature flag for legacy tests."""
    with temporary_feature_override(
        FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False
    ):
      yield

  def test_string_input(self):
    def simple_function(input_str: str) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'STRING'

  def test_int_input(self):
    def simple_function(input_str: int) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'INTEGER'

  def test_float_input(self):
    def simple_function(input_str: float) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'NUMBER'

  def test_bool_input(self):
    def simple_function(input_str: bool) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'BOOLEAN'

  def test_array_input(self):
    def simple_function(input_str: list[str]) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'ARRAY'

  def test_dict_input(self):
    def simple_function(input_str: dict[str, str]) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'OBJECT'

  def test_basemodel_input(self):
    class CustomInput(BaseModel):
      input_str: str

    def simple_function(input: CustomInput) -> str:
      return {'result': input}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input'].type == 'OBJECT'
    assert (
        function_decl.parameters.properties['input']
        .properties['input_str']
        .type
        == 'STRING'
    )

  def test_basemodel_required_fields(self):
    class SearchRequest(BaseModel):
      query: str
      max_results: int
      filter: str = ''

    def search(request: SearchRequest) -> list:
      return []

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=search
    )

    inner = function_decl.parameters.properties['request']
    assert set(inner.required) == {'query', 'max_results'}
    assert 'filter' not in (inner.required or [])

  def test_basemodel_all_optional_fields_no_required(self):
    class Config(BaseModel):
      timeout: int = 30
      retries: int = 3

    def run(config: Config) -> str:
      return ''

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=run
    )

    inner = function_decl.parameters.properties['config']
    assert not inner.required

  def test_nested_basemodel_required_fields(self):
    class Inner(BaseModel):
      x: int
      y: int = 0

    class Outer(BaseModel):
      inner: Inner
      label: str = ''

    def process(data: Outer) -> str:
      return ''

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=process
    )

    outer = function_decl.parameters.properties['data']
    assert set(outer.required) == {'inner'}
    assert 'label' not in (outer.required or [])

    inner = outer.properties['inner']
    assert set(inner.required) == {'x'}
    assert 'y' not in (inner.required or [])

  def test_toolcontext_ignored(self):
    def simple_function(input_str: str, tool_context: ToolContext) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function, ignore_params=['tool_context']
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'STRING'
    assert 'tool_context' not in function_decl.parameters.properties

  def test_basemodel(self):
    class SimpleFunction(BaseModel):
      input_str: str
      custom_input: int

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=SimpleFunction, ignore_params=['custom_input']
    )

    assert function_decl.name == 'SimpleFunction'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'STRING'
    assert 'custom_input' not in function_decl.parameters.properties

  def test_nested_basemodel_input(self):
    class ChildInput(BaseModel):
      input_str: str

    class CustomInput(BaseModel):
      child: ChildInput

    def simple_function(input: CustomInput) -> str:
      return {'result': input}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input'].type == 'OBJECT'
    assert (
        function_decl.parameters.properties['input'].properties['child'].type
        == 'OBJECT'
    )
    assert (
        function_decl.parameters.properties['input']
        .properties['child']
        .properties['input_str']
        .type
        == 'STRING'
    )

  def test_basemodel_with_nested_basemodel(self):
    class ChildInput(BaseModel):
      input_str: str

    class CustomInput(BaseModel):
      child: ChildInput

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=CustomInput, ignore_params=['custom_input']
    )

    assert function_decl.name == 'CustomInput'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['child'].type == 'OBJECT'
    assert (
        function_decl.parameters.properties['child']
        .properties['input_str']
        .type
        == 'STRING'
    )
    assert 'custom_input' not in function_decl.parameters.properties

  def test_list(self):
    def simple_function(
        input_str: list[str], input_dir: list[dict[str, str]]
    ) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'ARRAY'
    assert (
        function_decl.parameters.properties['input_str'].items.type == 'STRING'
    )
    assert function_decl.parameters.properties['input_dir'].type == 'ARRAY'
    assert (
        function_decl.parameters.properties['input_dir'].items.type == 'OBJECT'
    )

  def test_enums(self):

    class InputEnum(Enum):
      AGENT = 'agent'
      TOOL = 'tool'

    def simple_function(input: InputEnum = InputEnum.AGENT):
      return input.value

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input'].type == 'STRING'
    assert function_decl.parameters.properties['input'].default == 'agent'
    assert function_decl.parameters.properties['input'].enum == [
        'agent',
        'tool',
    ]

    def simple_function_with_wrong_enum(input: InputEnum = 'WRONG_ENUM'):
      return input.value

    with pytest.raises(ValueError):
      _automatic_function_calling_util.build_function_declaration(
          func=simple_function_with_wrong_enum
      )

  def test_basemodel_list(self):
    class ChildInput(BaseModel):
      input_str: str

    class CustomInput(BaseModel):
      child: ChildInput

    def simple_function(input_str: list[CustomInput]) -> str:
      return {'result': input_str}

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=simple_function
    )

    assert function_decl.name == 'simple_function'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['input_str'].type == 'ARRAY'
    assert (
        function_decl.parameters.properties['input_str'].items.type == 'OBJECT'
    )
    assert (
        function_decl.parameters.properties['input_str']
        .items.properties['child']
        .type
        == 'OBJECT'
    )
    assert (
        function_decl.parameters.properties['input_str']
        .items.properties['child']
        .properties['input_str']
        .type
        == 'STRING'
    )

  # TODO: comment out this test for now as crewai requires python 3.10 as minimum
  # def test_crewai_tool():
  #   docs_tool = CrewaiTool(
  #       name='directory_read_tool',
  #       description='use this to find files for you.',
  #       tool=FileReadTool(),
  #   )
  #   function_decl = docs_tool.get_declaration()
  #   assert function_decl.name == 'directory_read_tool'
  #   assert function_decl.parameters.type == 'OBJECT'
  #   assert function_decl.parameters.properties['file_path'].type == 'STRING'

  def test_function_no_return_annotation_gemini_api(self):
    """Test function with no return annotation using GEMINI_API variant."""

    def function_no_return(param: str):
      """A function with no return annotation."""
      return None

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=function_no_return, variant=GoogleLLMVariant.GEMINI_API
    )

    assert function_decl.name == 'function_no_return'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['param'].type == 'STRING'
    # GEMINI_API should not have response schema
    assert function_decl.response is None

  def test_function_no_return_annotation_vertex_ai(self):
    """Test function with no return annotation using VERTEX_AI variant."""

    def function_no_return(param: str):
      """A function with no return annotation."""
      return None

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=function_no_return, variant=GoogleLLMVariant.VERTEX_AI
    )

    assert function_decl.name == 'function_no_return'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['param'].type == 'STRING'
    # VERTEX_AI should have response schema for functions with no return annotation
    # Changed: Now uses Any type instead of NULL for no return annotation
    assert function_decl.response is not None
    assert (
        function_decl.response.type is None
    )  # Any type maps to None in schema

  def test_function_explicit_none_return_vertex_ai(self):
    """Test function with explicit None return annotation using VERTEX_AI variant."""

    def function_none_return(param: str) -> None:
      """A function that explicitly returns None."""
      pass

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=function_none_return, variant=GoogleLLMVariant.VERTEX_AI
    )

    assert function_decl.name == 'function_none_return'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['param'].type == 'STRING'
    # VERTEX_AI should have response schema for explicit None return
    assert function_decl.response is not None
    assert function_decl.response.type == types.Type.NULL

  def test_function_explicit_none_return_gemini_api(self):
    """Test function with explicit None return annotation using GEMINI_API variant."""

    def function_none_return(param: str) -> None:
      """A function that explicitly returns None."""
      pass

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=function_none_return, variant=GoogleLLMVariant.GEMINI_API
    )

    assert function_decl.name == 'function_none_return'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['param'].type == 'STRING'
    # GEMINI_API should not have response schema
    assert function_decl.response is None

  def test_function_regular_return_type_vertex_ai(self):
    """Test function with regular return type using VERTEX_AI variant."""

    def function_string_return(param: str) -> str:
      """A function that returns a string."""
      return param

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=function_string_return, variant=GoogleLLMVariant.VERTEX_AI
    )

    assert function_decl.name == 'function_string_return'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['param'].type == 'STRING'
    # VERTEX_AI should have response schema for string return
    assert function_decl.response is not None
    assert function_decl.response.type == types.Type.STRING

  def test_function_with_no_response_annotations(self):
    """Test a function that has no response annotations."""

    def transfer_to_agent(agent_name: str, tool_context: ToolContext):
      """Transfer the question to another agent."""
      tool_context.actions.transfer_to_agent = agent_name

    function_decl = _automatic_function_calling_util.build_function_declaration(
        func=transfer_to_agent,
        ignore_params=['tool_context'],
        variant=GoogleLLMVariant.VERTEX_AI,
    )

    assert function_decl.name == 'transfer_to_agent'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['agent_name'].type == 'STRING'
    assert 'tool_context' not in function_decl.parameters.properties
    # This function has no return annotation, so it gets Any type instead of NULL
    # Changed: Now uses Any type instead of NULL for no return annotation
    assert function_decl.response is not None
    assert (
        function_decl.response.type is None
    )  # Any type maps to None in schema

  def test_transfer_to_agent_tool_with_enum_constraint(self):
    """Test TransferToAgentTool adds enum constraint to agent_name."""
    from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool

    agent_names = ['agent_a', 'agent_b', 'agent_c']
    tool = TransferToAgentTool(agent_names=agent_names)

    function_decl = tool._get_declaration()

    assert function_decl.name == 'transfer_to_agent'
    assert function_decl.parameters.type == 'OBJECT'
    assert function_decl.parameters.properties['agent_name'].type == 'STRING'
    assert function_decl.parameters.properties['agent_name'].enum == agent_names
    assert 'tool_context' not in function_decl.parameters.properties


class TestBuildFunctionDeclarationWithJsonSchema:
  """Tests for build_function_declaration when JSON_SCHEMA_FOR_FUNC_DECL is enabled."""

  @pytest.fixture(autouse=True)
  def enable_feature_flag(self):
    """Enable the JSON_SCHEMA_FOR_FUNC_DECL feature flag for all tests."""
    with temporary_feature_override(
        FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True
    ):
      yield

  def test_basic_string_parameter(self):
    """Test basic string parameter with feature flag enabled."""

    def greet(name: str) -> str:
      """Greet someone."""
      return f'Hello, {name}!'

    decl = _automatic_function_calling_util.build_function_declaration(greet)

    assert decl.name == 'greet'
    assert decl.description == 'Greet someone.'
    assert decl.parameters_json_schema == {
        'properties': {'name': {'title': 'Name', 'type': 'string'}},
        'required': ['name'],
        'title': 'greetParams',
        'type': 'object',
    }

  def test_multiple_parameter_types(self):
    """Test multiple parameter types with feature flag enabled."""

    def create_user(name: str, age: int, active: bool) -> str:
      """Create a new user."""
      return f'Created {name}'

    decl = _automatic_function_calling_util.build_function_declaration(
        create_user
    )

    schema = decl.parameters_json_schema
    assert schema['properties'] == {
        'name': {'title': 'Name', 'type': 'string'},
        'age': {'title': 'Age', 'type': 'integer'},
        'active': {'title': 'Active', 'type': 'boolean'},
    }
    assert set(schema['required']) == {'name', 'age', 'active'}

  def test_list_parameter(self):
    """Test list parameter with feature flag enabled."""

    def sum_numbers(numbers: list[int]) -> int:
      """Sum a list of numbers."""
      return sum(numbers)

    decl = _automatic_function_calling_util.build_function_declaration(
        sum_numbers
    )

    schema = decl.parameters_json_schema
    assert schema['properties']['numbers'] == {
        'items': {'type': 'integer'},
        'title': 'Numbers',
        'type': 'array',
    }

  def test_dict_parameter(self):
    """Test dict parameter with feature flag enabled."""

    def process_data(data: dict[str, str]) -> str:
      """Process a dictionary."""
      return str(data)

    decl = _automatic_function_calling_util.build_function_declaration(
        process_data
    )

    schema = decl.parameters_json_schema
    assert schema['properties']['data'] == {
        'additionalProperties': {'type': 'string'},
        'title': 'Data',
        'type': 'object',
    }

  def test_optional_parameter(self):
    """Test optional parameter with feature flag enabled."""

    def search(query: str, limit: int | None = None) -> str:
      """Search for something."""
      return query

    decl = _automatic_function_calling_util.build_function_declaration(search)

    schema = decl.parameters_json_schema
    assert schema['required'] == ['query']
    assert 'query' in schema['properties']
    assert 'limit' in schema['properties']

  def test_enum_parameter(self):
    """Test enum parameter with feature flag enabled."""

    class Color(Enum):
      RED = 'red'
      GREEN = 'green'
      BLUE = 'blue'

    def set_color(color: Color) -> str:
      """Set the color."""
      return color.value

    decl = _automatic_function_calling_util.build_function_declaration(
        set_color
    )

    schema = decl.parameters_json_schema
    assert schema['properties']['color'] == {
        '$ref': '#/$defs/Color',
    }
    assert schema['$defs']['Color'] == {
        'enum': ['red', 'green', 'blue'],
        'title': 'Color',
        'type': 'string',
    }

  def test_tool_context_ignored(self):
    """Test that tool_context is ignored."""

    def my_tool(query: str, tool_context: ToolContext) -> str:
      """A tool that uses context."""
      return query

    decl = _automatic_function_calling_util.build_function_declaration(
        my_tool, ignore_params=['tool_context']
    )

    schema = decl.parameters_json_schema
    assert set(schema['properties'].keys()) == {'query'}
    assert 'tool_context' not in schema['properties']

  def test_gemini_api_no_response_schema(self):
    """Test that GEMINI_API variant does not include response schema."""

    def get_data() -> dict[str, int]:
      """Get some data."""
      return {'count': 42}

    decl = _automatic_function_calling_util.build_function_declaration(
        get_data, variant=GoogleLLMVariant.GEMINI_API
    )

    # GEMINI_API should not have response_json_schema due to bug b/421991354
    assert decl.response_json_schema is None

  @pytest.mark.parametrize(
      'variant, expect_response_schema',
      [
          (GoogleLLMVariant.GEMINI_API, False),
          (GoogleLLMVariant.VERTEX_AI, True),
      ],
  )
  def test_response_schema_by_variant(self, variant, expect_response_schema):
    """Test response schema generation based on the LLM variant."""

    def get_data() -> dict[str, int]:
      """Get some data."""
      return {'count': 42}

    decl = _automatic_function_calling_util.build_function_declaration(
        get_data, variant=variant
    )

    assert (decl.response_json_schema is not None) == expect_response_schema

  def test_pydantic_model_parameter(self):
    """Test Pydantic model parameter with feature flag enabled."""

    class Address(BaseModel):
      street: str
      city: str

    def save_address(address: Address) -> str:
      """Save an address."""
      return f'Saved address in {address.city}'

    decl = _automatic_function_calling_util.build_function_declaration(
        save_address
    )

    assert decl.parameters_json_schema is not None
    assert 'address' in decl.parameters_json_schema['properties']

  def test_no_parameters(self):
    """Test function with no parameters."""

    def get_time() -> str:
      """Get current time."""
      return '12:00'

    decl = _automatic_function_calling_util.build_function_declaration(get_time)

    assert decl.name == 'get_time'
    assert decl.parameters_json_schema is None

  def test_docstring_preserved(self):
    """Test that docstring is preserved as description."""

    def well_documented(x: int) -> int:
      """This is a well-documented function.

      It does something useful.
      """
      return x

    decl = _automatic_function_calling_util.build_function_declaration(
        well_documented
    )

    assert 'well-documented function' in decl.description
    assert 'something useful' in decl.description

  def test_default_values(self):
    """Test parameters with default values."""

    def greet(name: str = 'World') -> str:
      """Greet someone."""
      return f'Hello, {name}!'

    decl = _automatic_function_calling_util.build_function_declaration(greet)

    schema = decl.parameters_json_schema
    assert schema['properties']['name']['default'] == 'World'
    assert 'name' not in schema.get('required', [])