#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright 2019 The ChromiumOS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""A crontab script to delete night test data."""

from __future__ import print_function

__author__ = 'shenhan@google.com (Han Shen)'

import argparse
import datetime
import os
import re
import shutil
import stat
import sys
import time
import traceback
from pathlib import Path
from typing import Callable

from cros_utils import command_executer, constants, misc

DIR_BY_WEEKDAY = ('Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun')
NIGHTLY_TESTS_WORKSPACE = os.path.join(constants.CROSTC_WORKSPACE,
                                       'nightly-tests')


def CleanNumberedDir(s, dry_run=False):
  """Deleted directories under each dated_dir."""
  chromeos_dirs = [
      os.path.join(s, x) for x in os.listdir(s)
      if misc.IsChromeOsTree(os.path.join(s, x))
  ]
  ce = command_executer.GetCommandExecuter(log_level='none')
  all_succeeded = True
  for cd in chromeos_dirs:
    if misc.DeleteChromeOsTree(cd, dry_run=dry_run):
      print(f'Successfully removed chromeos tree {cd!r}.')
    else:
      all_succeeded = False
      print(f'Failed to remove chromeos tree {cd!r}, please check.')

  if not all_succeeded:
    print('Failed to delete at least one chromeos tree, please check.')
    return False

  ## Now delete the numbered dir Before forcibly removing the directory, just
  ## check 's' to make sure it matches the expected pattern.  A valid dir to be
  ## removed must be '/usr/local/google/crostc/(SUN|MON|TUE...|SAT)'.
  valid_dir_pattern = ('^' + NIGHTLY_TESTS_WORKSPACE + '/(' +
                       '|'.join(DIR_BY_WEEKDAY) + ')')
  if not re.search(valid_dir_pattern, s):
    print(f'Trying to delete an invalid dir {s!r} (must match '
          f'{valid_dir_pattern!r}), please check.')
    return False

  cmd = f'rm -fr {s}'
  if dry_run:
    print(cmd)
  else:
    if ce.RunCommand(cmd, print_to_console=False, terminated_timeout=480) == 0:
      print(f'Successfully removed {s!r}.')
    else:
      all_succeeded = False
      print(f'Failed to remove {s!r}, please check.')
  return all_succeeded


def CleanDatedDir(dated_dir, dry_run=False):
  # List subdirs under dir
  subdirs = [
      os.path.join(dated_dir, x) for x in os.listdir(dated_dir)
      if os.path.isdir(os.path.join(dated_dir, x))
  ]
  all_succeeded = True
  for s in subdirs:
    if not CleanNumberedDir(s, dry_run):
      all_succeeded = False
  return all_succeeded


def ProcessArguments(argv):
  """Process arguments."""
  parser = argparse.ArgumentParser(
      description='Automatically delete nightly test data directories.',
      usage='auto_delete_nightly_test_data.py options')
  parser.add_argument('-d',
                      '--dry_run',
                      dest='dry_run',
                      default=False,
                      action='store_true',
                      help='Only print command line, do not execute anything.')
  parser.add_argument('--days_to_preserve',
                      dest='days_to_preserve',
                      default=3,
                      help=('Specify the number of days (not including today),'
                            ' test data generated on these days will *NOT* be '
                            'deleted. Defaults to 3.'))
  options = parser.parse_args(argv)
  return options


def RemoveAllSubdirsMatchingPredicate(
    base_dir: Path, days_to_preserve: int, dry_run: bool,
    is_name_removal_worthy: Callable[[str], bool]) -> bool:
  """Removes all subdirs of base_dir that match the given predicate."""
  secs_to_preserve = 60 * 60 * 24 * days_to_preserve
  now = time.time()
  remove_older_than_time = now - secs_to_preserve

  had_errors = False
  for file in base_dir.iterdir():
    if not is_name_removal_worthy(file.name):
      continue

    try:
      # Take the stat here and use that later, so we only need to check for a
      # nonexistent file once.
      st = file.stat()
    except FileNotFoundError:
      # This was deleted while were checking; ignore it.
      continue

    if not stat.S_ISDIR(st.st_mode):
      continue

    if st.st_atime >= remove_older_than_time:
      continue

    if dry_run:
      print(f'Would remove {file}')
      continue

    this_iteration_had_errors = False

    def OnError(_func, path_name, excinfo):
      nonlocal this_iteration_had_errors
      this_iteration_had_errors = True
      print(f'Failed removing path at {path_name}; traceback:')
      traceback.print_exception(*excinfo)

    shutil.rmtree(file, onerror=OnError)

    # Some errors can be other processes racing with us to delete things. Don't
    # count those as an error which we complain loudly about.
    if this_iteration_had_errors:
      if file.exists():
        had_errors = True
      else:
        print(f'Discarding removal errors for {file}; dir was still removed.')

  return 1 if had_errors else 0


def IsChromeOsTmpDeletionCandidate(file_name: str):
  """Returns whether the given basename can be deleted from a chroot's /tmp."""
  name_prefixes = (
      'test_that_',
      'cros-update',
      'CrAU_temp_data',
  )
  if any(file_name.startswith(x) for x in name_prefixes):
    return True
  # Remove files that look like `tmpABCDEFGHI`.
  return len(file_name) == 9 and file_name.startswith('tmp')


def CleanChromeOsTmpFiles(chroot_tmp: str, days_to_preserve: int,
                          dry_run: bool) -> int:
  # Clean chroot/tmp/test_that_* and chroot/tmp/tmpxxxxxx, that were last
  # accessed more than specified time ago.
  return RemoveAllSubdirsMatchingPredicate(
      Path(chroot_tmp),
      days_to_preserve,
      dry_run,
      IsChromeOsTmpDeletionCandidate,
  )


def CleanChromeOsImageFiles(chroot_tmp, subdir_suffix, days_to_preserve,
                            dry_run):
  # Clean files that were last accessed more than the specified time.
  seconds_delta = days_to_preserve * 24 * 3600
  now = time.time()
  errors = 0

  for tmp_dir in os.listdir(chroot_tmp):
    # Directory under /tmp
    tmp_dir = os.path.join(chroot_tmp, tmp_dir)
    if tmp_dir.endswith(subdir_suffix):
      # Tmp directory which ends with subdir_suffix.
      for subdir in os.listdir(tmp_dir):
        # Subdirectories targeted for deletion.
        subdir_path = os.path.join(tmp_dir, subdir)
        if now - os.path.getatime(subdir_path) > seconds_delta:
          if dry_run:
            print(f'Will run:\nshutil.rmtree({subdir_path!r})')
          else:
            try:
              shutil.rmtree(subdir_path)
              print('Successfully cleaned chromeos image autotest directories '
                    f'from {subdir_path!r}.')
            except OSError:
              print('Some image autotest directories were not removed from '
                    f'"{subdir_path}".')
              errors += 1

  return errors


def CleanChromeOsTmpAndImages(days_to_preserve=1, dry_run=False):
  """Delete temporaries, images under crostc/chromeos."""
  chromeos_chroot_tmp = os.path.join(constants.CROSTC_WORKSPACE, 'chromeos',
                                     'chroot', 'tmp')
  # Clean files in tmp directory
  rv = CleanChromeOsTmpFiles(chromeos_chroot_tmp, days_to_preserve, dry_run)
  # Clean image files in *-tryjob directories
  rv += CleanChromeOsImageFiles(chromeos_chroot_tmp, '-tryjob',
                                days_to_preserve, dry_run)
  # Clean image files in *-release directories
  rv += CleanChromeOsImageFiles(chromeos_chroot_tmp, '-release',
                                days_to_preserve, dry_run)
  # Clean image files in *-pfq directories
  rv += CleanChromeOsImageFiles(chromeos_chroot_tmp, '-pfq', days_to_preserve,
                                dry_run)
  # Clean image files in *-llvm-next-nightly directories
  rv += CleanChromeOsImageFiles(chromeos_chroot_tmp, '-llvm-next-nightly',
                                days_to_preserve, dry_run)

  return rv


def CleanOldCLs(days_to_preserve='1', dry_run=False):
  """Abandon old CLs created by automation tooling."""
  ce = command_executer.GetCommandExecuter()
  chromeos_root = os.path.join(constants.CROSTC_WORKSPACE, 'chromeos')
  # Find Old CLs.
  old_cls_cmd = ('gerrit --raw search "owner:me status:open age:%sd"' %
                 days_to_preserve)
  _, cls, _ = ce.ChrootRunCommandWOutput(chromeos_root,
                                         old_cls_cmd,
                                         print_to_console=False)
  # Convert any whitespaces to spaces.
  cls = ' '.join(cls.split())
  if not cls:
    return 0

  abandon_cls_cmd = ('gerrit abandon %s' % cls)
  if dry_run:
    print('Going to execute: %s' % abandon_cls_cmd)
    return 0

  return ce.ChrootRunCommand(chromeos_root,
                             abandon_cls_cmd,
                             print_to_console=False)


def CleanChromeTelemetryTmpFiles(dry_run: bool) -> int:
  return RemoveAllSubdirsMatchingPredicate(
      Path(constants.CROSTC_WORKSPACE),
      days_to_preserve=1,
      dry_run=dry_run,
      is_name_removal_worthy=lambda x: x.startswith('tmp') and x.endswith(
          'telemetry_Crosperf'),
  )


def Main(argv):
  """Delete nightly test data directories, tmps and test images."""
  options = ProcessArguments(argv)
  # Function 'isoweekday' returns 1(Monday) - 7 (Sunday).
  d = datetime.datetime.today().isoweekday()
  # We go back 1 week, delete from that day till we are
  # options.days_to_preserve away from today.
  s = d - 7
  e = d - int(options.days_to_preserve)
  rv = 0
  for i in range(s + 1, e):
    if i <= 0:
      ## Wrap around if index is negative.  6 is from i + 7 - 1, because
      ## DIR_BY_WEEKDAY starts from 0, while isoweekday is from 1-7.
      dated_dir = DIR_BY_WEEKDAY[i + 6]
    else:
      dated_dir = DIR_BY_WEEKDAY[i - 1]

    rv += 0 if CleanDatedDir(os.path.join(NIGHTLY_TESTS_WORKSPACE, dated_dir),
                             options.dry_run) else 1

  ## Clean temporaries, images under crostc/chromeos
  rv2 = CleanChromeOsTmpAndImages(int(options.days_to_preserve),
                                  options.dry_run)

  # Clean CLs that are not updated in last 2 weeks.
  rv3 = CleanOldCLs('14', options.dry_run)

  # Clean telemetry temporaries from chrome source tree inside chroot.
  rv4 = CleanChromeTelemetryTmpFiles(options.dry_run)

  return rv + rv2 + rv3 + rv4


if __name__ == '__main__':
  retval = Main(sys.argv[1:])
  sys.exit(retval)