# -*- coding: utf-8 -*- #
# Copyright 2026 Google LLC. All Rights Reserved.
#
# 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.
"""Implementation of gcloud vectorsearch collections data-objects search."""


import json

from apitools.base.py import encoding
from apitools.base.py import list_pager
from googlecloudsdk.api_lib.vector_search import args as common_args
from googlecloudsdk.api_lib.vector_search import clients
from googlecloudsdk.calliope import arg_parsers
from googlecloudsdk.calliope import base
from googlecloudsdk.calliope import exceptions as calliope_exceptions
from googlecloudsdk.core import log
from googlecloudsdk.core import properties
from googlecloudsdk.core.util import files


SEMANTIC_SEARCH_TASK_TYPE_CHOICES = {
    'classification': 'Specifies that the given text will be classified.',
    'clustering': 'Specifies that the embeddings will be used for clustering.',
    'code-retrieval-query': (
        'Specifies that the embeddings will be used for code retrieval.'
    ),
    'fact-verification': (
        'Specifies that the embeddings will be used for fact verification.'
    ),
    'question-answering': (
        'Specifies that the embeddings will be used for question answering.'
    ),
    'retrieval-document': (
        'Specifies the given text is a document from the corpus being searched.'
    ),
    'retrieval-query': (
        'Specifies the given text is a query in a search/retrieval setting.'
    ),
    'semantic-similarity': 'Specifies the given text will be used for STS.',
}


VECTOR_SEARCH_DISTANCE_METRIC_CHOICES = {
    'dot-product': 'Dot product distance metric.',
    'cosine-distance': 'Cosine distance metric.',
}


@base.ReleaseTracks(base.ReleaseTrack.BETA, base.ReleaseTrack.GA)
@base.DefaultUniverseOnly
class Search(base.Command):
  """Search data objects from a Vector Search collection."""

  detailed_help = {
      'DESCRIPTION': (
          """
          Search data objects from a Vector Search collection.
      """
      ),
      'EXAMPLES': (
          """
          To search data objects from collection `my-collection` in location `us-central1` using text search and return 10 results, run:

          $ {command} --collection=my-collection --location=us-central1 --text-search-text="test" --text-search-data-fields="text_field" --top-k=10

          To search data objects from collection `my-collection` in location `us-central1` using semantic search and return 5 results, run:

          $ {command} --collection=my-collection --location=us-central1 --semantic-search-text="sci-fi" --semantic-search-field="plot_embedding" --semantic-task-type="retrieval-query" --top-k=5

          To search data objects from collection `my-collection` in location `us-central1` using vector search with an index hint and return 7 results, run:

          $ {command} --collection=my-collection --location=us-central1 --vector-search-field="genre_embedding" --vector-from-file="vector.json" --use-index="my-index" --top-k=7

          To search data objects from collection `my-collection` in location `us-central1` using vector search with KNN for exact results, run:

          $ {command} --collection=my-collection --location=us-central1 --vector-search-field="genre_embedding" --vector-from-file="vector.json" --use-knn --top-k=7

      """
      ),
  }

  @classmethod
  def Args(cls, parser):
    """Register flags for this command."""
    common_args.AddDataObjectFlags(parser, 'search')
    common_args.AddOutputFieldsFlags(parser)

    parser.add_argument(
        '--top-k',
        type=int,
        help='The number of nearest neighbors to return. Default is 10.',
    )

    search_hint_group = parser.add_mutually_exclusive_group('Search Hint')

    use_index_group = search_hint_group.add_group('Use Index Options')
    use_index_group.add_argument(
        '--use-index',
        metavar='INDEX',
        required=True,
        help=(
            """
            Full resource name or ID of the index to use for the search.
            This flag is compatible only with Semantic Search and Vector Search.
            """
        ),
    )
    if cls.ReleaseTrack() == base.ReleaseTrack.BETA:
      use_index_group.add_argument(
          '--dense-scann-search-leaves-pct',
          type=int,
          metavar='PERCENTAGE',
          help=(
              'The percentage of leaves to search for dense ScaNN, in the range'
              ' [0, 100].'
          ),
      )
      use_index_group.add_argument(
          '--dense-scann-initial-candidate-count',
          type=int,
          metavar='CANDIDATE_COUNT',
          help='The number of initial candidates for dense ScaNN.',
      )

    search_hint_group.add_argument(
        '--use-knn',
        action='store_true',
        help=(
            """
            If set to true, the search will use the system's default K-Nearest
            Neighbor (KNN) index engine.
            This flag is compatible only with Semantic Search and Vector Search.
            """
        ),
    )

    search_type_group = parser.add_mutually_exclusive_group(
        'Search type', required=True
    )
    semantic_search_group = search_type_group.add_argument_group(
        'Semantic Search'
    )
    vector_search_group = search_type_group.add_argument_group('Vector Search')
    text_search_group = search_type_group.add_argument_group('Text Search')

    # Semantic search flags
    semantic_search_group.add_argument(
        '--semantic-search-text',
        required=True,
        help='The query text for semantic search.',
    )
    semantic_search_group.add_argument(
        '--semantic-search-field',
        required=True,
        help='The vector field to search.',
    )
    semantic_search_group.add_argument(
        '--semantic-task-type',
        required=True,
        choices=SEMANTIC_SEARCH_TASK_TYPE_CHOICES,
        help='The task type of the query embedding for semantic search.',
    )

    # Vector search flags
    vector_search_group.add_argument(
        '--vector-search-field',
        required=True,
        help='The vector field to search.',
    )
    vector_search_group.add_argument(
        '--vector-from-file',
        required=True,
        help="""Path to a JSON file containing dense or sparse vector to search with.

* Example file content for dense vector:

```
{
  "dense": {
    "values": [
      0.7,
      0.6,
      0.5,
      0.4
    ]
  }
}
```

* Example file content for sparse vector:

```
{
  "sparse": {
    "indices": [1, 5, 10],
    "values": [0.1, 0.5, 0.21]
  }
}
```
""",
    )
    vector_search_group.add_argument(
        '--distance-metric',
        choices=VECTOR_SEARCH_DISTANCE_METRIC_CHOICES,
        help=(
            'The distance metric to use for the KNN search. If not specified,'
            ' dot-product will be used as the default.'
        ),
    )

    # Text search flags
    text_search_group.add_argument(
        '--text-search-text',
        required=True,
        help='The query text for text search.',
    )
    text_search_group.add_argument(
        '--text-search-data-fields',
        type=arg_parsers.ArgList(),
        required=True,
        help='The data field names to search.',
        metavar='DATA_FIELD_NAME',
    )

  def _GetOutputFields(self, args, client):
    if (
        args.output_data_fields
        or args.output_vector_fields
        or args.output_metadata_fields
    ):
      return common_args.ParseOutputFields(args, client)
    return None

  def _GetSearchHint(self, args, client, parent):
    if args.use_index:
      # if this is not full resource name, add parent path to the index name
      index_name = args.use_index
      if '/' not in index_name:
        index_name = '{}/indexes/{}'.format(
            parent, index_name
        )
      index_hint_msg = (
          client.GetMessage('SearchHint.IndexHint')(
              name=index_name
          )
      )
      if self.ReleaseTrack() == base.ReleaseTrack.BETA and (
          args.dense_scann_search_leaves_pct
          or args.dense_scann_initial_candidate_count
      ):
        dense_scann_params = (
            client.GetMessage('SearchHint.IndexHint.DenseScannParams')()
        )
        if args.dense_scann_search_leaves_pct:
          dense_scann_params.searchLeavesPct = (
              args.dense_scann_search_leaves_pct
          )
        if args.dense_scann_initial_candidate_count:
          dense_scann_params.initialCandidateCount = (
              args.dense_scann_initial_candidate_count
          )
        index_hint_msg.denseScannParams = dense_scann_params
      return client.GetMessage('SearchHint')(
          indexHint=index_hint_msg
      )
    elif args.use_knn:
      return client.GetMessage('SearchHint')(
          knnHint=client.GetMessage('SearchHint.KnnHint')()
      )
    return None

  def _GetFilterDict(self, args):
    if args.json_filter:
      try:
        return json.loads(args.json_filter)
      except json.JSONDecodeError as e:
        raise calliope_exceptions.InvalidArgumentException(
            '--json-filter', f'Invalid JSON: {e}'
        )
    return None

  def _BuildSemanticSearchMessage(
      self, args, client, filter_dict, search_hint, output_fields
  ):
    """Builds a SemanticSearch message."""
    semantic_search = client.GetMessage('SemanticSearch')(
        searchText=args.semantic_search_text,
        searchField=args.semantic_search_field,
        topK=args.top_k,
    )
    if args.semantic_task_type:
      task_type_code_name = args.semantic_task_type.replace('-', '_').upper()
      semantic_search.taskType = client.GetMessage(
          'SemanticSearch'
      ).TaskTypeValueValuesEnum.lookup_by_name(task_type_code_name)
    if filter_dict:
      try:
        semantic_search.filter = encoding.DictToMessage(
            filter_dict,
            client.GetMessage('SemanticSearch').FilterValue,
        )
      except Exception as e:
        raise calliope_exceptions.InvalidArgumentException(
            '--json-filter', f'Error converting JSON filter to message: {e}'
        )
    if search_hint:
      semantic_search.searchHint = search_hint
    if output_fields:
      semantic_search.outputFields = output_fields
    return semantic_search

  def _FillInVectorFromFile(self, args, vector_search, client):
    try:
      file_content = files.ReadFileContents(args.vector_from_file)
      vector_dict = json.loads(file_content)
      if len(vector_dict) != 1 or not (
          'dense' in vector_dict or 'sparse' in vector_dict
      ):
        raise calliope_exceptions.InvalidArgumentException(
            '--vector-from-file',
            'Vector file must contain exactly one vector field under key'
            ' "dense" or "sparse".',
        )
      if 'dense' in vector_dict:
        vector_search.vector = encoding.DictToMessage(
            vector_dict['dense'],
            client.GetMessage('DenseVector'),
        )
      elif 'sparse' in vector_dict:
        vector_search.sparseVector = encoding.DictToMessage(
            vector_dict['sparse'],
            client.GetMessage('SparseVector'),
        )
    except files.Error as e:
      raise calliope_exceptions.InvalidArgumentException(
          '--vector-from-file', f'Could not read file: {e}'
      )
    except json.JSONDecodeError as e:
      raise calliope_exceptions.InvalidArgumentException(
          '--vector-from-file', f'Invalid JSON: {e}'
      )
    except Exception as e:
      raise calliope_exceptions.InvalidArgumentException(
          '--vector-from-file',
          f'Error parsing vector file: {e}',
      )

  def _BuildVectorSearchMessage(
      self, args, client, filter_dict, search_hint, output_fields
  ):
    """Builds a VectorSearch message."""
    vector_search = client.GetMessage('VectorSearch')(
        searchField=args.vector_search_field,
        topK=args.top_k,
    )
    self._FillInVectorFromFile(args, vector_search, client)
    if args.distance_metric:
      distance_metric_code_name = args.distance_metric.replace('-', '_').upper()
      vector_search.distanceMetric = client.GetMessage(
          'VectorSearch'
      ).DistanceMetricValueValuesEnum.lookup_by_name(distance_metric_code_name)
    if filter_dict:
      try:
        vector_search.filter = encoding.DictToMessage(
            filter_dict,
            client.GetMessage('VectorSearch').FilterValue,
        )
      except Exception as e:
        raise calliope_exceptions.InvalidArgumentException(
            '--json-filter', f'Error converting JSON filter to message: {e}'
        )

    if search_hint:
      vector_search.searchHint = search_hint
    if output_fields:
      vector_search.outputFields = output_fields
    return vector_search

  def _LogWarningsIfNeeded(self, args):
    """Logs warnings if incompatible flags are used."""
    if args.use_index:
      log.warning(
          'Search hint is not supported for Text Search. Ignoring'
          ' --use-index=%s',
          args.use_index,
      )
    if args.use_knn:
      log.warning(
          'Search hint is not supported for Text Search. Ignoring --use-knn',
      )

  def _BuildTextSearchMessage(self, args, client, output_fields, filter_dict):
    """Builds a TextSearch message."""
    text_search = client.GetMessage('TextSearch')(
        searchText=args.text_search_text,
        dataFieldNames=args.text_search_data_fields,
        topK=args.top_k,
    )
    if output_fields:
      text_search.outputFields = output_fields
    if filter_dict:
      try:
        text_search.filter = encoding.DictToMessage(
            filter_dict,
            client.GetMessage('TextSearch').FilterValue,
        )
      except Exception as e:
        raise calliope_exceptions.InvalidArgumentException(
            '--json-filter', f'Error converting JSON filter to message: {e}'
        )

    return text_search

  def Run(self, args):
    """Run the search command."""
    client = clients.DataObjectsClient(self.ReleaseTrack())
    project = properties.VALUES.core.project.GetOrFail()
    parent = 'projects/{}/locations/{}/collections/{}'.format(
        project, args.location, args.collection
    )

    search_request_body = client.GetMessage('SearchDataObjectsRequest')()

    output_fields = self._GetOutputFields(args, client)
    search_hint = self._GetSearchHint(args, client, parent)
    filter_dict = self._GetFilterDict(args)

    if (
        args.semantic_search_text
        or args.semantic_search_field
        or args.semantic_task_type
    ):
      search_request_body.semanticSearch = self._BuildSemanticSearchMessage(
          args, client, filter_dict, search_hint, output_fields
      )
    elif args.vector_from_file or args.vector_search_field:
      search_request_body.vectorSearch = self._BuildVectorSearchMessage(
          args, client, filter_dict, search_hint, output_fields
      )
    elif args.text_search_text or args.text_search_data_fields:
      self._LogWarningsIfNeeded(args)
      search_request_body.textSearch = self._BuildTextSearchMessage(
          args, client, output_fields, filter_dict
      )

    search_req_body_field = client.GetRequestField('SearchDataObjectsRequest')
    full_req_message = 'VectorsearchProjectsLocationsCollectionsDataObjectsSearchRequest'
    full_req = getattr(client.messages, full_req_message)(
        parent=parent,
        **{search_req_body_field: search_request_body},
    )

    return list_pager.YieldFromList(
        client.service,
        full_req,
        batch_size_attribute=(
            search_req_body_field,
            'pageSize',
        ),
        field='results',
        method='Search',
        current_token_attribute=(
            search_req_body_field,
            'pageToken',
        ),
        next_token_attribute='nextPageToken',
    )
