# 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 __future__ import annotations

import asyncio
from collections.abc import Mapping
from collections.abc import Sequence
import datetime
from functools import lru_cache
import logging
from typing import Optional
from typing import TYPE_CHECKING

from google.genai import types
from typing_extensions import override

from ..utils.vertex_ai_utils import get_express_mode_api_key
from .base_memory_service import BaseMemoryService
from .base_memory_service import SearchMemoryResponse
from .memory_entry import MemoryEntry

if TYPE_CHECKING:
  import vertexai
  from vertexai import types as vertex_types

  from ..events.event import Event
  from ..sessions.session import Session

logger = logging.getLogger('google_adk.' + __name__)

# Strong references to fire-and-forget tasks to prevent garbage collection.
# See https://docs.python.org/3/library/asyncio-task.html#creating-tasks
_background_tasks: set[asyncio.Task] = set()

_GENERATE_MEMORIES_CONFIG_FALLBACK_KEYS = frozenset({
    'disable_consolidation',
    'disable_memory_revisions',
    'http_options',
    'metadata',
    'metadata_merge_strategy',
    'revision_expire_time',
    'revision_labels',
    'revision_ttl',
    'ttl',
    'wait_for_completion',
})

_CREATE_MEMORY_CONFIG_FALLBACK_KEYS = frozenset({
    'description',
    'disable_memory_revisions',
    'display_name',
    'expire_time',
    'http_options',
    'metadata',
    'revision_labels',
    'revision_expire_time',
    'revision_ttl',
    'topics',
    'ttl',
    'wait_for_completion',
})

_INGEST_EVENTS_CONFIG_FALLBACK_KEYS = frozenset({
    'force_flush',
    'generation_trigger_config',
    'stream_id',
})

_ENABLE_CONSOLIDATION_KEY = 'enable_consolidation'


def _should_use_generate_memories(
    custom_metadata: Mapping[str, object] | None,
) -> bool:
  """Returns True if custom_metadata contains keys only GenerateMemories supports.

  If any key in custom_metadata is recognized by GenerateMemories but NOT by
  IngestEvents, the generate_memories API path is used.  Otherwise
  ingest_events is the default.
  """
  if not custom_metadata:
    return False
  ingest_keys = _INGEST_EVENTS_CONFIG_FALLBACK_KEYS
  generate_keys = _GENERATE_MEMORIES_CONFIG_FALLBACK_KEYS
  for key in custom_metadata:
    if key not in ingest_keys and key in generate_keys:
      return True
  return False


# Vertex docs for GenerateMemoriesRequest.DirectMemoriesSource allow
# at most 5 direct_memories per request.
_MAX_DIRECT_MEMORIES_PER_GENERATE_CALL = 5


def _supports_generate_memories_metadata() -> bool:
  """Returns whether installed Vertex SDK supports config.metadata."""
  try:
    from vertexai import types as vertex_types
  except ImportError:
    return False
  return (
      'metadata' in vertex_types.GenerateAgentEngineMemoriesConfig.model_fields
  )


def _supports_create_memory_metadata() -> bool:
  """Returns whether installed Vertex SDK supports create config.metadata."""
  try:
    from vertexai import types as vertex_types
  except ImportError:
    return False
  return 'metadata' in vertex_types.AgentEngineMemoryConfig.model_fields


@lru_cache(maxsize=1)
def _get_generate_memories_config_keys() -> frozenset[str]:
  """Returns supported config keys for memories.generate.

  Uses SDK runtime model fields when available and falls back to a static
  allowlist to preserve compatibility when introspection is unavailable.
  """
  try:
    from vertexai import types as vertex_types
  except ImportError:
    return _GENERATE_MEMORIES_CONFIG_FALLBACK_KEYS

  try:
    model_fields = vertex_types.GenerateAgentEngineMemoriesConfig.model_fields
  except AttributeError:
    return _GENERATE_MEMORIES_CONFIG_FALLBACK_KEYS

  if not isinstance(model_fields, Mapping):
    return _GENERATE_MEMORIES_CONFIG_FALLBACK_KEYS
  return frozenset(model_fields.keys())


@lru_cache(maxsize=1)
def _get_create_memory_config_keys() -> frozenset[str]:
  """Returns supported config keys for memories.create.

  Uses SDK runtime model fields when available and falls back to a static
  allowlist to preserve compatibility when introspection is unavailable.
  """
  try:
    from vertexai import types as vertex_types
  except ImportError:
    return _CREATE_MEMORY_CONFIG_FALLBACK_KEYS

  try:
    model_fields = vertex_types.AgentEngineMemoryConfig.model_fields
  except AttributeError:
    return _CREATE_MEMORY_CONFIG_FALLBACK_KEYS

  if not isinstance(model_fields, Mapping):
    return _CREATE_MEMORY_CONFIG_FALLBACK_KEYS
  return frozenset(model_fields.keys())


class VertexAiMemoryBankService(BaseMemoryService):
  """Implementation of the BaseMemoryService using Vertex AI Memory Bank."""

  def __init__(
      self,
      project: Optional[str] = None,
      location: Optional[str] = None,
      agent_engine_id: Optional[str] = None,
      *,
      express_mode_api_key: Optional[str] = None,
  ):
    """Initializes a VertexAiMemoryBankService.

    Args:
      project: The project ID of the Memory Bank to use.
      location: The location of the Memory Bank to use.
      agent_engine_id: The ID of the agent engine to use for the Memory Bank,
        e.g. '456' in
        'projects/my-project/locations/us-central1/reasoningEngines/456'. To
        extract from api_resource.name, use:
        ``agent_engine.api_resource.name.split('/')[-1]``
      express_mode_api_key: The API key to use for Express Mode. If not
        provided, the API key from the GOOGLE_API_KEY environment variable will
        be used. It will only be used if GOOGLE_GENAI_USE_ENTERPRISE is true. Do
        not use Google AI Studio API key for this field. For more details, visit
        https://cloud.google.com/vertex-ai/generative-ai/docs/start/express-mode/overview
    """
    if not agent_engine_id:
      raise ValueError(
          'agent_engine_id is required for VertexAiMemoryBankService.'
      )

    try:
      import vertexai  # noqa: F401
    except ImportError as e:
      from ..utils._dependency import missing_extra

      raise missing_extra('google-cloud-aiplatform', 'gcp') from e

    self._project = project
    self._location = location
    self._agent_engine_id = agent_engine_id
    self._express_mode_api_key = get_express_mode_api_key(
        project, location, express_mode_api_key
    )

    if agent_engine_id and '/' in agent_engine_id:
      logger.warning(
          "agent_engine_id appears to be a full resource path: '%s'. "
          "Expected just the ID (e.g., '456'). "
          "Extract the ID using: agent_engine.api_resource.name.split('/')[-1]",
          agent_engine_id,
      )

  @override
  async def add_session_to_memory(self, session: Session) -> None:
    await self._add_events_to_memory_from_events(
        app_name=session.app_name,
        user_id=session.user_id,
        events_to_process=session.events,
    )

  @override
  async def add_events_to_memory(
      self,
      *,
      app_name: str,
      user_id: str,
      events: Sequence[Event],
      session_id: str | None = None,
      custom_metadata: Mapping[str, object] | None = None,
  ) -> None:
    """Adds events to Vertex AI Memory Bank.

    Uses ``memories.ingest_events`` by default. If ``custom_metadata`` contains
    keys supported only by ``memories.generate`` (e.g. ``ttl``,
    ``revision_ttl``, ``metadata``, ``wait_for_completion``), the generate path
    is used instead.

    Args:
      app_name: The application name for memory scope.
      user_id: The user ID for memory scope.
      events: The events to process for memory generation.
      session_id: Optional session ID. Currently unused.
      custom_metadata: Optional service-specific metadata. Supported keys
        depend on the API path chosen:

        **IngestEvents keys** (default path):
          stream_id: Identifier for the event stream.
          force_flush: If True, forces flushing buffered events.
          generation_trigger_config: Configuration for triggering memory
            generation, e.g.
            ``{"generation_rule": {"idle_duration": "60s"}}``.

        **GenerateMemories keys** (used when any of these are present):
          ttl: Time-to-live for generated memories, e.g. ``"6000s"``.
          revision_ttl: Time-to-live for memory revisions.
          metadata: A mapping of custom metadata key-value pairs.
          wait_for_completion: Whether to wait for generation to complete.
          disable_consolidation: Disable memory consolidation.
          disable_memory_revisions: Disable memory revisions.
    """
    _ = session_id
    await self._add_events_to_memory_from_events(
        app_name=app_name,
        user_id=user_id,
        events_to_process=events,
        custom_metadata=custom_metadata,
    )

  @override
  async def add_memory(
      self,
      *,
      app_name: str,
      user_id: str,
      memories: Sequence[MemoryEntry],
      custom_metadata: Mapping[str, object] | None = None,
  ) -> None:
    """Adds explicit memory items using Vertex Memory Bank.

    By default, this writes directly via `memories.create`.
    If `custom_metadata["enable_consolidation"]` is set to True, this uses
    `memories.generate` with `direct_memories_source` so provided memories are
    consolidated server-side.
    """
    if _is_consolidation_enabled(custom_metadata):
      await self._add_memories_via_generate_direct_memories_source(
          app_name=app_name,
          user_id=user_id,
          memories=memories,
          custom_metadata=custom_metadata,
      )
      return

    await self._add_memories_via_create(
        app_name=app_name,
        user_id=user_id,
        memories=memories,
        custom_metadata=custom_metadata,
    )

  async def _add_events_to_memory_from_events(
      self,
      *,
      app_name: str,
      user_id: str,
      events_to_process: Sequence[Event],
      custom_metadata: Mapping[str, object] | None = None,
  ) -> None:
    # The generate_memories API is used only when custom_metadata contains
    # keys exclusive to GenerateMemories.  Otherwise, ingest_events is the
    # default path, as its behavior is consistent with GenerateMemories
    # (trigger immediately) and supports additional parameters like
    # generation_trigger_config.
    if _should_use_generate_memories(custom_metadata):
      import vertexai

      direct_events = []
      for event in events_to_process:
        if _should_filter_out_event(event.content):
          continue
        if event.content:
          direct_events.append(
              vertexai.types.GenerateMemoriesRequestDirectContentsSourceEvent(
                  content=event.content
              )
          )
      if direct_events:
        api_client = self._get_api_client()
        config = _build_generate_memories_config(custom_metadata)
        operation = await api_client.agent_engines.memories.generate(
            name='reasoningEngines/' + self._agent_engine_id,
            direct_contents_source=vertexai.types.GenerateMemoriesRequestDirectContentsSource(
                events=direct_events
            ),
            scope={
                'app_name': app_name,
                'user_id': user_id,
            },
            config=config,
        )
        logger.info('Generate memory response received.')
        logger.debug('Generate memory response: %s', operation)
      else:
        logger.info('No events to add to memory.')
      return

    await self._add_events_to_memory_via_ingest(
        app_name=app_name,
        user_id=user_id,
        events_to_process=events_to_process,
        custom_metadata=custom_metadata,
    )

  async def _add_events_to_memory_via_ingest(
      self,
      *,
      app_name: str,
      user_id: str,
      events_to_process: Sequence[Event],
      custom_metadata: Mapping[str, object] | None = None,
  ) -> None:
    """Adds events to Vertex AI Memory Bank via memories.ingest_events.

    Args:
      app_name: The application name for memory scope.
      user_id: The user ID for memory scope.
      events_to_process: The events to process for memory ingestion.
      custom_metadata: Optional service-specific metadata. Supported keys:
        stream_id: Identifier for the event stream.
        force_flush: If True, forces flushing buffered events (passed as
          part of the ingest_events config).
        generation_trigger_config: Configuration for triggering memory
          generation, e.g.
          ``{"generation_rule": {"idle_duration": "60s"}}``.
    """
    import vertexai

    direct_events = []
    for event in events_to_process:
      if _should_filter_out_event(event.content):
        continue
      if event.content:
        event_time = None
        if event.timestamp is not None:
          event_time = datetime.datetime.fromtimestamp(
              event.timestamp, tz=datetime.timezone.utc
          )
        direct_events.append(
            vertexai.types.IngestionDirectContentsSourceEvent(
                content=event.content,
                event_id=event.id,
                event_time=event_time,
            )
        )

    api_client = self._get_api_client()

    stream_id = custom_metadata.get('stream_id') if custom_metadata else None
    force_flush = (
        custom_metadata.get('force_flush') if custom_metadata else None
    )
    generation_trigger_config = (
        custom_metadata.get('generation_trigger_config')
        if custom_metadata
        else None
    )

    request_kwargs: dict[str, object] = {
        'name': 'reasoningEngines/' + self._agent_engine_id,
        'scope': {
            'app_name': app_name,
            'user_id': user_id,
        },
    }
    # No-events requests are valid for trigger config updates, but
    # won't trigger an events flush.
    if direct_events:
      request_kwargs['direct_contents_source'] = (
          vertexai.types.IngestionDirectContentsSource(events=direct_events)
      )
    if stream_id:
      request_kwargs['stream_id'] = stream_id
    # force_flush is part of the ingest_events config, not a
    # top-level request parameter.
    config: dict[str, object] = {}
    if force_flush is not None:
      config['force_flush'] = force_flush
    if config:
      request_kwargs['config'] = config
    if generation_trigger_config:
      request_kwargs['generation_trigger_config'] = generation_trigger_config

    # Fire the ingest request without blocking. IngestEvents latency
    # (~800ms to trigger) makes awaiting unnecessary outside debugging.
    task = asyncio.create_task(
        api_client.agent_engines.memories.ingest_events(**request_kwargs)
    )
    _background_tasks.add(task)
    task.add_done_callback(_background_tasks.discard)
    task.add_done_callback(_log_ingest_task_error)
    logger.info('Ingest events request triggered.')

  async def _add_memories_via_create(
      self,
      *,
      app_name: str,
      user_id: str,
      memories: Sequence[MemoryEntry],
      custom_metadata: Mapping[str, object] | None = None,
  ) -> None:
    """Adds direct memory items without server-side extraction."""
    normalized_memories = _normalize_memories_for_create(memories)
    api_client = self._get_api_client()
    for index, memory in enumerate(normalized_memories):
      memory_fact = _memory_entry_to_fact(memory, index=index)
      memory_metadata = _merge_custom_metadata_for_memory(
          custom_metadata=custom_metadata,
          memory=memory,
      )
      memory_revision_labels = _revision_labels_for_memory(memory)
      config = _build_create_memory_config(
          memory_metadata,
          memory_revision_labels=memory_revision_labels,
      )
      operation = await api_client.agent_engines.memories.create(
          name='reasoningEngines/' + self._agent_engine_id,
          fact=memory_fact,
          scope={
              'app_name': app_name,
              'user_id': user_id,
          },
          config=config,
      )
      logger.info('Create memory response received.')
      logger.debug('Create memory response: %s', operation)

  async def _add_memories_via_generate_direct_memories_source(
      self,
      *,
      app_name: str,
      user_id: str,
      memories: Sequence[MemoryEntry],
      custom_metadata: Mapping[str, object] | None = None,
  ) -> None:
    """Adds memories via generate API with direct_memories_source."""
    normalized_memories = _normalize_memories_for_create(memories)
    memory_texts = [
        _memory_entry_to_fact(m, index=i)
        for i, m in enumerate(normalized_memories)
    ]
    api_client = self._get_api_client()
    config = _build_generate_memories_config(custom_metadata)
    for memory_batch in _iter_memory_batches(memory_texts):
      operation = await api_client.agent_engines.memories.generate(
          name='reasoningEngines/' + self._agent_engine_id,
          direct_memories_source={
              'direct_memories': [
                  {'fact': memory_text} for memory_text in memory_batch
              ]
          },
          scope={
              'app_name': app_name,
              'user_id': user_id,
          },
          config=config,
      )
      logger.info('Generate direct memory response received.')
      logger.debug('Generate direct memory response: %s', operation)

  @override
  async def search_memory(self, *, app_name: str, user_id: str, query: str):
    api_client = self._get_api_client()
    retrieved_memories_iterator = (
        await api_client.agent_engines.memories.retrieve(
            name='reasoningEngines/' + self._agent_engine_id,
            scope={
                'app_name': app_name,
                'user_id': user_id,
            },
            similarity_search_params={
                'search_query': query,
            },
        )
    )

    logger.info('Search memory response received.')

    memory_events: list[MemoryEntry] = []
    try:
      async for retrieved_memory in retrieved_memories_iterator:
        try:
          memory = retrieved_memory.memory
          if memory is None:
            logger.warning('Skipping memory entry with missing memory object.')
            continue
          fact = memory.fact
          if not fact:
            logger.warning('Skipping memory entry with empty or missing fact.')
            continue
          update_time = memory.update_time
          memory_events.append(
              MemoryEntry(
                  author='user',
                  content=types.Content(
                      parts=[types.Part(text=fact)],
                      role='user',
                  ),
                  timestamp=update_time.isoformat() if update_time else None,
              )
          )
        except AttributeError:
          logger.warning(
              'Skipping malformed memory entry: %s', retrieved_memory
          )
    except Exception:
      logger.exception(
          'Error while iterating memory results. Returning %d partial results.',
          len(memory_events),
      )
    return SearchMemoryResponse(memories=memory_events)

  async def retrieve_profiles(
      self,
      *,
      app_name: str,
      user_id: str,
  ) -> list[vertex_types.MemoryProfile]:
    """Retrieves structured user profiles for the scope, one per schema.

    Profiles are a Vertex Memory Bank capability distinct from memory search:
    a scope-keyed lookup, not a semantic query.

    Args:
      app_name: The application name for the profile scope.
      user_id: The user ID for the profile scope.

    Returns:
      The structured profiles for the scope, one per registered schema.
    """
    api_client = self._get_api_client()
    response = await api_client.agent_engines.memories.retrieve_profiles(
        name='reasoningEngines/' + self._agent_engine_id,
        scope={
            'app_name': app_name,
            'user_id': user_id,
        },
    )
    profiles = list((response.profiles or {}).values())
    if profiles:
      logger.info('Retrieved %d memory profiles.', len(profiles))
    else:
      logger.info('Retrieved no memory profiles.')
    return profiles

  def _get_api_client(self) -> vertexai.AsyncClient:
    """Instantiates an API client for the given project and location.

    It needs to be instantiated inside each request so that the event loop
    management can be properly propagated.
    Returns:
      An async API client for the given project and location or express mode api
      key.
    """
    import vertexai

    if self._express_mode_api_key:
      return vertexai.Client(api_key=self._express_mode_api_key).aio
    return vertexai.Client(project=self._project, location=self._location).aio


def _log_ingest_task_error(task: asyncio.Task) -> None:
  """Logs errors from fire-and-forget ingest_events tasks."""
  if task.cancelled():
    return
  exception = task.exception()
  if exception:
    logger.error('Background ingest_events task failed: %s', exception)


def _should_filter_out_event(content: types.Content) -> bool:
  """Returns whether the event should be filtered out."""
  if not content or not content.parts:
    return True
  for part in content.parts:
    if (
        part.text
        or part.inline_data
        or part.file_data
        or part.function_call
        or part.function_response
        or part.executable_code
        or part.code_execution_result
        or part.tool_call
        or part.tool_response
    ):
      return False
  return True


def _build_generate_memories_config(
    custom_metadata: Mapping[str, object] | None,
) -> dict[str, object]:
  """Builds a valid memories.generate config from caller metadata."""
  config: dict[str, object] = {'wait_for_completion': False}
  supports_metadata = _supports_generate_memories_metadata()
  config_keys = _get_generate_memories_config_keys()
  if not custom_metadata:
    return config

  logger.debug('Memory generation metadata: %s', custom_metadata)

  metadata_by_key: dict[str, object] = {}
  for key, value in custom_metadata.items():
    if key == _ENABLE_CONSOLIDATION_KEY:
      continue
    if key == 'ttl':
      if value is None:
        continue
      if custom_metadata.get('revision_ttl') is None:
        config['revision_ttl'] = value
      continue
    if key == 'metadata':
      if value is None:
        continue
      if not supports_metadata:
        logger.warning(
            'Ignoring metadata because installed Vertex SDK does not support'
            ' config.metadata.'
        )
        continue
      if isinstance(value, Mapping):
        config['metadata'] = _build_vertex_metadata(value)
      else:
        logger.warning(
            'Ignoring metadata because custom_metadata["metadata"] is not a'
            ' mapping.'
        )
      continue
    if key in config_keys:
      if value is None:
        continue
      config[key] = value
    else:
      metadata_by_key[key] = value

  if not metadata_by_key:
    return config

  if not supports_metadata:
    logger.warning(
        'Ignoring custom metadata keys %s because installed Vertex SDK does '
        'not support config.metadata.',
        sorted(metadata_by_key.keys()),
    )
    return config

  existing_metadata = config.get('metadata')
  if existing_metadata is None:
    config['metadata'] = _build_vertex_metadata(metadata_by_key)
    return config

  if isinstance(existing_metadata, Mapping):
    merged_metadata = dict(existing_metadata)
    merged_metadata.update(_build_vertex_metadata(metadata_by_key))
    config['metadata'] = merged_metadata
    return config

  logger.warning(
      'Ignoring custom metadata keys %s because config.metadata is not a'
      ' mapping.',
      sorted(metadata_by_key.keys()),
  )
  return config


def _build_create_memory_config(
    custom_metadata: Mapping[str, object] | None,
    *,
    memory_revision_labels: Mapping[str, str] | None = None,
) -> dict[str, object]:
  """Builds a valid memories.create config from caller metadata."""
  config: dict[str, object] = {'wait_for_completion': False}
  supports_metadata = _supports_create_memory_metadata()
  config_keys = _get_create_memory_config_keys()
  supports_revision_labels = 'revision_labels' in config_keys

  if custom_metadata:
    logger.debug('Memory creation metadata: %s', custom_metadata)

  metadata_by_key: dict[str, object] = {}
  custom_revision_labels: dict[str, str] = {}
  for key, value in (custom_metadata or {}).items():
    if key == _ENABLE_CONSOLIDATION_KEY:
      continue
    if key == 'metadata':
      if value is None:
        continue
      if not supports_metadata:
        logger.warning(
            'Ignoring metadata because installed Vertex SDK does not support'
            ' create config.metadata.'
        )
        continue
      if isinstance(value, Mapping):
        config['metadata'] = _build_vertex_metadata(value)
      else:
        logger.warning(
            'Ignoring metadata because custom_metadata["metadata"] is not a'
            ' mapping.'
        )
      continue
    if key == 'revision_labels':
      if value is None:
        continue
      extracted_labels = _extract_revision_labels(
          value,
          source='custom_metadata["revision_labels"]',
      )
      if extracted_labels:
        custom_revision_labels.update(extracted_labels)
      continue
    if key in config_keys:
      if value is None:
        continue
      config[key] = value
    else:
      metadata_by_key[key] = value

  if metadata_by_key:
    if not supports_metadata:
      logger.warning(
          'Ignoring custom metadata keys %s because installed Vertex SDK does '
          'not support create config.metadata.',
          sorted(metadata_by_key.keys()),
      )
    else:
      existing_metadata = config.get('metadata')
      if existing_metadata is None:
        config['metadata'] = _build_vertex_metadata(metadata_by_key)
      elif isinstance(existing_metadata, Mapping):
        merged_metadata = dict(existing_metadata)
        merged_metadata.update(_build_vertex_metadata(metadata_by_key))
        config['metadata'] = merged_metadata
      else:
        logger.warning(
            'Ignoring custom metadata keys %s because config.metadata is not a'
            ' mapping.',
            sorted(metadata_by_key.keys()),
        )

  revision_labels = dict(custom_revision_labels)
  if memory_revision_labels:
    revision_labels.update(memory_revision_labels)
  if revision_labels:
    if supports_revision_labels:
      config['revision_labels'] = revision_labels
    else:
      logger.warning(
          'Ignoring revision labels %s because installed Vertex SDK does not '
          'support create config.revision_labels.',
          sorted(revision_labels.keys()),
      )
  return config


def _normalize_memories_for_create(
    memories: Sequence[MemoryEntry],
) -> list[MemoryEntry]:
  """Validates add_memory inputs."""
  if isinstance(memories, str):
    raise TypeError('memories must be a sequence of memory items.')
  if not isinstance(memories, Sequence):
    raise TypeError('memories must be a sequence of memory items.')

  validated_memories: list[MemoryEntry] = []
  for index, raw_memory in enumerate(memories):
    if not isinstance(raw_memory, MemoryEntry):
      raise TypeError(f'memories[{index}] must be a MemoryEntry.')
    validated_memories.append(raw_memory)
  if not validated_memories:
    raise ValueError('memories must contain at least one entry.')
  return validated_memories


def _memory_entry_to_fact(
    memory: MemoryEntry,
    *,
    index: int,
) -> str:
  """Builds a memories.create fact payload from MemoryEntry text content."""
  if _should_filter_out_event(memory.content):
    raise ValueError(f'memories[{index}] must include text.')

  text_parts: list[str] = []
  for part in memory.content.parts:
    if part.inline_data or part.file_data:
      raise ValueError(
          f'memories[{index}] must include text only; inline_data and '
          'file_data are not supported.'
      )

    if not part.text:
      continue
    stripped_text = part.text.strip()
    if stripped_text:
      text_parts.append(stripped_text)

  if not text_parts:
    raise ValueError(f'memories[{index}] must include non-whitespace text.')
  return '\n'.join(text_parts)


def _merge_custom_metadata_for_memory(
    *,
    custom_metadata: Mapping[str, object] | None,
    memory: MemoryEntry,
) -> Mapping[str, object] | None:
  """Merges write-level metadata with MemoryEntry metadata."""
  merged_metadata: dict[str, object] = {}

  if custom_metadata:
    merged_metadata.update(dict(custom_metadata))
  if memory.custom_metadata:
    merged_metadata.update(memory.custom_metadata)

  if not merged_metadata:
    return None
  return merged_metadata


def _revision_labels_for_memory(
    memory: MemoryEntry,
) -> Mapping[str, str] | None:
  """Builds revision labels from MemoryEntry revision metadata."""
  revision_labels: dict[str, str] = {}
  if memory.author is not None:
    revision_labels['author'] = memory.author
  if memory.timestamp is not None:
    revision_labels['timestamp'] = memory.timestamp

  if not revision_labels:
    return None
  return revision_labels


def _extract_revision_labels(
    value: object,
    *,
    source: str,
) -> Mapping[str, str] | None:
  """Extracts revision labels from config metadata."""
  if not isinstance(value, Mapping):
    logger.warning('Ignoring %s because it is not a mapping.', source)
    return None

  revision_labels: dict[str, str] = {}
  for key, label_value in value.items():
    if not isinstance(key, str):
      logger.warning(
          'Ignoring revision label with non-string key %r from %s.',
          key,
          source,
      )
      continue
    if not isinstance(label_value, str):
      logger.warning(
          'Ignoring revision label %s from %s because its value is not a '
          'string.',
          key,
          source,
      )
      continue
    revision_labels[key] = label_value

  if not revision_labels:
    return None
  return revision_labels


def _is_consolidation_enabled(
    custom_metadata: Mapping[str, object] | None,
) -> bool:
  """Returns whether direct memories should be consolidated via generate API."""
  if not custom_metadata:
    return False
  enable_consolidation = custom_metadata.get(_ENABLE_CONSOLIDATION_KEY)
  if enable_consolidation is None:
    return False
  if not isinstance(enable_consolidation, bool):
    raise TypeError(
        f'custom_metadata["{_ENABLE_CONSOLIDATION_KEY}"] must be a bool.'
    )
  return enable_consolidation


def _iter_memory_batches(memories: Sequence[str]) -> Sequence[Sequence[str]]:
  """Returns memory slices that comply with direct_memories limits."""
  memory_batches: list[Sequence[str]] = []
  for index in range(0, len(memories), _MAX_DIRECT_MEMORIES_PER_GENERATE_CALL):
    memory_batches.append(
        memories[index : index + _MAX_DIRECT_MEMORIES_PER_GENERATE_CALL]
    )
  return memory_batches


def _build_vertex_metadata(
    metadata_by_key: Mapping[str, object],
) -> dict[str, object]:
  """Converts metadata values to Vertex MemoryMetadataValue objects."""
  vertex_metadata: dict[str, object] = {}
  for key, value in metadata_by_key.items():
    converted_value = _to_vertex_metadata_value(key, value)
    if converted_value is None:
      continue
    vertex_metadata[key] = converted_value
  return vertex_metadata


def _to_vertex_metadata_value(
    key: str,
    value: object,
) -> dict[str, object] | None:
  """Converts a metadata value to Vertex MemoryMetadataValue shape."""
  if isinstance(value, bool):
    return {'bool_value': value}
  if isinstance(value, (int, float)):
    return {'double_value': float(value)}
  if isinstance(value, str):
    return {'string_value': value}
  if isinstance(value, datetime.datetime):
    return {'timestamp_value': value}
  if isinstance(value, Mapping):
    if value.keys() <= {
        'bool_value',
        'double_value',
        'string_value',
        'timestamp_value',
    }:
      return dict(value)
    return {'string_value': str(dict(value))}
  if value is None:
    logger.warning(
        'Ignoring custom metadata key %s because its value is None.',
        key,
    )
    return None
  return {'string_value': str(value)}
