#!/bin/sh
"." "`dirname $0`/../../vendor/google/aosp/scripts/envsetup.sh"; "exec" "$PY3" "$0" "$@"

"""Unpacks a prebuilt to look like an Android build directory."""

import argparse
import glob
import itertools
import os
import shutil
import stat
import sys
import tempfile
import zipfile


SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))

def main():
    """Unpacks a prebuilt to look like an Android build directory.

    Inputs, (specifically adb and qemu_trusty_arm64-target-files-buildid.zip)
    must be placed in the pristine directory before running the script, or
    passed in with the --artifacts-dir argument.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-a", "--artifacts-dir",
        type=str,
        default=os.path.join(SCRIPT_DIR, "pristine"),
        help="Path to the directory containing the input artifacts",
    )
    parser.add_argument(
        "-o", "--output-dir",
        type=str,
        default=os.path.join(SCRIPT_DIR, "android"),
        help="Path to the output directory",
    )
    args = parser.parse_args()

    adb_source_path = os.path.join(args.artifacts_dir, "adb")
    # Make sure the relevant files are here - currently just adb + target files.
    if not os.path.isfile(adb_source_path):
        print(f"ADB missing from {adb_source_path}")
        sys.exit(1)

    # Make sure we have exactly one target zip
    target_files_search = os.path.join(args.artifacts_dir,
                                       "qemu_trusty_arm64-target_files-*.zip")
    target_files_zip_glob = glob.glob(target_files_search)
    if not target_files_zip_glob:
        print("No Android target_files found")
        sys.exit(1)
    if len(target_files_zip_glob) > 1:
        print ("Multiple Android target_files archives found. Please clean up "
               "old inputs before updating the prebuilt.")
        sys.exit(1)

    if os.path.exists(args.output_dir):
        print(f"Output directory {args.output_dir} already exists.")
        print("Please remove it manually before running this script.")
        sys.exit(1)

    # Unpack the target_files archive
    with tempfile.TemporaryDirectory() as target_files_path:
        zipfile.ZipFile(target_files_zip_glob[0]).extractall(target_files_path)

        # Build our directory structure
        host_bin_path = os.path.join(args.output_dir, "out", "host", "linux-x86", "bin")
        image_path = os.path.join(args.output_dir, "out", "target", "product", "trusty")
        test_path = os.path.join(image_path, "data")
        os.makedirs(host_bin_path)
        os.makedirs(image_path)
        os.makedirs(test_path)

        # Install adb to the host bin directory
        shutil.copy(adb_source_path, host_bin_path)

        # Install images
        image_source_path = os.path.join(target_files_path, "IMAGES")
        shutil.move(os.path.join(image_source_path, "system.img"), image_path)
        shutil.move(os.path.join(image_source_path, "vendor.img"), image_path)
        shutil.move(os.path.join(image_source_path, "userdata.img"), image_path)

        # Install test apps
        test_source_path = os.path.join(target_files_path, "DATA", "nativetest64")
        shutil.move(test_source_path, test_path)

    # Fix up permissions
    exec_perms = (stat.S_IXUSR | stat.S_IRUSR | stat.S_IXGRP | stat.S_IRGRP
                  | stat.S_IXOTH | stat.S_IROTH)
    for (dirpath, _, filenames) in itertools.chain(os.walk(host_bin_path),
                                                   os.walk(test_path)):
        for filename in filenames:
            os.chmod(os.path.join(dirpath, filename), exec_perms)

    sys.exit(0)


if __name__ == "__main__":
    main()