#!/usr/bin/env python3
#
# Copyright (C) 2024 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.
#
"""Downloads etm modules from the build server."""

import argparse

import logging
import os
from pathlib import Path
import shutil
import stat
import textwrap


ETM_MODULES = ['coresight.ko', 'coresight-etm4x.ko', 'coresight-funnel.ko', 'coresight-tmc.ko']


def logger():
    """Returns the main logger for this module."""
    return logging.getLogger(__name__)


def check_call(cmd):
    """Proxy for subprocess.check_call with logging."""
    import subprocess
    logger().debug('check_call `%s`', ' '.join(cmd))
    subprocess.check_call(cmd)


def fetch_artifact(branch: str, build: str, target: str, name: str) -> None:
    """Fetches and artifact from the build server."""
    logger().info('Fetching %s from %s %s (artifacts matching %s)', build,
                  target, branch, name)
    fetch_artifact_path = '/google/data/ro/projects/android/fetch_artifact'
    cmd = [fetch_artifact_path, '--branch', branch, '--target', target,
           '--bid', build, name]
    check_call(cmd)


def get_etm_module_target(branch: str) -> str:
    if branch == 'kernel-pixel-android16-gs-pixel-6.12':
        return 'muzel'
    return 'etm_kernel_modules'


def fetch_etm_modules(branch: str, build: str, install_dir: str) -> None:
    """Installs the device specific components of the release."""
    install_dir = Path(install_dir)
    install_dir.mkdir(exist_ok=True)
    target = get_etm_module_target(branch)
    for name in ETM_MODULES:
        fetch_artifact(branch, build, target, name)
        dest_file = install_dir / name
        if dest_file.is_file():
            dest_file.unlink()
        shutil.move(name, dest_file)


def get_args():
    """Parses and returns command line arguments."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('-b', '--branch', required=True,
                        choices=['kernel-pixel-android16-gs-pixel-6.12',
                                 'kernel-pixel-android15-gs-pixel-6.6',
                                 'kernel-android14-gs-pixel-6.1',
                                 'kernel-android14-gs-pixel-5.15-24Q3'],
                        help='Branch to pull build from.')
    parser.add_argument('--build', required=True, help='Build number to pull.')
    parser.add_argument('-o', dest='install_dir', required=True,
                        help='Directory to store etm modules')
    parser.add_argument('-v', '--verbose', action='count', default=0,
                        help='Increase output verbosity.')
    return parser.parse_args()


def main():
    """Program entry point."""
    args = get_args()
    verbose_map = (logging.WARNING, logging.INFO, logging.DEBUG)
    verbosity = args.verbose
    if verbosity > 2:
        verbosity = 2
    logging.basicConfig(level=verbose_map[verbosity])
    logging.debug(f'args={args}')
    fetch_etm_modules(args.branch, args.build, args.install_dir)


if __name__ == '__main__':
    main()