# -*- coding: utf-8 -*- #
# Copyright 2024 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.
"""The command lists the models in Model Garden and their supported functionalities."""

from googlecloudsdk.api_lib.ai.model_garden import client as client_mg
from googlecloudsdk.calliope import base
from googlecloudsdk.calliope import exceptions as c_exceptions
from googlecloudsdk.command_lib.ai import constants
from googlecloudsdk.command_lib.ai import endpoint_util
from googlecloudsdk.command_lib.ai.region_util import (
    _IsDefaultUniverse,
)
from googlecloudsdk.core import properties

_SHORT_NAME_FORMAT = (
    'format("{0:s}@{1:s}/{2:s}", name, versionId,'
    ' name.regex("publishers/hf-.*", "@hf", "@mg")).sub("publishers/hf-",'
    ' "").sub("publishers/", "").sub("models/", "").sub("@001/@hf", "").'
    ' sub("/@mg", ""):sort=1'
)
_FULL_RESOURCE_NAME_FORMAT = 'format("{0:s}@{1:s}", name, versionId):sort=1'
_MODEL_ID_LABEL = ':label=MODEL_ID'

_CAN_DEPLOY_FILTER = 'supportedActions.multiDeployVertex.yesno(yes=Yes)'
_CAN_DEPLOY_LABEL = ':label=CAN_DEPLOY'

_CAN_PREDICT_FILTER = 'publisherModelTemplate.yesno(yes=Yes)'
_CAN_PREDICT_LABEL = ':label=CAN_PREDICT'

_DEFAULT_TABLE_FORMAT = (
    f'table({_SHORT_NAME_FORMAT}{_MODEL_ID_LABEL},'
    f' {_CAN_DEPLOY_FILTER}{_CAN_DEPLOY_LABEL},'
    f' {_CAN_PREDICT_FILTER}{_CAN_PREDICT_LABEL})'
)
_SHORT_MODEL_NAME_ONLY_TABLE_FORMAT = (
    f'table({_SHORT_NAME_FORMAT}{_MODEL_ID_LABEL})'
)
_FULL_RESOURCE_NAME_ONLY_TABLE_FORMAT = (
    f'table({_FULL_RESOURCE_NAME_FORMAT}{_MODEL_ID_LABEL})'
)
_FULL_RESOURCE_NAME_TABLE_FORMAT = (
    f'table({_FULL_RESOURCE_NAME_FORMAT}{_MODEL_ID_LABEL},'
    f' {_CAN_DEPLOY_FILTER}{_CAN_DEPLOY_LABEL},'
    f' {_CAN_PREDICT_FILTER}{_CAN_PREDICT_LABEL})'
)


@base.ReleaseTracks(
    base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA, base.ReleaseTrack.GA
)
@base.UniverseCompatible
class List(base.ListCommand):
  """List the publisher models in Model Garden.

  This command lists either all models in Model Garden or all Hugging
  Face models supported by Model Garden.

  ## EXAMPLES

  To list all models in Model Garden, run:

    $ gcloud ai model-garden models list

  To list Hugging Face models that can be deployed in Model Garden, run:

    $ gcloud ai model-garden models list --can-deploy-hugging-face-models

  To list models with `gemma` in their names, run:

    $ gcloud ai model-garden models list --model-filter=gemma

  Note: Since the number of Hugging Face models is large, the default limit is
  set to 500 with a page size of 100 when listing supported Hugging Face models.
  To override the limit or page size, specify the --limit or --page-size flags,
  respectively. To list all models in Model Garden, use `--limit=unlimited`.
  """

  @staticmethod
  def Args(parser):
    # The default display format is set in Run() rather than here so that an
    # explicit --format flag (e.g., --format="json(field)") is used as-is
    # without being composed with the default table projection. Format
    # composition causes default table columns to leak into json(...) output via
    # shared parent paths in the projection tree.
    parser.add_argument(
        '--can-deploy-hugging-face-models',
        action='store_true',
        default=False,
        required=False,
        help='Whether to only list Hugging Face models that can be deployed.',
    )
    parser.add_argument(
        '--model-filter',
        action='store',
        default=None,
        required=False,
        help=(
            'Filter to apply to the model names or the display names of the'
            ' list of models.'
        ),
    )
    parser.add_argument(
        '--full-resource-name',
        action='store_true',
        default=False,
        required=False,
        help='Whether to return the full resource name of the model.',
    )
    base.URI_FLAG.RemoveFromParser(parser)
    base.LIMIT_FLAG.SetDefault(parser, 1000)

  def Run(self, args):
    version = constants.BETA_VERSION

    # Set the default display format only when the user hasn't provided an
    # explicit --format.  Without this guard, the default table(...) projection
    # is composed with the explicit format, and the projection tree's INNER
    # flag on shared parent paths (e.g. supportedActions) causes sibling
    # branches from the table projection (e.g. multiDeployVertex) to leak
    # into json(...) output.
    if not args.IsSpecified('format'):
      if args.can_deploy_hugging_face_models:
        args.GetDisplayInfo().AddFormat(
            _FULL_RESOURCE_NAME_ONLY_TABLE_FORMAT
            if args.full_resource_name
            else _SHORT_MODEL_NAME_ONLY_TABLE_FORMAT
        )
      elif args.full_resource_name:
        args.GetDisplayInfo().AddFormat(_FULL_RESOURCE_NAME_TABLE_FORMAT)
      else:
        args.GetDisplayInfo().AddFormat(_DEFAULT_TABLE_FORMAT)

    # Set the default page size to 100 if the user requests to list supported
    # Hugging Face models, since there are tens of thousands of Hugging Face
    # models and the call will take a long time.
    if args.can_deploy_hugging_face_models:
      if args.page_size is None:
        args.page_size = 100

    if _IsDefaultUniverse():
      api_region = 'us-central1'
    else:
      api_region = properties.VALUES.ai.region.Get()
      if not api_region:
        raise c_exceptions.RequiredArgumentException(
            'ai/region',
            'The [ai/region] property must be set in this environment.',
        )
    with endpoint_util.AiplatformEndpointOverrides(
        version, region=api_region, use_info_log=True
    ):
      mg_client = client_mg.ModelGardenClient(version)
      return mg_client.ListPublisherModels(
          limit=args.limit,
          batch_size=args.page_size,
          list_hf_models=args.can_deploy_hugging_face_models,
          model_filter=args.model_filter,
      )
