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

# Code generated by the Google Gen AI SDK generator DO NOT EDIT.

from functools import partial
import json
import logging
from typing import Any, Optional, Union
from urllib.parse import urlencode

from . import _api_module
from . import _common
from . import types
from ._common import get_value_by_path as getv
from ._common import set_value_by_path as setv
from .pagers import AsyncPager, Pager

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


def _DeleteDocumentConfig_to_mldev(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
  to_object: dict[str, Any] = {}

  if getv(from_object, ['force']) is not None:
    setv(parent_object, ['_query', 'force'], getv(from_object, ['force']))

  return to_object


def _DeleteDocumentParameters_to_mldev(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
  to_object: dict[str, Any] = {}
  if getv(from_object, ['name']) is not None:
    setv(to_object, ['_url', 'name'], getv(from_object, ['name']))

  if getv(from_object, ['config']) is not None:
    _DeleteDocumentConfig_to_mldev(getv(from_object, ['config']), to_object)

  return to_object


def _GetDocumentParameters_to_mldev(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
  to_object: dict[str, Any] = {}
  if getv(from_object, ['name']) is not None:
    setv(to_object, ['_url', 'name'], getv(from_object, ['name']))

  return to_object


def _ListDocumentsConfig_to_mldev(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
  to_object: dict[str, Any] = {}

  if getv(from_object, ['page_size']) is not None:
    setv(
        parent_object, ['_query', 'pageSize'], getv(from_object, ['page_size'])
    )

  if getv(from_object, ['page_token']) is not None:
    setv(
        parent_object,
        ['_query', 'pageToken'],
        getv(from_object, ['page_token']),
    )

  return to_object


def _ListDocumentsParameters_to_mldev(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
  to_object: dict[str, Any] = {}
  if getv(from_object, ['parent']) is not None:
    setv(to_object, ['_url', 'parent'], getv(from_object, ['parent']))

  if getv(from_object, ['config']) is not None:
    _ListDocumentsConfig_to_mldev(getv(from_object, ['config']), to_object)

  return to_object


def _ListDocumentsResponse_from_mldev(
    from_object: Union[dict[str, Any], object],
    parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
  to_object: dict[str, Any] = {}
  if getv(from_object, ['sdkHttpResponse']) is not None:
    setv(
        to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
    )

  if getv(from_object, ['nextPageToken']) is not None:
    setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))

  if getv(from_object, ['documents']) is not None:
    setv(
        to_object,
        ['documents'],
        [item for item in getv(from_object, ['documents'])],
    )

  return to_object


class Documents(_api_module.BaseModule):

  def get(
      self, *, name: str, config: Optional[types.GetDocumentConfigOrDict] = None
  ) -> types.Document:
    """Gets metadata about a Document.

    Args:
      name (str): The resource name of the Document.
        Example: ragStores/rag-store-foo/documents/documents-bar
      config (GetDocumentConfig | None): Optional parameters for the request.

    Returns:
      The Document.
    """

    parameter_model = types._GetDocumentParameters(
        name=name,
        config=config,
    )

    request_url_dict: Optional[dict[str, str]]
    if self._api_client.vertexai:
      raise ValueError(
          'This method is only supported in Gemini Developer API mode, not in'
          ' Gemini Enterprise Agent Platform mode.'
      )
    else:
      request_dict = _GetDocumentParameters_to_mldev(parameter_model)
      request_url_dict = request_dict.get('_url')
      if request_url_dict:
        path = '{name}'.format_map(request_url_dict)
      else:
        path = '{name}'

    query_params = request_dict.get('_query')
    if query_params:
      path = f'{path}?{urlencode(query_params)}'
    # TODO: remove the hack that pops config.
    request_dict.pop('config', None)

    http_options: Optional[types.HttpOptions] = None
    if (
        parameter_model.config is not None
        and parameter_model.config.http_options is not None
    ):
      http_options = parameter_model.config.http_options

    request_dict = _common.convert_to_dict(request_dict)
    request_dict = _common.encode_unserializable_types(request_dict)

    response = self._api_client.request('get', path, request_dict, http_options)

    response_dict = {} if not response.body else json.loads(response.body)

    return_value = types.Document._from_response(
        response=response_dict,
        kwargs={
            'config': {
                'response_schema': getattr(
                    parameter_model.config, 'response_schema', None
                ),
                'response_json_schema': getattr(
                    parameter_model.config, 'response_json_schema', None
                ),
                'include_all_fields': getattr(
                    parameter_model.config, 'include_all_fields', None
                ),
            }
        }
        if getattr(parameter_model, 'config', None)
        else {},
    )

    self._api_client._verify_response(return_value)
    return return_value

  def delete(
      self,
      *,
      name: str,
      config: Optional[types.DeleteDocumentConfigOrDict] = None,
  ) -> None:
    """Deletes a Document.

    Args:
      name (str): The resource name of the Document.
        Example: ragStores/rag-store-foo/documents/documents-bar
      config (DeleteDocumentConfig | None): Optional parameters for the request.

    Returns:
      None
    """

    parameter_model = types._DeleteDocumentParameters(
        name=name,
        config=config,
    )

    request_url_dict: Optional[dict[str, str]]
    if self._api_client.vertexai:
      raise ValueError(
          'This method is only supported in Gemini Developer API mode, not in'
          ' Gemini Enterprise Agent Platform mode.'
      )
    else:
      request_dict = _DeleteDocumentParameters_to_mldev(parameter_model)
      request_url_dict = request_dict.get('_url')
      if request_url_dict:
        path = '{name}'.format_map(request_url_dict)
      else:
        path = '{name}'

    query_params = request_dict.get('_query')
    if query_params:
      path = f'{path}?{urlencode(query_params)}'
    # TODO: remove the hack that pops config.
    request_dict.pop('config', None)

    http_options: Optional[types.HttpOptions] = None
    if (
        parameter_model.config is not None
        and parameter_model.config.http_options is not None
    ):
      http_options = parameter_model.config.http_options

    request_dict = _common.convert_to_dict(request_dict)
    request_dict = _common.encode_unserializable_types(request_dict)

    self._api_client.request('delete', path, request_dict, http_options)

  def _list(
      self,
      *,
      parent: str,
      config: Optional[types.ListDocumentsConfigOrDict] = None,
  ) -> types.ListDocumentsResponse:
    parameter_model = types._ListDocumentsParameters(
        parent=parent,
        config=config,
    )

    request_url_dict: Optional[dict[str, str]]
    if self._api_client.vertexai:
      raise ValueError(
          'This method is only supported in Gemini Developer API mode, not in'
          ' Gemini Enterprise Agent Platform mode.'
      )
    else:
      request_dict = _ListDocumentsParameters_to_mldev(parameter_model)
      request_url_dict = request_dict.get('_url')
      if request_url_dict:
        path = '{parent}/documents'.format_map(request_url_dict)
      else:
        path = '{parent}/documents'

    query_params = request_dict.get('_query')
    if query_params:
      path = f'{path}?{urlencode(query_params)}'
    # TODO: remove the hack that pops config.
    request_dict.pop('config', None)

    http_options: Optional[types.HttpOptions] = None
    if (
        parameter_model.config is not None
        and parameter_model.config.http_options is not None
    ):
      http_options = parameter_model.config.http_options

    request_dict = _common.convert_to_dict(request_dict)
    request_dict = _common.encode_unserializable_types(request_dict)

    response = self._api_client.request('get', path, request_dict, http_options)

    response_dict = {} if not response.body else json.loads(response.body)

    if not self._api_client.vertexai:
      response_dict = _ListDocumentsResponse_from_mldev(response_dict)

    return_value = types.ListDocumentsResponse._from_response(
        response=response_dict,
        kwargs={
            'config': {
                'response_schema': getattr(
                    parameter_model.config, 'response_schema', None
                ),
                'response_json_schema': getattr(
                    parameter_model.config, 'response_json_schema', None
                ),
                'include_all_fields': getattr(
                    parameter_model.config, 'include_all_fields', None
                ),
            }
        }
        if getattr(parameter_model, 'config', None)
        else {},
    )

    self._api_client._verify_response(return_value)
    return return_value

  def list(
      self,
      *,
      parent: str,
      config: Optional[types.ListDocumentsConfigOrDict] = None,
  ) -> Pager[types.Document]:
    """Lists documents.

    Args:
      parent (str): The name of the RagStore containing the Documents.
      config (ListDocumentsConfig): Optional configuration for the list request.

    Returns:
      A Pager object that contains one page of documents. When iterating over
      the pager, it automatically fetches the next page if there are more.
    Usage:
    .. code-block:: python
      for document in client.documents.list(parent='rag_store_name'):
        print(f"document: {document.name} - {document.display_name}")
    """

    list_request = partial(self._list, parent=parent)
    return Pager(
        'documents',
        list_request,
        self._list(parent=parent, config=config),
        config,
    )


class AsyncDocuments(_api_module.BaseModule):

  async def get(
      self, *, name: str, config: Optional[types.GetDocumentConfigOrDict] = None
  ) -> types.Document:
    """Gets metadata about a Document.

    Args:
      name (str): The resource name of the Document.
        Example: ragStores/rag-store-foo/documents/documents-bar
      config (GetDocumentConfig | None): Optional parameters for the request.

    Returns:
      The Document.
    """

    parameter_model = types._GetDocumentParameters(
        name=name,
        config=config,
    )

    request_url_dict: Optional[dict[str, str]]
    if self._api_client.vertexai:
      raise ValueError(
          'This method is only supported in Gemini Developer API mode, not in'
          ' Gemini Enterprise Agent Platform mode.'
      )
    else:
      request_dict = _GetDocumentParameters_to_mldev(parameter_model)
      request_url_dict = request_dict.get('_url')
      if request_url_dict:
        path = '{name}'.format_map(request_url_dict)
      else:
        path = '{name}'

    query_params = request_dict.get('_query')
    if query_params:
      path = f'{path}?{urlencode(query_params)}'
    # TODO: remove the hack that pops config.
    request_dict.pop('config', None)

    http_options: Optional[types.HttpOptions] = None
    if (
        parameter_model.config is not None
        and parameter_model.config.http_options is not None
    ):
      http_options = parameter_model.config.http_options

    request_dict = _common.convert_to_dict(request_dict)
    request_dict = _common.encode_unserializable_types(request_dict)

    response = await self._api_client.async_request(
        'get', path, request_dict, http_options
    )

    response_dict = {} if not response.body else json.loads(response.body)

    return_value = types.Document._from_response(
        response=response_dict,
        kwargs={
            'config': {
                'response_schema': getattr(
                    parameter_model.config, 'response_schema', None
                ),
                'response_json_schema': getattr(
                    parameter_model.config, 'response_json_schema', None
                ),
                'include_all_fields': getattr(
                    parameter_model.config, 'include_all_fields', None
                ),
            }
        }
        if getattr(parameter_model, 'config', None)
        else {},
    )

    self._api_client._verify_response(return_value)
    return return_value

  async def delete(
      self,
      *,
      name: str,
      config: Optional[types.DeleteDocumentConfigOrDict] = None,
  ) -> None:
    """Deletes a Document.

    Args:
      name (str): The resource name of the Document.
        Example: ragStores/rag-store-foo/documents/documents-bar
      config (DeleteDocumentConfig | None): Optional parameters for the request.

    Returns:
      None
    """

    parameter_model = types._DeleteDocumentParameters(
        name=name,
        config=config,
    )

    request_url_dict: Optional[dict[str, str]]
    if self._api_client.vertexai:
      raise ValueError(
          'This method is only supported in Gemini Developer API mode, not in'
          ' Gemini Enterprise Agent Platform mode.'
      )
    else:
      request_dict = _DeleteDocumentParameters_to_mldev(parameter_model)
      request_url_dict = request_dict.get('_url')
      if request_url_dict:
        path = '{name}'.format_map(request_url_dict)
      else:
        path = '{name}'

    query_params = request_dict.get('_query')
    if query_params:
      path = f'{path}?{urlencode(query_params)}'
    # TODO: remove the hack that pops config.
    request_dict.pop('config', None)

    http_options: Optional[types.HttpOptions] = None
    if (
        parameter_model.config is not None
        and parameter_model.config.http_options is not None
    ):
      http_options = parameter_model.config.http_options

    request_dict = _common.convert_to_dict(request_dict)
    request_dict = _common.encode_unserializable_types(request_dict)

    await self._api_client.async_request(
        'delete', path, request_dict, http_options
    )

  async def _list(
      self,
      *,
      parent: str,
      config: Optional[types.ListDocumentsConfigOrDict] = None,
  ) -> types.ListDocumentsResponse:
    parameter_model = types._ListDocumentsParameters(
        parent=parent,
        config=config,
    )

    request_url_dict: Optional[dict[str, str]]
    if self._api_client.vertexai:
      raise ValueError(
          'This method is only supported in Gemini Developer API mode, not in'
          ' Gemini Enterprise Agent Platform mode.'
      )
    else:
      request_dict = _ListDocumentsParameters_to_mldev(parameter_model)
      request_url_dict = request_dict.get('_url')
      if request_url_dict:
        path = '{parent}/documents'.format_map(request_url_dict)
      else:
        path = '{parent}/documents'

    query_params = request_dict.get('_query')
    if query_params:
      path = f'{path}?{urlencode(query_params)}'
    # TODO: remove the hack that pops config.
    request_dict.pop('config', None)

    http_options: Optional[types.HttpOptions] = None
    if (
        parameter_model.config is not None
        and parameter_model.config.http_options is not None
    ):
      http_options = parameter_model.config.http_options

    request_dict = _common.convert_to_dict(request_dict)
    request_dict = _common.encode_unserializable_types(request_dict)

    response = await self._api_client.async_request(
        'get', path, request_dict, http_options
    )

    response_dict = {} if not response.body else json.loads(response.body)

    if not self._api_client.vertexai:
      response_dict = _ListDocumentsResponse_from_mldev(response_dict)

    return_value = types.ListDocumentsResponse._from_response(
        response=response_dict,
        kwargs={
            'config': {
                'response_schema': getattr(
                    parameter_model.config, 'response_schema', None
                ),
                'response_json_schema': getattr(
                    parameter_model.config, 'response_json_schema', None
                ),
                'include_all_fields': getattr(
                    parameter_model.config, 'include_all_fields', None
                ),
            }
        }
        if getattr(parameter_model, 'config', None)
        else {},
    )

    self._api_client._verify_response(return_value)
    return return_value

  async def list(
      self,
      *,
      parent: str,
      config: Optional[types.ListDocumentsConfigOrDict] = None,
  ) -> AsyncPager[types.Document]:
    """Lists documents asynchronously.

    Args:
      parent (str): The name of the RagStore containing the Documents.
      config (ListDocumentsConfig): Optional configuration for the list request.

    Returns:
      A Pager object that contains one page of documents. When iterating over
      the pager, it automatically fetches the next page if there are more.
    Usage:
    .. code-block:: python
      async for document in await
      client.aio.documents.list(parent='rag_store_name'):
        print(f"document: {document.name} - {document.display_name}")
    """

    list_request = partial(self._list, parent=parent)
    return AsyncPager(
        'documents',
        list_request,
        await self._list(parent=parent, config=config),
        config,
    )
