#!/usr/bin/env python

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

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


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

    Assumes that it is run from trusty/prebuilts/aosp.
    Remove the pristine directory before use.

    Inputs, (specifically adb and qemu_trusty_arm64-target-files-buildid.zip)
    must be placed in the pristine directory before running the script.
    """
    pristine_dir = "pristine"
    adb_source_path = pristine_dir + "/adb"
    # Make sure the relevant files are here - currently just adb + target files.
    if not os.path.isfile(adb_source_path):
        print "ADB missing from prebuilt"
        sys.exit(1)

    # Make sure we have exactly one target zip
    target_files_search = pristine_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)

    # Unpack the target_files archive
    target_files_path = "target"
    shutil.rmtree(target_files_path, ignore_errors=True)
    zipfile.ZipFile(target_files_zip_glob[0]).extractall(target_files_path)

    # Build our directory structure
    base = "android"
    # Clean the unpacked build directory if it already exists
    shutil.rmtree(base, ignore_errors=True)
    host_bin_path = base + "/out/host/linux-x86/bin"
    image_path = base + "/out/target/product/trusty"
    test_path = base + "/out/target/product/trusty/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 = target_files_path + "/IMAGES"
    shutil.move(image_source_path + "/system.img", image_path)
    shutil.move(image_source_path + "/vendor.img", image_path)
    shutil.move(image_source_path + "/userdata.img", image_path)

    # Install test apps
    test_source_path = target_files_path + "/DATA/nativetest64"
    shutil.move(test_source_path, test_path)

    # Clean up unpacked but unneeded files
    shutil.rmtree(target_files_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()