#!/usr/bin/env python3

# Copyright 2020 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Manages IREE Docker image definitions.

Includes information on their dependency graph and GCR URL.

See the README for more information on how to add and update images.

Example usage:

Rebuild the cmake image and all images that transitively on depend on it,
tagging them with `latest` and updating all references to their sha digests:
  python3 build_tools/docker/manage_images.py --image cmake

Print out output for rebuilding the cmake image and all images that
transitively on depend on it, but don't take side-effecting actions:
  python3 build_tools/docker/manage_images.py --image cmake --dry-run

Rebuild and push all images and update references to them in the repository:
  python3 build_tools/docker/manage_images.py --images all
"""

import argparse
import fileinput
import os
import posixpath
import re
import subprocess
import sys
from typing import Dict, List, Sequence, Union

import utils

IREE_GCR_URL = 'gcr.io/iree-oss/'
DIGEST_REGEX = r'sha256:[a-zA-Z0-9]+'
DOCKER_DIR = 'build_tools/docker/'.replace('/', os.sep)

# Map from image names to images that they depend on.
IMAGES_TO_DEPENDENCIES = {
    'base': [],
    'bazel': ['base', 'util'],
    'cmake': ['base', 'util'],
    'cmake-android': ['cmake-python', 'util'],
    'cmake-emscripten': ['cmake'],
    'cmake-python': ['cmake'],
    'cmake-python-vulkan': ['cmake-python', 'vulkan'],
    'cmake-python-swiftshader': ['cmake-python-vulkan', 'swiftshader'],
    'cmake-python-nvidia': ['cmake-python-vulkan', 'util'],
    'cmake-riscv': ['cmake', 'util'],
    'cmake-bazel-frontends': ['cmake-python', 'bazel'],
    'cmake-bazel-frontends-android': ['cmake-bazel-frontends', 'cmake-android'],
    'cmake-bazel-frontends-vulkan': ['cmake-bazel-frontends', 'vulkan'],
    'cmake-bazel-frontends-swiftshader': [
        'cmake-bazel-frontends-vulkan', 'swiftshader'
    ],
    'cmake-bazel-frontends-nvidia': ['cmake-bazel-frontends-vulkan'],
    'rbe-toolchain': ['vulkan'],
    'samples': ['cmake-python-swiftshader'],
    'swiftshader': ['cmake'],
    'util': [],
    'vulkan': ['util'],
}

IMAGES_TO_DEPENDENT_IMAGES = {k: [] for k in IMAGES_TO_DEPENDENCIES}
for image, dependencies in IMAGES_TO_DEPENDENCIES.items():
  for dependency in dependencies:
    IMAGES_TO_DEPENDENT_IMAGES[dependency].append(image)

IMAGES_HELP = [f'`{name}`' for name in IMAGES_TO_DEPENDENCIES]
IMAGES_HELP = f'{", ".join(IMAGES_HELP)} or `all`'


def parse_arguments():
  """Parses command-line options."""
  parser = argparse.ArgumentParser(
      description="Build IREE's Docker images and optionally push them to GCR.")
  parser.add_argument('--images',
                      '--image',
                      type=str,
                      required=True,
                      action='append',
                      help=f'Name of the image to build: {IMAGES_HELP}.')
  parser.add_argument(
      '--dry_run',
      '--dry-run',
      '-n',
      action='store_true',
      help='Print output without building or pushing any images.')

  args = parser.parse_args()
  for image in args.images:
    if image == 'all':
      # Sort for a determinstic order
      args.images = sorted(IMAGES_TO_DEPENDENCIES.keys())
    elif image not in IMAGES_TO_DEPENDENCIES:
      raise parser.error('Expected --image to be one of:\n'
                         f'  {IMAGES_HELP}\n'
                         f'but got `{image}`.')
  return args


def _dag_dfs(input_nodes: Sequence[str],
             node_to_child_nodes: Dict[str, Sequence[str]]) -> List[str]:
  # Python doesn't have a builtin OrderedSet, but we don't have many images, so
  # we just use a list.
  ordered_nodes = []

  def add_children(parent_node: str):
    if parent_node not in ordered_nodes:
      for child_node in node_to_child_nodes[parent_node]:
        add_children(child_node)
      ordered_nodes.append(parent_node)

  for node in input_nodes:
    add_children(node)
  return ordered_nodes


def get_ordered_images_to_process(images: Sequence[str]) -> List[str]:
  dependents = _dag_dfs(images, IMAGES_TO_DEPENDENT_IMAGES)
  dependents.reverse()
  return dependents


def get_dependencies(images: Sequence[str]) -> List[str]:
  return _dag_dfs(images, IMAGES_TO_DEPENDENCIES)


def get_repo_digest(tagged_image_url: str, dry_run: bool = False) -> str:
  inspect_command = [
      'docker',
      'image',
      'inspect',
      tagged_image_url,
      '-f',
      '{{index .RepoDigests 0}}',
  ]
  try:
    completed_process = utils.run_command(
        inspect_command,
        dry_run=False,  # Run even if --dry_run is True.
        capture_output=True,
        timeout=10)
  except subprocess.CalledProcessError as error:
    if dry_run:
      return ""
    else:
      raise RuntimeError(
          f'Computing the repository digest for {tagged_image_url} failed. Has '
          'it been pushed to GCR?') from error
  _, repo_digest = completed_process.stdout.strip().split('@')
  return repo_digest


def update_rbe_reference(digest: str, dry_run: bool = False):
  print('Updating WORKSPACE file for rbe-toolchain')
  digest_updates = 0
  for line in fileinput.input(files=['WORKSPACE'], inplace=True):
    if line.strip().startswith('digest ='):
      digest_updates += 1
      if dry_run:
        print(line, end='')
      else:
        print(re.sub(DIGEST_REGEX, digest, line), end='')
    else:
      print(line, end='')

  if digest_updates > 1:
    raise RuntimeError(
        "There is more than one instance of 'digest =' in the WORKSPACE file. "
        "This means that more than just the 'rbe_toolchain' digest was "
        "overwritten, and the file should be restored.")


def update_references(image_url: str, digest: str, dry_run: bool = False):
  """Updates all references to 'image_url' with a sha256 digest."""
  print(f'Updating references to {image_url}')

  grep_command = ['git', 'grep', '-l', f'{image_url}@sha256']
  try:
    completed_process = utils.run_command(grep_command,
                                          capture_output=True,
                                          timeout=5)
  except subprocess.CalledProcessError as error:
    if error.returncode == 1:
      print(f'Found no references to {image_url}')
      return
    raise error

  # Update references in all grepped files.
  files = completed_process.stdout.split()
  print(f'Updating references in {len(files)} files: {files}')
  if not dry_run:
    for line in fileinput.input(files=files, inplace=True):
      print(re.sub(f'{image_url}@{DIGEST_REGEX}', f'{image_url}@{digest}',
                   line),
            end='')


def parse_prod_digests() -> Dict[str, str]:
  image_urls_to_prod_digests = {}
  with open(utils.PROD_DIGESTS_PATH, "r") as f:
    for line in f:
      image_url, digest = line.strip().split("@")
      image_urls_to_prod_digests[image_url] = digest
  return image_urls_to_prod_digests


if __name__ == '__main__':
  args = parse_arguments()

  # Ensure the user has the correct authorization to push to GCR.
  utils.check_gcloud_auth(dry_run=args.dry_run)

  images_to_process = get_ordered_images_to_process(args.images)
  print(f'Also processing dependent images. Will process: {images_to_process}')

  dependencies = get_dependencies(images_to_process)
  print(f'Pulling image dependencies: {dependencies}')
  image_urls_to_prod_digests = parse_prod_digests()
  for dependency in dependencies:
    dependency_url = posixpath.join(IREE_GCR_URL, dependency)
    # If `dependency` is a new image then it may not have a prod digest yet.
    if dependency_url in image_urls_to_prod_digests:
      digest = image_urls_to_prod_digests[dependency_url]
      dependency_with_digest = f'{dependency_url}@{digest}'
      utils.run_command(["docker", "pull", dependency_with_digest],
                        dry_run=args.dry_run)

  for image in images_to_process:
    print('\n' * 5 + f'Processing image {image}')
    image_url = posixpath.join(IREE_GCR_URL, image)
    tagged_image_url = f'{image_url}'
    image_path = os.path.join(DOCKER_DIR, image)

    utils.run_command(
        ['docker', 'build', '--tag', tagged_image_url, image_path],
        dry_run=args.dry_run)

    utils.run_command(['docker', 'push', tagged_image_url],
                      dry_run=args.dry_run)

    digest = get_repo_digest(tagged_image_url, args.dry_run)

    # Check that the image is in 'prod_digests.txt' and append it to the list
    # in the file if it isn't.
    if image_url not in image_urls_to_prod_digests:
      image_with_digest = f'{image_url}@{digest}'
      print(
          f'Adding new image {image_with_digest} to {utils.PROD_DIGESTS_PATH}')
      if not args.dry_run:
        with open(utils.PROD_DIGESTS_PATH, 'a') as f:
          f.write(f'{image_with_digest}\n')

    # Just hardcode this oddity
    if image == 'rbe-toolchain':
      update_rbe_reference(digest, dry_run=args.dry_run)
    update_references(image_url, digest, dry_run=args.dry_run)