/*
 * Copyright (C) 2016 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.
 */

package com.android.aetest.tradefed.targetprep;

import static com.android.tradefed.util.BuildTestsZipUtils.getApkFile;

import com.android.tradefed.build.IBuildInfo;
import com.android.tradefed.config.Option;
import com.android.tradefed.config.OptionClass;
import com.android.tradefed.device.DeviceNotAvailableException;
import com.android.tradefed.device.ITestDevice;
import com.android.tradefed.log.LogUtil.CLog;
import com.android.tradefed.targetprep.AltDirBehavior;
import com.android.tradefed.targetprep.BaseTargetPreparer;
import com.android.tradefed.targetprep.BuildError;
import com.android.tradefed.targetprep.ITargetCleaner;
import com.android.tradefed.targetprep.ITargetPreparer;
import com.android.tradefed.targetprep.TargetSetupError;
import com.android.tradefed.targetprep.TestAppInstallSetup;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * A {@link ITargetPreparer} that creates a managed profile in the testing device.
 *
 * <p>This was forked off Android Enterprise team's target preparer, where its method of resolving
 * the test APK was substituted with the platform's implementation. It is also modified to use the
 * standard TradeFed way of disabling a preparer. The two classes are otherwise the same.
 */
@OptionClass(alias = "ae-test-create-managed-profile")
public class AeTestManagedProfileCreator extends BaseTargetPreparer implements ITargetCleaner {
    public static final String MANAGED_PROFILE_USER_ID = "managed_profile_user_id";

    private final TestAppInstallSetup mInstallPreparer = new TestAppInstallSetup();

    @Option(
            name = "profile-owner-component",
            description = "Profile owner component to set; optional")
    private String mProfileOwnerComponent = null;

    @Option(name = "profile-owner-apk", description = "Profile owner apk path; optional")
    private String mProfileOwnerApk = null;

    @Option(
            name = "alt-dir",
            description =
                    "Alternate directory to look for the profile owner apk if the apk is not in "
                            + "the tests zip file. For each alternate dir, will look in //, "
                            + "//data/app, //DATA/app, //DATA/app/apk_name/ and "
                            + "//DATA/priv-app/apk_name/. Can be repeated. Look for apks in "
                            + "last alt-dir first.")
    private List<File> mAltDirs = new ArrayList<>();

    @Option(
            name = "alt-dir-behavior",
            description =
                    "The order of alternate directory to be used "
                            + "when searching for apks to install")
    private AltDirBehavior mAltDirBehavior = AltDirBehavior.FALLBACK;

    @Option(
            name = "test-app-file-name",
            description =
                    "the name of an apk file to be installed in work profile. Can be repeated. Items "
                            + "that are directories will have any APKs contained therein, "
                            + "including subdirectories, grouped by package name and installed.")
    private List<File> mTestFiles = new ArrayList<>();

    @Option(
            name = "test-app-alt-dir",
            description =
                    "Alternate directory to look for the apk if the apk is not in the tests "
                            + "zip file. For each alternate dir, will look in //, //data/app, "
                            + "//DATA/app, //DATA/app/apk_name/ and //DATA/priv-app/apk_name/. "
                            + "Can be repeated. Look for apks in last alt-dir first.")
    private List<File> mAltTestDirs = new ArrayList<>();

    /** UserID for user in managed profile. */
    private int mManagedProfileUserId;

    /** {@inheritDoc} */
    @Override
    public void setUp(ITestDevice device, IBuildInfo buildInfo)
            throws TargetSetupError, DeviceNotAvailableException, BuildError {
        String pmCommand =
                "pm create-user --profileOf 0 --managed "
                        + "TestProfile_"
                        + System.currentTimeMillis();
        String pmCommandOutput = device.executeShellCommand(pmCommand);

        // Extract the id of the new user.
        String[] pmCmdTokens = pmCommandOutput.split("\\s+");
        if (!pmCmdTokens[0].contains("Success:")) {
            throw new TargetSetupError("Managed profile creation failed.");
        }
        mManagedProfileUserId = Integer.parseInt(pmCmdTokens[pmCmdTokens.length - 1]);

        // Start managed profile user.
        device.startUser(mManagedProfileUserId);

        CLog.i(String.format("New user created: %d", mManagedProfileUserId));
        buildInfo.addBuildAttribute(MANAGED_PROFILE_USER_ID, String.valueOf(mManagedProfileUserId));

        // Install required test APKs in the managed profile
        installTestApks(device, buildInfo);

        if (mProfileOwnerComponent != null && mProfileOwnerApk != null) {
            device.installPackageForUser(
                    getApk(device, buildInfo, mProfileOwnerApk), true, mManagedProfileUserId);
            String command =
                    String.format(
                            "dpm set-profile-owner --user %d %s",
                            mManagedProfileUserId, mProfileOwnerComponent);
            String commandOutput = device.executeShellCommand(command);
            String[] cmdTokens = commandOutput.split("\\s+");
            if (!cmdTokens[0].contains("Success:")) {
                throw new RuntimeException(
                        String.format(
                                "Fail to setup %s as profile owner.", mProfileOwnerComponent));
            }

            CLog.i(
                    String.format(
                            "%s was set as profile owner of user %d",
                            mProfileOwnerComponent, mManagedProfileUserId));
        }
        /*
        Commenting this out, as the reboot somehow evaporates the permissions granted to
        managed profile apps during installation. Creation of apps is not being impacted by this.
        */
        // Reboot device to create the apps in managed profile.
        // device.reboot();
        // device.waitForDeviceAvailable();
    }

    /** {@inheritDoc} */
    @Override
    public void tearDown(ITestDevice testDevice, IBuildInfo buildInfo, Throwable throwable)
            throws DeviceNotAvailableException {
        mInstallPreparer.tearDown(testDevice, buildInfo, throwable);
        testDevice.removeUser(mManagedProfileUserId);
    }

    private void installTestApks(ITestDevice device, IBuildInfo buildInfo)
            throws DeviceNotAvailableException, TargetSetupError, BuildError {
        if (mTestFiles.isEmpty()) {
            return;
        }
        CLog.i(
                String.format(
                        "Installing the following test APKs in user %d: \n%s",
                        mManagedProfileUserId, mTestFiles));
        mInstallPreparer.setUserId(mManagedProfileUserId);
        mInstallPreparer.setShouldGrantPermission(true);
        for (File file : mTestFiles) {
            mInstallPreparer.addTestFile(file);
        }
        for (File dir : mAltTestDirs) {
            mInstallPreparer.setAltDir(dir);
        }
        mInstallPreparer.setUp(device, buildInfo);
    }

    /**
     * Get the {@link File} object that refers to the APK to install.
     *
     * <p>This is a replacement of the method with a similar signature in the original class's super
     * class, so that the above code can align as closely to the original class as possible.
     */
    private File getApk(ITestDevice device, IBuildInfo buildInfo, String fileName)
            throws TargetSetupError {
        try {
            File apkFile =
                    getApkFile(
                            buildInfo,
                            mProfileOwnerApk,
                            mAltDirs,
                            mAltDirBehavior,
                            false /* use resource as fallback */,
                            null /* device signing key */);
            if (apkFile == null) {
                throw new TargetSetupError(
                        String.format("Test app %s was not found.", mProfileOwnerApk),
                        device.getDeviceDescriptor());
            }
            return apkFile;
        } catch (IOException e) {
            throw new TargetSetupError(
                    String.format(
                            "failed to resolve apk path for apk %s in build %s",
                            mProfileOwnerApk, buildInfo.toString()),
                    e,
                    device.getDeviceDescriptor());
        }
    }
}