#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# Copyright (C) 2024 by Arm Limited.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
# TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''
This script is a utility installer that will download the latest release from
the Arm Artifactory download site, and unpack it. It has no dependencies on
third-party Python modules or external native software binaries.

The script exposes three run subcommands:

  - list
  - download
  - install

Run subcommands with --help for full documentation:

    get-streamline-cli.py [subcommand] --help

List
====

    get-streamline-cli.py list

When run with the "list" subcommand, this tool will list the available versions
on the Arm download server.

Download
========

    get-streamline-cli.py download [args]

When run with the "download" subcommand this script will download a release,
but will not install it. By default the script will download the latest
release, but a specific version can be manually specified using --tool-version.

By default, the tool will error if there is a local file matching the name of
the download. Run with --overwrite to allow the tool to overwrite local files.

Install
=======

    get-streamline-cli.py install [args]

When run with the "install" subcommand this script will download and extract
a release in to a local directory. By default the script will download the
latest release and extract it into the current working directory, but both
options can be manually specified by the user.

By default, the tool will error if there is a local file matching the name of
the download, or a local directory matching the extracted tool directory. Run
with --overwrite to allow the tool to overwrite local files.
'''

import argparse
import datetime
import logging
import os
import pathlib
import re
import shutil
import sys
import tarfile
import tempfile
import urllib.error
import urllib.request
from typing import Optional, List, Tuple

VERSION = 'v1.1.0'

DESCRIPTION = f'''
Streamline CLI tools downloader

A utility tool that can download and unpack releases of the Streamline
command line profiling tools from the Arm tools download server.

Use: get-streamline-cli.py list
    Lists available versions on the server.

Use: get-streamline-cli.py download [--tool-version VER]
    Downloads a tool release from the server. If no version is specified
    the latest version is downloaded.

Use: get-streamline-cli.py install [--tool-version VER] [--install-dir DIR]
    Downloads and installs a tool release. If no version is specified the
    latest version is downloaded. If no install directory is specified the
    release is unpacked to the current working directory.
'''

REPO = 'https://artifacts.tools.arm.com/arm-performance-studio'
PRODUCT = 'Streamline_CLI_Tools'
DIR_INSIDE_TAR = 'streamline_cli_tools'  # TODO: Programmatic extract from tar
VERBOSE = False

logging.basicConfig(format='%(message)s')
log = logging.getLogger(__name__)


def get_copy_yr() -> str:
    '''
    Return the year range for the header string.
    '''
    startyear = 2024
    endyear = datetime.datetime.now().year

    year = f"{startyear}"
    if startyear != endyear:
        year = f"{year}-{endyear}"

    return year


def download(artifact: str, output_file: str) -> bool:
    '''
    Download a file from the Arm public tools Artifactory download site.

    Args:
        artifact: The URL to download.
        output_file: The local file to create.

    Return:
        True on success, False otherwise.
    '''
    try:
        urllib.request.urlretrieve(artifact, output_file)
    except urllib.error.HTTPError as ex:
        log.info(f'  - HTTP Error: {ex.url} ({ex.code})')
        return False

    return True


def download_release_listing() -> Optional[List[Tuple[str, str]]]:
    '''
    Download sorted list of all published releases.

    Return:
        The list of releases (file, version) or None if no releases found.
    '''
    log.info('\nFinding available releases ...')

    # Download the file listing page from Artifactory
    download_url = f'{REPO}/{PRODUCT}/'
    tmp_file = tempfile.NamedTemporaryFile()

    download_ok = download(download_url, tmp_file.name)

    with open(tmp_file.name, 'r', encoding='utf-8') as file_handle:
        report = file_handle.read()

    # Error check after we read the file so we trigger the TempFile cleanup
    if not download_ok:
        log.error('ERROR: Download of release listing failed')
        return None

    # Extract the tool versions from the directory listing file
    pattern = re.compile(r'^.*<a href="'
                         r'(Arm_Streamline_CLI_Tools_(.*?)_linux_arm64\.tgz)'
                         r'">.*')

    found_releases = []
    for line in report.splitlines():
        match = pattern.match(line)
        if not match:
            continue

        file_name = match.group(1)
        version = match.group(2)

        try:
            versioni = [int(x) for x in version.split('.')]
            found_releases.append((versioni, version, file_name))

        # Could not convert to integer, so skip release
        except ValueError:
            log.warning(f'WARNING: Ignoring {file_name}, unknown version ...')
            continue

    found_releases.sort(key=lambda x: x[0])

    clean_releases = []
    for _, release_version, release_path in found_releases:
        clean_releases.append((release_path, release_version))

    if not clean_releases:
        log.error('ERROR: No releases found on server')
        return None

    return clean_releases


def determine_selected_release(
        version: Optional[str]) -> Optional[Tuple[str, str]]:
    '''
    Determine which release to download.

    Args:
        version: User-specified version, or None to select latest.

    Return:
        The release version and file name on the remote server, or None if no
        releases were found or if the user version was not found.
    '''
    release_list = download_release_listing()
    if not release_list:
        return None

    for release_path, _ in release_list:
        log.info(f'  - Found {release_path}')

    # No user version, so return the latest
    if not version:
        release = release_list[-1]
        log.info(f'  - Selected {release[0]}')
        return release

    # User version, so return on exact match
    for release_path, release_version in release_list:
        if release_version == version:
            log.info(f'  - Selected {release_path}')
            return (release_path, release_version)

    log.error(f'ERROR: No release available for version {version}')
    return None


def download_selected_release(release: str, overwrite: bool) -> bool:
    '''
    Download the latest release.

    Args:
        release: The release file name.
        overwrite: True if we should overwrite the local file.

    Return:
        True on success, False otherwise.
    '''
    log.info(f'\nDownloading {release} ...')

    if os.path.exists(release):
        if not overwrite:
            log.error(f'ERROR: Local file {release} already exists')
            return False

        os.remove(release)

    # Download the file listing page from Artifactory
    download_url = f'{REPO}/{PRODUCT}/{release}'
    download_ok = download(download_url, release)
    if not download_ok:
        log.error(f'ERROR: Download of {release} failed')

    return download_ok


def extract_selected_release(
        release: str, install_dir: str, overwrite: bool) -> bool:
    '''
    Extract the selected release.

    Args:
        release: The release file name.
        install_dir: The directory to extract the tool archive into.
        overwrite: True if we should overwrite the local directory.

    Return:
        True on success, False otherwise.
    '''
    log.info(f'\nExtracting {release} ...')

    # Determine the root directory name in the tar file as it may change to
    # add versioning in a future release
    keyfile = '/sl-record'
    with tarfile.open(release) as file_handle:
        for member in file_handle.getmembers():
            if member.name.endswith(keyfile):
                break
        else:
            log.error(f'ERROR: Cannot identify tar directory layout')
            return False

    # Only handle the case where the directory layout is <dir>/bin/sl-record
    slrecord_path = pathlib.Path(member.name)
    if len(slrecord_path.parts) != 3:
        log.error(f'ERROR: Cannot identify tar directory layout')
        return False

    tar_root_dir = slrecord_path.parts[0]
    install_dir = os.path.abspath(install_dir)
    unpack_dir = os.path.join(install_dir, tar_root_dir)

    # Check we can unpack
    if os.path.exists(unpack_dir):
        if overwrite:
            log.debug('  - Local file: overwriting directory - OK')
            shutil.rmtree(unpack_dir)
        else:
            log.error(f'ERROR: Local directory {unpack_dir} already exists')
            return False
    else:
        log.debug('  - Local file: new directory - OK')

    # Ensure precursor directories exist
    os.makedirs(install_dir, exist_ok=True)

    # Use Python 3.12-style data filters for security if they are available
    try:
        if hasattr(tarfile, 'data_filter'):
            with tarfile.open(release) as file_handle:
                file_handle.extractall(path=install_dir, filter='data')
        else:
            with tarfile.open(release) as file_handle:
                file_handle.extractall(path=install_dir)
    except tarfile.TarError:
        log.error(f'ERROR: Extract of {release} failed')
        return False

    return True


def download_release(args: argparse.Namespace) -> Optional[str]:
    '''
    Download a selected release.

    Args:
        args: Command line arguments

    Return:
        Release file path on the local machine.
    '''
    # Determine the file name of the release on the server
    release = determine_selected_release(args.tool_version)
    if not release:
        return None

    # Download the release
    try:
        download_ok = download_selected_release(release[0], args.overwrite)
        if not download_ok:
            return None
    except PermissionError:
        log.error('ERROR: No file permission to download to this location')
        return None

    return release[0]


def list_main(args: argparse.Namespace) -> int:
    '''
    Main function for the list operation mode.

    Args:
        args: Command line arguments

    Return:
        Process exit code, zero on success, non-zero on error.
    '''
    release_list = download_release_listing()
    if not release_list:
        return 1

    print('Available versions:')
    for release in release_list:
        print(f'  - {release[0]} (version = {release[1]})')

    return 0


def download_main(args: argparse.Namespace) -> int:
    '''
    Main function for the download operation mode.

    Args:
        args: Command line arguments

    Return:
        Process exit code, zero on success, non-zero on error.
    '''
    # Download the release from the server ...
    release = download_release(args)
    if not release:
        return 1

    print(f'SUCCESS: Downloaded Streamline CLI Tools {release}')
    return 0


def install_main(args: argparse.Namespace) -> int:
    '''
    Main function for the install operation mode.

    Args:
        args: Command line arguments

    Return:
        Process exit code, zero on success, non-zero on error.
    '''
    local_download = download_release(args)
    if not local_download:
        return 1

    # Extract the release
    try:
        extract_ok = extract_selected_release(
            local_download, args.install_dir, args.overwrite)
        if not extract_ok:
            return 1
    except PermissionError:
        log.error('ERROR: No file permission to install to this location')
        return 1
    finally:
        # Remove the tar file
        os.remove(local_download)

    release = os.path.basename(local_download)
    print(f'SUCCESS: Installed Streamline CLI Tools {release}')
    return 0


def parse_cli() -> argparse.Namespace:
    '''
    Parse the command line.

    Return:
        The argparse results object.
    '''
    parser = argparse.ArgumentParser(
        prog='get-streamline-cli',
        description=DESCRIPTION,
        formatter_class=argparse.RawTextHelpFormatter)

    subparsers = parser.add_subparsers(
        title='subcommands',
        description='valid subcommands',
        help='sub-command help',
        required=True)

    # List mode
    parser_list = subparsers.add_parser(
        'list', help='list available releases')

    parser_list.set_defaults(func=list_main)

    # Download mode
    parser_download = subparsers.add_parser(
        'download', help='download a release')

    parser_download.set_defaults(func=download_main)

    parser_download.add_argument(
        '--overwrite', action='store_true', default=False,
        help='overwrite local download if it already exists')

    parser_download.add_argument(
        '--tool-version', default=None,
        help='the tool version to download')

    parser_download.add_argument(
        '--verbose', action='store_true', default=False,
        help='enable verbose logging')

    # Install mode
    parser_install = subparsers.add_parser(
        'install', help='download and install a release')

    parser_install.set_defaults(func=install_main)

    parser_install.add_argument(
        '--install-dir', default='.',
        help='the tool installation directory')

    parser_install.add_argument(
        '--tool-version', default=None,
        help='the tool version to install')

    parser_install.add_argument(
        '--overwrite', action='store_true', default=False,
        help='overwrite local download and install if they already exist')

    parser_install.add_argument(
        '--verbose', action='store_true', default=False,
        help='enable verbose logging')

    return parser.parse_args()


def main() -> int:
    '''
    The script main function.

    Return:
        Process exit code.
    '''
    print(f"get-streamline-cli {VERSION}")
    print(f"Copyright (c) {get_copy_yr()} Arm Limited. All rights reserved.\n")

    args = parse_cli()

    # Configure logging
    if getattr(args, 'verbose', False):
        log.setLevel(logging.INFO)

    return args.func(args)


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