# 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.
"""A small subset of commands for interaction with Google Cloud Storage.
It's used when py_utils from the catapult repo are not available (as is the case
on autorollers).
"""

import hashlib
import subprocess
import shutil

PUBLIC_BUCKET = 'chromium-telemetry'
INTERNAL_BUCKET = 'chrome-telemetry'


def _RunCommand(args):
  gcloud_path = shutil.which('gcloud')
  if not gcloud_path:
    raise RuntimeError("gcloud executable not found in PATH.")
  try:
    result = subprocess.run(
      [gcloud_path, 'storage'] + args,
      capture_output=True,
      text=True,
      check=True,
    )
    return result.stdout
  except subprocess.CalledProcessError as e:
    raise RuntimeError(e.stderr)


def Get(bucket, remote_path, local_path):
  url = 'gs://%s/%s' % (bucket, remote_path)
  _RunCommand(['cp', url, local_path])


def Exists(bucket, remote_path):
  try:
    _RunCommand(['ls', 'gs://%s/%s' % (bucket, remote_path)])
    return True
  except RuntimeError:
    return False


def ListFiles(bucket, path='', sort_by='name'):
  """Returns files matching the given path in bucket.

  Args:
    bucket: Name of cloud storage bucket to look at.
    path: Path within the bucket to filter to. Path can include wildcards.
    sort_by: 'name' (default), 'time' or 'size'.

  Returns:
    A sorted list of files.
  """
  bucket_prefix = 'gs://%s' % bucket
  full_path = '%s/%s' % (bucket_prefix, path)
  stdout = _RunCommand(['ls', '-l', full_path])

  # Filter out directories and the summary line.
  file_infos = [
    line.split(None, 2)
    for line in stdout.splitlines()
    if len(line) > 0 and not line.startswith("TOTAL") and not line.endswith('/')
  ]

  # The first field in the info is size, the second is time, the third is name.
  if sort_by == 'size':
    file_infos.sort(key=lambda info: int(info[0]))
  elif sort_by == 'time':
    file_infos.sort(key=lambda info: info[1])
  elif sort_by == 'name':
    file_infos.sort(key=lambda info: info[2])
  else:
    raise ValueError("Wrong sort_by value: %s" % sort_by)

  return [url[len(bucket_prefix) :] for _, _, url in file_infos]


def Insert(bucket, remote_path, local_path, publicly_readable):
  """Upload file in |local_path| to cloud storage.

  Newer version of 'Insert()' returns an object instead of a string.

  Args:
    bucket: the google cloud storage bucket name.
    remote_path: the remote file path in |bucket|.
    local_path: path of the local file to be uploaded.
    publicly_readable: whether the uploaded file has publicly readable
    permission.

  Returns:
    A CloudFilepath object providing the location of the object in various
    formats.
  """
  url = 'gs://%s/%s' % (bucket, remote_path)
  command_and_args = ['cp']
  if publicly_readable:
    command_and_args += ['-a', 'public-read']
  command_and_args += [local_path, url]
  _RunCommand(command_and_args)


def CalculateHash(file_path):
  """Calculates and returns the hash of the file at file_path."""
  sha1 = hashlib.sha1()
  with open(file_path, 'rb') as f:
    while True:
      # Read in 1mb chunks, so it doesn't all have to be loaded into memory.
      chunk = f.read(1024 * 1024)
      if not chunk:
        break
      sha1.update(chunk)
  return sha1.hexdigest()