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

"""Utility script to install APKs from the command line quickly."""

import argparse
import glob
import logging
import os
import sys

import devil_chromium
from devil.android import apk_helper
from devil.android import device_denylist
from devil.android import device_errors
from devil.android import device_utils
from devil.utils import run_tests_helper
from pylib import constants


def main():
    parser = argparse.ArgumentParser()

    apk_group = parser.add_mutually_exclusive_group(required=True)
    apk_group.add_argument(
        '--apk',
        dest='apk_name',
        help='DEPRECATED The name of the apk containing the'
        ' application (with the .apk extension).',
    )
    apk_group.add_argument(
        'apk_path', nargs='?', help='The path to the APK to install.'
    )

    # TODO(jbudorick): Remove once no clients pass --apk_package
    parser.add_argument('--apk_package', help='DEPRECATED unused')
    parser.add_argument(
        '--split',
        action='append',
        dest='splits',
        help='A glob matching the apk splits. Can be specified multiple times.',
    )
    parser.add_argument(
        '--keep_data',
        action='store_true',
        default=False,
        help='Keep the package data when installing the application.',
    )
    parser.add_argument(
        '--debug',
        action='store_const',
        const='Debug',
        dest='build_type',
        default=os.environ.get('BUILDTYPE', 'Debug'),
        help='If set, run test suites under out/Debug. '
        'Default is env var BUILDTYPE or Debug',
    )
    parser.add_argument(
        '--release',
        action='store_const',
        const='Release',
        dest='build_type',
        help='If set, run test suites under out/Release. '
        'Default is env var BUILDTYPE or Debug.',
    )
    parser.add_argument(
        '-d',
        '--device',
        dest='devices',
        action='append',
        default=[],
        help='Target device for apk to install on. Enter multiple'
        ' times for multiple devices.',
    )
    parser.add_argument(
        '--adb-path',
        type=os.path.abspath,
        help='Absolute path to the adb binary to use.',
    )
    parser.add_argument('--denylist-file', help='Device denylist JSON file.')
    parser.add_argument(
        '-v',
        '--verbose',
        action='count',
        help='Enable verbose logging.',
        default=0,
    )
    parser.add_argument(
        '--downgrade',
        action='store_true',
        help='If set, allows downgrading of apk.',
    )
    parser.add_argument(
        '--timeout',
        type=int,
        default=device_utils.DeviceUtils.INSTALL_DEFAULT_TIMEOUT,
        help='Seconds to wait for APK installation. (default: %(default)s)',
    )

    args = parser.parse_args()

    run_tests_helper.SetLogLevel(args.verbose)
    constants.SetBuildType(args.build_type)

    devil_chromium.Initialize(
        output_directory=constants.GetOutDirectory(), adb_path=args.adb_path
    )

    apk = args.apk_path or args.apk_name
    if not apk.endswith('.apk'):
        apk += '.apk'
    if not os.path.exists(apk):
        apk = os.path.join(constants.GetOutDirectory(), 'apks', apk)
        if not os.path.exists(apk):
            parser.error('%s not found.' % apk)

    if args.splits:
        splits = []
        base_apk_package = apk_helper.ApkHelper(apk).GetPackageName()
        for split_glob in args.splits:
            apks = [f for f in glob.glob(split_glob) if f.endswith('.apk')]
            if not apks:
                logging.warning('No apks matched for %s.', split_glob)
            for f in apks:
                helper = apk_helper.ApkHelper(f)
                if (
                    helper.GetPackageName() == base_apk_package
                    and helper.GetSplitName()
                ):
                    splits.append(f)

    denylist = (
        device_denylist.Denylist(args.denylist_file)
        if args.denylist_file
        else None
    )
    devices = device_utils.DeviceUtils.HealthyDevices(
        denylist=denylist, device_arg=args.devices
    )

    def denylisting_install(device):
        try:
            if args.splits:
                device.InstallSplitApk(
                    apk,
                    splits,
                    reinstall=args.keep_data,
                    allow_downgrade=args.downgrade,
                )
            else:
                device.Install(
                    apk,
                    reinstall=args.keep_data,
                    allow_downgrade=args.downgrade,
                    timeout=args.timeout,
                )
        except (
            device_errors.CommandFailedError,
            device_errors.DeviceUnreachableError,
        ):
            logging.exception('Failed to install %s', apk)
            if denylist:
                denylist.Extend([str(device)], reason='install_failure')
                logging.warning('Denylisting %s', str(device))
        except device_errors.CommandTimeoutError:
            logging.exception('Timed out while installing %s', apk)
            if denylist:
                denylist.Extend([str(device)], reason='install_timeout')
                logging.warning('Denylisting %s', str(device))

    device_utils.DeviceUtils.parallel(devices).pMap(denylisting_install)


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