#!/usr/bin/env python3
# Copyright 2022 The ChromiumOS Authors.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Copies rust-bootstrap artifacts from an SDK build to localmirror.

We use localmirror to host these artifacts, but they've changed a bit over
time, so simply `gsutil.py cp $FROM $TO` doesn't work. This script allows the
convenience of the old `cp` command.
"""

import argparse
import logging
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile
from typing import List


_LOCALMIRROR_ROOT = 'gs://chromeos-localmirror/distfiles/'


def _is_in_chroot() -> bool:
  return Path('/etc/cros_chroot_version').exists()


def _ensure_pbzip2_is_installed():
  if shutil.which('pbzip2'):
    return

  logging.info('Auto-installing pbzip2...')
  subprocess.run(['sudo', 'emerge', '-G', 'pbzip2'], check=True)


def _determine_target_path(sdk_path: str) -> str:
  """Determine where `sdk_path` should sit in localmirror."""
  gs_prefix = 'gs://'
  if not sdk_path.startswith(gs_prefix):
    raise ValueError(f'Invalid GS path: {sdk_path!r}')

  file_name = Path(sdk_path[len(gs_prefix):]).name
  return _LOCALMIRROR_ROOT + file_name


def _download(remote_path: str, local_file: Path):
  """Downloads the given gs:// path to the given local file."""
  logging.info('Downloading %s -> %s', remote_path, local_file)
  subprocess.run(
      ['gsutil.py', 'cp', remote_path,
       str(local_file)],
      check=True,
  )


def _debinpkgify(binpkg_file: Path) -> Path:
  """Converts a binpkg into the files it installs.

  Note that this function makes temporary files in the same directory as
  `binpkg_file`. It makes no attempt to clean them up.
  """
  logging.info('Converting %s from a binpkg...', binpkg_file)

  # The SDK builder produces binary packages:
  # https://wiki.gentoo.org/wiki/Binary_package_guide
  #
  # Which means that `binpkg_file` is in the XPAK format. We want to split
  # that out, and recompress it from zstd (which is the compression format
  # that CrOS uses) to bzip2 (which is what we've historically used, and
  # which is what our ebuild expects).
  tmpdir = binpkg_file.parent

  def _mkstemp(suffix=None) -> str:
    fd, file_path = tempfile.mkstemp(dir=tmpdir, suffix=suffix)
    os.close(fd)
    return Path(file_path)

  # First, split the actual artifacts that land in the chroot out to
  # `temp_file`.
  artifacts_file = _mkstemp()
  logging.info('Extracting artifacts from %s into %s...', binpkg_file,
               artifacts_file)
  with artifacts_file.open('wb') as f:
    subprocess.run(
        [
            'qtbz2',
            '-s',
            '-t',
            '-O',
            str(binpkg_file),
        ],
        check=True,
        stdout=f,
    )

  decompressed_artifacts_file = _mkstemp()
  decompressed_artifacts_file.unlink()
  logging.info('Decompressing artifacts from %s to %s...', artifacts_file,
               decompressed_artifacts_file)
  subprocess.run(
      [
          'zstd',
          '-d',
          str(artifacts_file),
          '-o',
          str(decompressed_artifacts_file),
      ],
      check=True,
  )

  # Finally, recompress it as a tbz2.
  tbz2_file = _mkstemp('.tbz2')
  logging.info(
      'Recompressing artifacts from %s to %s (this may take a while)...',
      decompressed_artifacts_file, tbz2_file)
  with tbz2_file.open('wb') as f:
    subprocess.run(
        [
            'pbzip2',
            '-9',
            '-c',
            str(decompressed_artifacts_file),
        ],
        check=True,
        stdout=f,
    )
  return tbz2_file


def _upload(local_file: Path, remote_path: str, force: bool):
  """Uploads the local file to the given gs:// path."""
  logging.info('Uploading %s -> %s', local_file, remote_path)
  cmd_base = ['gsutil.py', 'cp', '-a', 'public-read']
  if not force:
    cmd_base.append('-n')
  subprocess.run(
      cmd_base + [str(local_file), remote_path],
      check=True,
      stdin=subprocess.DEVNULL,
  )


def main(argv: List[str]):
  logging.basicConfig(
      format='>> %(asctime)s: %(levelname)s: %(filename)s:%(lineno)d: '
      '%(message)s',
      level=logging.INFO,
  )

  parser = argparse.ArgumentParser(
      description=__doc__,
      formatter_class=argparse.RawDescriptionHelpFormatter,
  )

  parser.add_argument(
      'sdk_artifact',
      help='Path to the SDK rust-bootstrap artifact to copy. e.g., '
      'gs://chromeos-prebuilt/host/amd64/amd64-host/'
      'chroot-2022.07.12.134334/packages/dev-lang/'
      'rust-bootstrap-1.59.0.tbz2.')
  parser.add_argument(
      '-n',
      '--dry-run',
      action='store_true',
      help='Do everything except actually uploading the artifact.')
  parser.add_argument(
      '--force',
      action='store_true',
      help='Upload the artifact even if one exists in localmirror already.')
  opts = parser.parse_args(argv)

  if not _is_in_chroot():
    parser.error('Run me from within the chroot.')
  _ensure_pbzip2_is_installed()

  target_path = _determine_target_path(opts.sdk_artifact)
  with tempfile.TemporaryDirectory() as tempdir:
    download_path = Path(tempdir) / 'sdk_artifact'
    _download(opts.sdk_artifact, download_path)
    file_to_upload = _debinpkgify(download_path)
    if opts.dry_run:
      logging.info('--dry-run specified; skipping upload of %s to %s',
                   file_to_upload, target_path)
    else:
      _upload(file_to_upload, target_path, opts.force)


if __name__ == '__main__':
  main(sys.argv[1:])