#!/usr/bin/env python3

# Copyright (C) 2023 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Script for running Android Gerrit-based Mobly tests locally.

Example:
    - Run a test module.
    local_mobly_runner.py -m my_test_module

    - Run a test module. Build the module and install test APKs before running the test.
    local_mobly_runner.py -m my_test_module -b -i

    - Run a test module with specific Android devices.
    local_mobly_runner.py -m my_test_module -s DEV00001,DEV00002

    - Run a list of zipped Mobly packages (built from `python_test_host`)
    local_mobly_runner.py -p test_pkg1,test_pkg2,test_pkg3

Please run `local_mobly_runner.py -h` for a full list of options.
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
from typing import List, Optional, Tuple
import zipfile

_LOCAL_SETUP_INSTRUCTIONS = (
    '\n\tcd <repo_root>; set -a; source build/envsetup.sh; set +a; lunch'
    ' <target>'
)

_tempdirs = []
_tempfiles = []


def _padded_print(line: str) -> None:
    print(f'\n-----{line}-----\n')


def _parse_args() -> argparse.Namespace:
    """Parses command line args."""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=__doc__)
    group1 = parser.add_mutually_exclusive_group(required=True)
    group1.add_argument(
        '-m', '--module', help='The Android build module of the test to run.'
    )
    group1.add_argument(
        '-p', '--packages',
        help='A comma-delimited list of test packages to run.'
    )
    group1.add_argument(
        '-t',
        '--test_paths',
        help=(
            'A comma-delimited list of test paths to run directly. Implies '
            'the --novenv option.'
        ),
    )
    parser.add_argument(
        '-b',
        '--build',
        action='store_true',
        help='Build/rebuild the specified module. Requires the -m option.',
    )
    parser.add_argument(
        '-i',
        '--install_apks',
        action='store_true',
        help=(
            'Install all APKs associated with the module to all specified'
            ' devices. Requires the -m or -p options.'
        ),
    )
    group2 = parser.add_mutually_exclusive_group()
    group2.add_argument(
        '-s',
        '--serials',
        help=(
            'Specify the devices to test with a comma-delimited list of device '
            'serials.'
        ),
    )
    group2.add_argument(
        '-c', '--config', help='Provide a custom Mobly config for the test.'
    )
    parser.add_argument('-lp', '--log_path',
                        help='Specify a path to store logs.')
    parser.add_argument(
        '--novenv',
        action='store_true',
        help=(
            "Run directly in the host's system Python, without setting up a "
            'virtualenv.'
        ),
    )
    args = parser.parse_args()
    if args.build and not args.module:
        parser.error('Option --build requires --module to be specified.')
    if args.install_apks and not (args.module or args.packages):
        parser.error('Option --install_apks requires --module or --packages.')

    args.novenv = args.novenv or (args.test_paths is not None)
    return args


def _build_module(module: str) -> None:
    """Builds the specified module."""
    _padded_print(f'Building test module {module}.')
    try:
        subprocess.check_call(f'm -j {module}', shell=True,
                              executable='/bin/bash')
    except subprocess.CalledProcessError as e:
        if e.returncode == 127:
            # `m` command not found
            print(
                '`m` command not found. Please set up your local environment '
                f'with {_LOCAL_SETUP_INSTRUCTIONS}.'
            )
        else:
            print(f'Failed to build module {module}.')
        exit(1)


def _get_module_artifacts(module: str) -> List[str]:
    """Return the list of artifacts generated from a module."""
    try:
        outmod_paths = (
            subprocess.check_output(
                f'outmod {module}', shell=True, executable='/bin/bash'
            )
            .decode('utf-8')
            .splitlines()
        )
    except subprocess.CalledProcessError as e:
        if e.returncode == 127:
            # `outmod` command not found
            print(
                '`outmod` command not found. Please set up your local '
                f'environment with {_LOCAL_SETUP_INSTRUCTIONS}.'
            )
        if str(e.output).startswith('Could not find module'):
            print(
                f'Cannot find the build output of module {module}. Ensure that '
                'the module list is up-to-date with `refreshmod`.'
            )
        exit(1)

    for path in outmod_paths:
        if not os.path.isfile(path):
            print(
                f'Declared file {path} does not exist. Please build your '
                'module with the -b option.'
            )
            exit(1)

    return outmod_paths


def _resolve_test_resources(
        args: argparse.Namespace,
) -> Tuple[List[str], List[str], List[str]]:
    """Resolve test resources from the given test module or package.

    Args:
      args: Parsed command-line args.

    Returns:
      Tuple of (mobly_bins, requirement_files, test_apks).
    """
    mobly_bins = []
    requirements_files = []
    test_apks = []
    if args.test_paths:
        mobly_bins.extend(args.test_paths.split(','))
    elif args.module:
        for path in _get_module_artifacts(args.module):
            if path.endswith(args.module):
                mobly_bins.append(path)
            if path.endswith('requirements.txt'):
                requirements_files.append(path)
            if path.endswith('.apk'):
                test_apks.append(path)
    elif args.packages:
        unzip_root = tempfile.mkdtemp(prefix='mobly_unzip_')
        _tempdirs.append(unzip_root)
        for package in args.packages.split(','):
            mobly_bins.append(package)
            unzip_dir = os.path.join(unzip_root, os.path.basename(package))
            print(f'Unzipping test package {package} to {unzip_dir}.')
            os.makedirs(unzip_dir)
            with zipfile.ZipFile(package) as zf:
                zf.extractall(unzip_dir)
            for path in os.listdir(unzip_dir):
                path = os.path.join(unzip_dir, path)
                if path.endswith('requirements.txt'):
                    requirements_files.append(path)
                if path.endswith('.apk'):
                    test_apks.append(path)
    else:
        print('No tests specified. Aborting.')
        exit(1)
    return mobly_bins, requirements_files, test_apks


def _setup_virtualenv(requirements_files: List[str]) -> str:
    """Creates a virtualenv and install dependencies into it.

    Args:
      requirements_files: List of paths of requirements.txt files.

    Returns:
      Path to the virtualenv's Python interpreter.
    """
    if not requirements_files:
        print('No requirements.txt file found. Aborting.')
        exit(1)

    venv_dir = tempfile.mkdtemp(prefix='venv_')
    _padded_print(f'Creating virtualenv at {venv_dir}.')
    subprocess.check_call([sys.executable, '-m', 'venv', venv_dir])
    _tempdirs.append(venv_dir)
    venv_executable = os.path.join(venv_dir, 'bin/python3')

    # Install requirements
    for requirements_file in requirements_files:
        print(f'Installing dependencies from {requirements_file}.')
        subprocess.check_call(
            [venv_executable, '-m', 'pip', 'install', '-r', requirements_file]
        )
    return venv_executable


def _install_apks(
        apks: List[str],
        serials: Optional[List[str]] = None,
) -> None:
    """Installs given APKS to specified devices.

    If no serials specified, installs APKs on all attached devices.

    Args:
      apks: List of paths to APKs.
      serials: List of device serials.
    """
    _padded_print('Installing test APKs.')
    if not serials:
        serials = (
            subprocess.check_output(
                'adb devices | tail -n +2 | cut -f 1', shell=True
            )
            .decode('utf-8')
            .strip()
            .splitlines()
        )
    for apk in apks:
        for serial in serials:
            print(f'Installing {apk} on device {serial}.')
            subprocess.check_call(
                ['adb', '-s', serial, 'install', '-r', '-g', apk]
            )


def _generate_mobly_config(serials: Optional[List[str]] = None) -> str:
    """Generates a Mobly config for the provided device serials.

    If no serials specified, generate a wildcard config (test loads all attached
    devices).

    Args:
      serials: List of device serials.

    Returns:
      Path to the generated config.
    """
    config = {
        'TestBeds': [{
            'Name': 'LocalTestBed',
            'Controllers': {
                'AndroidDevice': serials if serials else '*',
            },
        }]
    }
    _, config_path = tempfile.mkstemp(prefix='mobly_config_')
    _padded_print(f'Generating Mobly config at {config_path}.')
    with open(config_path, 'w') as f:
        json.dump(config, f)
    _tempfiles.append(config_path)
    return config_path


def _run_mobly_tests(
        python_executable: str,
        mobly_bins: List[str],
        config: str,
        log_path: Optional[str] = None,
) -> None:
    """Runs the Mobly tests with the specified binary and config."""
    env = os.environ.copy()
    for mobly_bin in mobly_bins:
        bin_name = os.path.basename(mobly_bin)
        if log_path:
            env['MOBLY_LOGPATH'] = os.path.join(log_path, bin_name)
        cmd = [python_executable, mobly_bin, '-c', config]
        _padded_print(f'Running Mobly test {bin_name}.')
        print(f'Command: {cmd}\n')
        subprocess.run(cmd, env=env)


def _clean_up() -> None:
    """Cleans up temporary directories and files."""
    _padded_print('Cleaning up temporary directories/files.')
    for td in _tempdirs:
        shutil.rmtree(td, ignore_errors=True)
    _tempdirs.clear()
    for tf in _tempfiles:
        os.remove(tf)
    _tempfiles.clear()


def main() -> None:
    args = _parse_args()

    # Build the test module if requested by user
    if args.build:
        _build_module(args.module)

    serials = args.serials.split(',') if args.serials else None

    # Resolve test resources
    mobly_bins, requirements_files, test_apks = _resolve_test_resources(args)

    # Install test APKs, if necessary
    if args.install_apks:
        _install_apks(test_apks, serials)

    # Set up the Python virtualenv, if necessary
    python_executable = (
        sys.executable if args.novenv else _setup_virtualenv(requirements_files)
    )

    # Generate the Mobly config, if necessary
    config = args.config or _generate_mobly_config(serials)

    # Run the tests
    _run_mobly_tests(python_executable, mobly_bins, config, args.log_path)

    # Clean up temporary dirs/files
    _clean_up()


if __name__ == '__main__':
    main()