#!/usr/bin/env vpython3
# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Updates the Fuchsia product bundles to the given revision. Should be used
in a 'hooks_os' entry so that it only runs when .gclient's target_os includes
'fuchsia'."""

import argparse
import json
import logging
import os
import sys

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
                                             'test')))

import common
import update_sdk
from compatible_utils import running_unattended


# TODO(crbug.com/40863468): Remove when the old scripts have been deprecated.
_IMAGE_TO_PRODUCT_BUNDLE = {
    'qemu.arm64': 'terminal.qemu-arm64',
    'qemu.x64': 'terminal.x64',
}


# TODO(crbug.com/40863468): Remove when the old scripts have been deprecated.
def convert_to_products(images_list):
  """Convert image names in the SDK to product bundle names."""

  product_bundle_list = []
  for image in images_list:
    if image in _IMAGE_TO_PRODUCT_BUNDLE:
      logging.warning(f'Image name {image} has been deprecated. Use '
                      f'{_IMAGE_TO_PRODUCT_BUNDLE.get(image)} instead.')
      product_bundle_list.append(_IMAGE_TO_PRODUCT_BUNDLE[image])
    else:
      if image.endswith('-release'):
        image = image[:-len('-release')]
        logging.warning(f'Image name {image}-release has been deprecated. Use '
                        f'{image} instead.')
      product_bundle_list.append(image)
  return product_bundle_list


def remove_repositories(repo_names_to_remove):
  """Removes given repos from repo list.
  Repo MUST be present in list to succeed.

  Args:
    repo_names_to_remove: List of repo names (as strings) to remove.
  """
  for repo_name in repo_names_to_remove:
    common.run_ffx_command(cmd=('repository', 'remove', repo_name))


def get_repositories():
  """Lists repositories that are available on disk.

  Also prunes repositories that are listed, but do not have an actual packages
  directory.

  Returns:
    List of dictionaries containing info about the repositories. They have the
    following structure:
    {
      'name': <repo name>,
      'spec': {
        'type': <type, usually pm>,
        'path': <path to packages directory>
      },
    }
  """

  repos = json.loads(
      common.run_ffx_command(cmd=('--machine', 'json', 'repository', 'list'),
                             capture_output=True).stdout.strip())
  to_prune = set()
  sdk_root_abspath = os.path.abspath(os.path.dirname(common.SDK_ROOT))
  for repo in repos:
    # Confirm the path actually exists. If not, prune list.
    # Also assert the product-bundle repository is for the current repo
    # (IE within the same directory).
    if not os.path.exists(repo['spec']['path']):
      to_prune.add(repo['name'])

    if not repo['spec']['path'].startswith(sdk_root_abspath):
      to_prune.add(repo['name'])

  repos = [repo for repo in repos if repo['name'] not in to_prune]

  remove_repositories(to_prune)
  return repos


def get_current_signature(image_dir):
  """Determines the current version of the image, if it exists.

  Returns:
    The current version, or None if the image is non-existent.
  """

  version_file = os.path.join(image_dir, 'product_bundle.json')
  if os.path.exists(version_file):
    with open(version_file) as f:
      try:
        data = json.load(f)
      except json.decoder.JSONDecodeError:
        logging.warning('product_bundle.json is not at the JSON format and may be empty.')
        return None
      if 'product_version' not in data:
        logging.warning('The key "product_version" does not exist in product_bundle.json')
        return None
      return data['product_version']
  return None


# VisibleForTesting
def internal_hash():
  hash_filename = os.path.join(os.path.dirname(__file__),
                               'linux_internal.sdk.sha1')
  return (open(hash_filename, 'r').read().strip()
          if os.path.exists(hash_filename) else '')


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('--verbose',
                      '-v',
                      action='store_true',
                      help='Enable debug-level logging.')
  parser.add_argument(
      'products',
      type=str,
      help='List of product bundles to download, represented as a comma '
      'separated list.')
  parser.add_argument(
      '--internal',
      action='store_true',
      help='Whether the images are coming from internal, it impacts version '
      'file, bucket and download location.')
  args = parser.parse_args()

  logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)

  # Check whether there's Fuchsia support for this platform.
  common.get_host_os()

  new_products = convert_to_products(args.products.split(','))
  logging.debug('Searching for the following products: %s', str(new_products))

  logging.debug('Getting new SDK hash')
  if args.internal:
    new_hash = internal_hash()
  else:
    new_hash = common.get_hash_from_sdk()

  auth_args = [
      '--auth',
      os.path.join(os.path.dirname(__file__), 'get_auth_token.py')
  ] if running_unattended() else []
  for product in new_products:
    prod, board = product.split('.', 1)
    if prod.startswith('smart_display_') and board in [
        'astro', 'sherlock', 'nelson'
    ]:
      # This is a hacky way of keeping the files into the folders matching
      # the original image name, since the definition is unfortunately in
      # src-internal. Likely we can download two copies for a smooth
      # transition, but it would be easier to keep it as-is during the ffx
      # product v2 migration.
      # TODO(crbug.com/40938340): Migrate the image download folder away from
      # the following hack.
      prod, board = board + '-release', prod
    if args.internal:
      # sdk_override.txt does not work for internal images.
      override_url = None
      image_dir = os.path.join(common.INTERNAL_IMAGES_ROOT, prod, board)
    else:
      override_url = update_sdk.GetSDKOverrideGCSPath()
      if override_url:
        # TODO(zijiehe): Convert to removesuffix once python 3.9 is supported.
        if override_url.endswith('/sdk'):
          override_url = override_url[:-len('/sdk')]
        logging.debug(f'Using {override_url} from override file.')
      image_dir = os.path.join(common.IMAGES_ROOT, prod, board)
    curr_signature = get_current_signature(image_dir)

    if not override_url and curr_signature == new_hash:
      continue

    common.make_clean_directory(image_dir)
    base_url = override_url or 'gs://{bucket}/development/{new_hash}'.format(
        bucket='fuchsia-sdk' if args.internal else 'fuchsia', new_hash=new_hash)
    effective_auth_args = auth_args if base_url.startswith(
        'gs://fuchsia-artifacts-internal/') or base_url.startswith(
            'gs://fuchsia-sdk/') else []
    lookup_output = common.run_ffx_command(cmd=[
        '--machine', 'json', 'product', 'lookup', product, new_hash,
        '--base-url', base_url
    ] + effective_auth_args,
                                           capture_output=True).stdout.strip()
    download_url = json.loads(lookup_output)['transfer_manifest_url']
    # The download_url is purely a timestamp based gs location and is fairly
    # meaningless, so we log the base_url instead which contains the sdk version
    # if it's not coming from the sdk_override.txt file.
    logging.info(f'Downloading {product} from {base_url} and {download_url}.')
    common.run_ffx_command(
        cmd=['product', 'download', download_url, image_dir] +
        effective_auth_args)

  return 0


if __name__ == '__main__':
  sys.exit(main())