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

"""Extra utils depending on types that are shared between sync and async modules."""

import asyncio
import inspect
import io
import logging
import sys
import typing
from typing import Any, Callable, Dict, Optional, Union, get_args, get_origin
import mimetypes
import os
import pydantic

from . import _common
from . import _mcp_utils
from . import _transformers as t
from . import errors
from . import types
from ._adapters import McpToGenAiToolAdapter


if sys.version_info >= (3, 10):
  from types import UnionType
else:
  UnionType = typing._UnionGenericAlias  # type: ignore[attr-defined]

if typing.TYPE_CHECKING:
  from mcp import ClientSession as McpClientSession
  from mcp.types import Tool as McpTool
else:
  McpClientSession: typing.Type = Any
  McpTool: typing.Type = Any
  try:
    from mcp import ClientSession as McpClientSession
    from mcp.types import Tool as McpTool
  except ImportError:
    McpClientSession = None
    McpTool = None

_DEFAULT_MAX_REMOTE_CALLS_AFC = 10

logger = logging.getLogger('google_genai.models')


def _create_generate_content_config_model(
    config: types.GenerateContentConfigOrDict,
) -> types.GenerateContentConfig:
  if isinstance(config, dict):
    return types.GenerateContentConfig(**config)
  else:
    return config


def _get_gcs_uri(
    src: Union[str, types.BatchJobSourceOrDict]
) -> Optional[str]:
  """Extracts the first GCS URI from the source, if available."""
  if isinstance(src, str) and src.startswith('gs://'):
    return src
  elif isinstance(src, dict) and src.get('gcs_uri'):
    return src['gcs_uri'][0] if src['gcs_uri'] else None
  elif isinstance(src, types.BatchJobSource) and src.gcs_uri:
    return src.gcs_uri[0] if src.gcs_uri else None
  return None


def _get_bigquery_uri(
    src: Union[str, types.BatchJobSourceOrDict]
) -> Optional[str]:
  """Extracts the BigQuery URI from the source, if available."""
  if isinstance(src, str) and src.startswith('bq://'):
    return src
  elif isinstance(src, dict) and src.get('bigquery_uri'):
    return src['bigquery_uri']
  elif isinstance(src, types.BatchJobSource) and src.bigquery_uri:
    return src.bigquery_uri
  return None


def format_destination(
    src: Union[str, types.BatchJobSource],
    config: Optional[types.CreateBatchJobConfig] = None,
) -> types.CreateBatchJobConfig:
  """Formats the destination uri based on the source uri for Vertex AI."""
  if config is None:
    config = types.CreateBatchJobConfig()

  unique_name = None
  if not config.display_name:
    unique_name = _common.timestamped_unique_name()
    config.display_name = f'genai_batch_job_{unique_name}'

  if not config.dest:
    gcs_source_uri = _get_gcs_uri(src)
    bigquery_source_uri = _get_bigquery_uri(src)

    if gcs_source_uri and gcs_source_uri.endswith('.jsonl'):
      config.dest = f'{gcs_source_uri[:-6]}/dest'
    elif bigquery_source_uri:
      unique_name = unique_name or _common.timestamped_unique_name()
      config.dest = f'{bigquery_source_uri}_dest_{unique_name}'

  return config


def find_afc_incompatible_tool_indexes(
    config: Optional[types.GenerateContentConfigOrDict] = None,
) -> list[int]:
  """Checks if the config contains any AFC incompatible tools.

  A `types.Tool` object that contains `function_declarations` is considered a
  non-AFC tool for this execution path.

  Args:
    config: The GenerateContentConfig to check for incompatible tools.

  Returns:
    A list of indexes of the incompatible tools in the config.
  """
  if not config:
    return []
  config_model = _create_generate_content_config_model(config)
  incompatible_tools_indexes: list[int] = []

  if not config_model or not config_model.tools:
    return incompatible_tools_indexes

  for index, tool in enumerate(config_model.tools):
    if isinstance(tool, types.Tool) and tool.function_declarations:
      incompatible_tools_indexes.append(index)

  return incompatible_tools_indexes


def get_function_map(
    config: Optional[types.GenerateContentConfigOrDict] = None,
    mcp_to_genai_tool_adapters: Optional[
        dict[str, McpToGenAiToolAdapter]
    ] = None,
    is_caller_method_async: bool = False,
) -> dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]]:
  """Returns a function map from the config."""
  function_map: dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]] = {}
  if not config:
    return function_map
  config_model = _create_generate_content_config_model(config)
  if config_model.tools:
    for tool in config_model.tools:
      if callable(tool):
        if inspect.iscoroutinefunction(tool) and not is_caller_method_async:
          raise errors.UnsupportedFunctionError(
              f'Function {tool.__name__} is a coroutine function, which is not'
              ' supported for automatic function calling. Please manually'
              f' invoke {tool.__name__} to get the function response.'
          )
        function_map[tool.__name__] = tool
  if mcp_to_genai_tool_adapters:
    if not is_caller_method_async:
      raise errors.UnsupportedFunctionError(
          'MCP tools are not supported in synchronous methods.'
      )
    for tool_name, _ in mcp_to_genai_tool_adapters.items():
      if function_map.get(tool_name):
        raise ValueError(
            f'Tool {tool_name} is already defined for the request.'
        )
    function_map.update(mcp_to_genai_tool_adapters)
  return function_map


def convert_number_values_for_dict_function_call_args(
    args: _common.StringDict,
) -> _common.StringDict:
  """Converts float values in dict with no decimal to integers."""
  return {
      key: convert_number_values_for_function_call_args(value)
      for key, value in args.items()
  }


def convert_number_values_for_function_call_args(
    args: Union[dict[str, object], list[object], object],
) -> Union[dict[str, object], list[object], object]:
  """Converts float values with no decimal to integers."""
  if isinstance(args, float) and args.is_integer():
    return int(args)
  if isinstance(args, dict):
    return {
        key: convert_number_values_for_function_call_args(value)
        for key, value in args.items()
    }
  if isinstance(args, list):
    return [
        convert_number_values_for_function_call_args(value) for value in args
    ]
  return args


def is_annotation_pydantic_model(annotation: Any) -> bool:
  try:
    return inspect.isclass(annotation) and issubclass(
        annotation, pydantic.BaseModel
    )
  # for python 3.10 and below, inspect.isclass(annotation) has inconsistent
  # results with versions above. for example, inspect.isclass(dict[str, int]) is
  # True in 3.10 and below but False in 3.11 and above.
  except TypeError:
    return False


def convert_if_exist_pydantic_model(
    value: Any, annotation: Any, param_name: str, func_name: str
) -> Any:
  if isinstance(value, dict) and is_annotation_pydantic_model(annotation):
    try:
      return annotation(**value)
    except pydantic.ValidationError as e:
      raise errors.UnknownFunctionCallArgumentError(
          f'Failed to parse parameter {param_name} for function'
          f' {func_name} from function call part because function call argument'
          f' value {value} is not compatible with parameter annotation'
          f' {annotation}, due to error {e}'
      )
  if isinstance(value, list) and get_origin(annotation) == list:
    item_type = get_args(annotation)[0]
    return [
        convert_if_exist_pydantic_model(item, item_type, param_name, func_name)
        for item in value
    ]
  if isinstance(value, dict) and get_origin(annotation) == dict:
    _, value_type = get_args(annotation)
    return {
        k: convert_if_exist_pydantic_model(v, value_type, param_name, func_name)
        for k, v in value.items()
    }
  # example 1: typing.Union[int, float]
  # example 2: int | float equivalent to UnionType[int, float]
  if get_origin(annotation) in (Union, UnionType):
    for arg in get_args(annotation):
      if (
          (get_args(arg) and get_origin(arg) is list)
          or isinstance(value, arg)
          or (isinstance(value, dict) and is_annotation_pydantic_model(arg))
      ):
        try:
          return convert_if_exist_pydantic_model(
              value, arg, param_name, func_name
          )
        # do not raise here because there could be multiple pydantic model types
        # in the union type.
        except pydantic.ValidationError:
          continue
    # if none of the union type is matched, raise error
    raise errors.UnknownFunctionCallArgumentError(
        f'Failed to parse parameter {param_name} for function'
        f' {func_name} from function call part because function call argument'
        f' value {value} cannot be converted to parameter annotation'
        f' {annotation}.'
    )
  # the only exception for value and annotation type to be different is int and
  # float. see convert_number_values_for_function_call_args function for context
  if isinstance(value, int) and annotation is float:
    return value
  if not isinstance(value, annotation):
    raise errors.UnknownFunctionCallArgumentError(
        f'Failed to parse parameter {param_name} for function {func_name} from'
        f' function call part because function call argument value {value} is'
        f' not compatible with parameter annotation {annotation}.'
    )
  return value


def convert_argument_from_function(
    args: _common.StringDict, function: Callable[..., Any]
) -> _common.StringDict:
  signature = inspect.signature(function)
  func_name = function.__name__
  converted_args = {}
  for param_name, param in signature.parameters.items():
    if param_name in args:
      converted_args[param_name] = convert_if_exist_pydantic_model(
          args[param_name],
          param.annotation,
          param_name,
          func_name,
      )
  return converted_args


def invoke_function_from_dict_args(
    args: _common.StringDict, function_to_invoke: Callable[..., Any]
) -> Any:
  converted_args = convert_argument_from_function(args, function_to_invoke)
  try:
    return function_to_invoke(**converted_args)
  except Exception as e:
    raise errors.FunctionInvocationError(
        f'Failed to invoke function {function_to_invoke.__name__} with'
        f' converted arguments {converted_args} from model returned function'
        f' call argument {args} because of error {e}'
    )


async def invoke_function_from_dict_args_async(
    args: _common.StringDict, function_to_invoke: Callable[..., Any]
) -> Any:
  converted_args = convert_argument_from_function(args, function_to_invoke)
  try:
    return await function_to_invoke(**converted_args)
  except Exception as e:
    raise errors.FunctionInvocationError(
        f'Failed to invoke function {function_to_invoke.__name__} with'
        f' converted arguments {converted_args} from model returned function'
        f' call argument {args} because of error {e}'
    )


def get_function_response_parts(
    response: types.GenerateContentResponse,
    function_map: dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]],
) -> list[types.Part]:
  """Returns the function response parts from the response."""
  func_response_parts = []
  if (
      response.candidates is not None
      and isinstance(response.candidates[0].content, types.Content)
      and response.candidates[0].content.parts is not None
  ):
    for part in response.candidates[0].content.parts:
      if not part.function_call:
        continue
      func_name = part.function_call.name
      if func_name is not None and part.function_call.args is not None:
        func = function_map[func_name]
        args = convert_number_values_for_dict_function_call_args(
            part.function_call.args
        )
        func_response: _common.StringDict
        try:
          if not isinstance(func, McpToGenAiToolAdapter):
            func_response = {
                'result': invoke_function_from_dict_args(args, func)
            }
        except Exception as e:  # pylint: disable=broad-except
          func_response = {'error': str(e)}
        func_response_part = types.Part.from_function_response(
            name=func_name, response=func_response
        )
        func_response_parts.append(func_response_part)
  return func_response_parts


async def get_function_response_parts_async(
    response: types.GenerateContentResponse,
    function_map: dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]],
) -> list[types.Part]:
  """Returns the function response parts from the response."""
  func_response_parts = []
  if (
      response.candidates is not None
      and isinstance(response.candidates[0].content, types.Content)
      and response.candidates[0].content.parts is not None
  ):
    for part in response.candidates[0].content.parts:
      if not part.function_call:
        continue
      func_name = part.function_call.name
      if func_name is not None and part.function_call.args is not None:
        func = function_map[func_name]
        args = convert_number_values_for_dict_function_call_args(
            part.function_call.args
        )
        func_response: _common.StringDict
        try:
          if isinstance(func, McpToGenAiToolAdapter):
            mcp_tool_response = await func.call_tool(
                types.FunctionCall(name=func_name, args=args)
            )
            if mcp_tool_response.isError:
              func_response = {'error': mcp_tool_response}
            else:
              func_response = {'result': mcp_tool_response}
          elif inspect.iscoroutinefunction(func):
            func_response = {
                'result': await invoke_function_from_dict_args_async(args, func)
            }
          else:
            func_response = {
                'result': await asyncio.to_thread(
                    invoke_function_from_dict_args, args, func
                )
            }
        except Exception as e:  # pylint: disable=broad-except
          func_response = {'error': str(e)}
        func_response_part = types.Part.from_function_response(
            name=func_name, response=func_response
        )
        func_response_parts.append(func_response_part)
  return func_response_parts


def should_disable_afc(
    config: Optional[types.GenerateContentConfigOrDict] = None,
) -> bool:
  """Returns whether automatic function calling is enabled."""
  if not config:
    return False
  config_model = _create_generate_content_config_model(config)
  # If max_remote_calls is less or equal to 0, warn and disable AFC.
  if (
      config_model
      and config_model.automatic_function_calling
      and config_model.automatic_function_calling.maximum_remote_calls
      is not None
      and int(config_model.automatic_function_calling.maximum_remote_calls) <= 0
  ):
    logger.warning(
        'max_remote_calls in automatic_function_calling_config'
        f' {config_model.automatic_function_calling.maximum_remote_calls} is'
        ' less than or equal to 0. Disabling automatic function calling.'
        ' Please set max_remote_calls to a positive integer.'
    )
    return True

  # Default to enable AFC if not specified.
  if (
      not config_model.automatic_function_calling
      or config_model.automatic_function_calling.disable is None
  ):
    return False

  if (
      config_model.automatic_function_calling.disable
      and config_model.automatic_function_calling.maximum_remote_calls
      is not None
      # exclude the case where max_remote_calls is set to 10 by default.
      and 'maximum_remote_calls'
      in config_model.automatic_function_calling.model_fields_set
      and int(config_model.automatic_function_calling.maximum_remote_calls) > 0
  ):
    logger.warning(
        '`automatic_function_calling.disable` is set to `True`. And'
        ' `automatic_function_calling.maximum_remote_calls` is a'
        ' positive number'
        f' {config_model.automatic_function_calling.maximum_remote_calls}.'
        ' Disabling automatic function calling. If you want to enable'
        ' automatic function calling, please set'
        ' `automatic_function_calling.disable` to `False` or leave it unset,'
        ' and set `automatic_function_calling.maximum_remote_calls` to a'
        ' positive integer or leave'
        ' `automatic_function_calling.maximum_remote_calls` unset.'
    )

  return config_model.automatic_function_calling.disable


def get_max_remote_calls_afc(
    config: Optional[types.GenerateContentConfigOrDict] = None,
) -> int:
  if not config:
    return _DEFAULT_MAX_REMOTE_CALLS_AFC
  """Returns the remaining remote calls for automatic function calling."""
  if should_disable_afc(config):
    raise ValueError(
        'automatic function calling is not enabled, but SDK is trying to get'
        ' max remote calls.'
    )
  config_model = _create_generate_content_config_model(config)
  if (
      not config_model.automatic_function_calling
      or config_model.automatic_function_calling.maximum_remote_calls is None
  ):
    return _DEFAULT_MAX_REMOTE_CALLS_AFC
  return int(config_model.automatic_function_calling.maximum_remote_calls)


def raise_error_for_afc_incompatible_config(config: Optional[types.GenerateContentConfig]
) -> None:
  """Raises an error if the config is not compatible with AFC."""
  if (
      not config
      or not config.tool_config
      or not config.tool_config.function_calling_config
  ):
    return
  afc_config = config.automatic_function_calling
  disable_afc_config = afc_config.disable if afc_config else False
  stream_function_call = (
      config.tool_config.function_calling_config.stream_function_call_arguments
  )

  if stream_function_call and not disable_afc_config:
    raise ValueError(
        'Running in streaming mode with stream_function_call_arguments'
        ' enabled, this feature is not compatible with automatic function'
        ' calling (AFC). Please set config.automatic_function_calling.disable'
        ' to True to disable AFC or leave config.tool_config.'
        ' function_calling_config.stream_function_call_arguments to be empty'
        ' or set to False to disable streaming function call arguments.'
    )

def should_append_afc_history(
    config: Optional[types.GenerateContentConfigOrDict] = None,
) -> bool:
  if not config:
    return True
  config_model = _create_generate_content_config_model(config)
  if not config_model.automatic_function_calling:
    return True
  return not config_model.automatic_function_calling.ignore_call_history


def parse_config_for_mcp_usage(
    config: Optional[types.GenerateContentConfigOrDict] = None,
) -> Optional[types.GenerateContentConfig]:
  """Returns a parsed config with an appended MCP header if MCP tools or sessions are used."""
  if not config:
    return None
  config_model = _create_generate_content_config_model(config)
  # Create a copy of the config model with the tools field cleared since some
  # tools may not be pickleable.
  config_model_copy = config_model.model_copy(update={'tools': None})
  config_model_copy.tools = config_model.tools
  if config_model.tools and _mcp_utils.has_mcp_tool_usage(config_model.tools):
    if config_model_copy.http_options is None:
      config_model_copy.http_options = types.HttpOptions(headers={})
    if config_model_copy.http_options.headers is None:
      config_model_copy.http_options.headers = {}
    _mcp_utils.set_mcp_usage_header(config_model_copy.http_options.headers)

  return config_model_copy


async def parse_config_for_mcp_sessions(
    config: Optional[types.GenerateContentConfigOrDict] = None,
) -> tuple[
    Optional[types.GenerateContentConfig],
    dict[str, McpToGenAiToolAdapter],
]:
  """Returns a parsed config with MCP sessions converted to GenAI tools.

  Also returns a map of MCP tools to GenAI tool adapters to be used for AFC.
  """
  mcp_to_genai_tool_adapters: dict[str, McpToGenAiToolAdapter] = {}
  parsed_config = parse_config_for_mcp_usage(config)
  if not parsed_config:
    return None, mcp_to_genai_tool_adapters
  # Create a copy of the config model with the tools field cleared as they will
  # be replaced with the MCP tools converted to GenAI tools.
  parsed_config_copy = parsed_config.model_copy(update={'tools': None})
  if parsed_config.tools:
    parsed_config_copy.tools = []
    for tool in parsed_config.tools:
      if McpClientSession is not None and isinstance(tool, McpClientSession):
        mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
            tool, await tool.list_tools()
        )
        # Extend the config with the MCP session tools converted to GenAI tools.
        parsed_config_copy.tools.extend(mcp_to_genai_tool_adapter.tools)
        for genai_tool in mcp_to_genai_tool_adapter.tools:
          if genai_tool.function_declarations:
            for function_declaration in genai_tool.function_declarations:
              if function_declaration.name:
                if mcp_to_genai_tool_adapters.get(function_declaration.name):
                  raise ValueError(
                      f'Tool {function_declaration.name} is already defined for'
                      ' the request.'
                  )
                mcp_to_genai_tool_adapters[function_declaration.name] = (
                    mcp_to_genai_tool_adapter
                )
      else:
        parsed_config_copy.tools.append(tool)

  return parsed_config_copy, mcp_to_genai_tool_adapters


def append_chunk_contents(
    contents: Union[types.ContentListUnion, types.ContentListUnionDict],
    chunk: types.GenerateContentResponse,
) -> Union[types.ContentListUnion, types.ContentListUnionDict]:
  """Appends the contents of the chunk to the contents list and returns it."""
  if chunk is not None and chunk.candidates is not None:
    chunk_content = chunk.candidates[0].content
    contents = t.t_contents(contents)  # type: ignore[assignment]
    if isinstance(contents, list) and chunk_content is not None:
      contents.append(chunk_content)  # type: ignore[arg-type]
  return contents


def prepare_resumable_upload(
    file: Union[str, os.PathLike[str], io.IOBase],
    user_http_options: Optional[types.HttpOptionsOrDict] = None,
    user_mime_type: Optional[str] = None,
) -> tuple[
    types.HttpOptions,
    int,
    str,
]:
  """Prepares the HTTP options, file bytes size and mime type for a resumable upload.

  This function inspects a file (from a path or an in-memory object) to
  determine its size and MIME type. It then constructs the necessary HTTP
  headers and options required to initiate a resumable upload session.
  """
  size_bytes = None
  mime_type = user_mime_type
  if isinstance(file, io.IOBase):
    if mime_type is None:
      raise ValueError(
          'Unknown mime type: Could not determine the mimetype for your'
          ' file\n please set the `mime_type` argument'
      )
    if hasattr(file, 'mode'):
      if 'b' not in file.mode:
        raise ValueError('The file must be opened in binary mode.')
    offset = file.tell()
    file.seek(0, os.SEEK_END)
    size_bytes = file.tell() - offset
    file.seek(offset, os.SEEK_SET)
  else:
    fs_path = os.fspath(file)
    if not fs_path or not os.path.isfile(fs_path):
      raise FileNotFoundError(f'{file} is not a valid file path.')
    size_bytes = os.path.getsize(fs_path)
    if mime_type is None:
      mime_type, _ = mimetypes.guess_type(fs_path)
    if mime_type is None:
      raise ValueError(
          'Unknown mime type: Could not determine the mimetype for your'
          ' file\n    please set the `mime_type` argument'
      )
  http_options: types.HttpOptions
  if user_http_options:
    if isinstance(user_http_options, dict):
      user_http_options = types.HttpOptions(**user_http_options)
    http_options = user_http_options
    http_options.api_version = ''
    http_options.headers = {
        'Content-Type': 'application/json',
        'X-Goog-Upload-Protocol': 'resumable',
        'X-Goog-Upload-Command': 'start',
        'X-Goog-Upload-Header-Content-Length': f'{size_bytes}',
        'X-Goog-Upload-Header-Content-Type': f'{mime_type}',
    }
  else:
    http_options = types.HttpOptions(
        api_version='',
        headers={
            'Content-Type': 'application/json',
            'X-Goog-Upload-Protocol': 'resumable',
            'X-Goog-Upload-Command': 'start',
            'X-Goog-Upload-Header-Content-Length': f'{size_bytes}',
            'X-Goog-Upload-Header-Content-Type': f'{mime_type}',
        },
    )
  if isinstance(file, (str, os.PathLike)):
    if http_options.headers is None:
        http_options.headers = {}
    http_options.headers['X-Goog-Upload-File-Name'] = os.path.basename(file)
  return http_options, size_bytes, mime_type