/*
 * Copyright (C) 2020 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 android.virt.test;

import static org.junit.Assert.*;

import com.android.tradefed.build.IBuildInfo;
import com.android.tradefed.testtype.DeviceTestCase;
import com.android.tradefed.testtype.IAbi;
import com.android.tradefed.testtype.IAbiReceiver;
import com.android.tradefed.testtype.IBuildReceiver;
import com.android.tradefed.log.LogUtil.CLog;

import org.apache.commons.compress.compressors.CompressorStreamFactory;

import org.junit.Before;
import org.junit.Test;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;

public abstract class VirtHostTestCase extends DeviceTestCase
                                       implements IBuildReceiver, IAbiReceiver {

    private static final String DEVICE_TMP_DIR = "/data/local/tmp/";
    private static final String KERNEL_DEVICE_PATH = DEVICE_TMP_DIR + "/virt-host-test-kernel";
    private static final String RAMDISK_DEVICE_PATH = DEVICE_TMP_DIR + "/virt-host-test-initramfs";
    private static final String HOST_BINARY_DEVICE_PATH = DEVICE_TMP_DIR + "/virt-host-test-host";

    private static final int EXIT_SUCCESS = 0;
    private static final int EXIT_FAILURE = 1;

    private static final String CMDLINE_TEST_NAME_PARAM = "virt_test_name";

    private static final int CID_RESERVED = 2;

    private IBuildInfo mBuildInfo = null;
    private IAbi mAbi = null;

    /**
     * Waits for device to be online, marks the most recent boottime of the device
     */
    @Before
    public void setUp() throws Exception {
        getDevice().waitForDeviceAvailable();
    }

    protected boolean enableRoot() throws Exception {
        return getDevice().enableAdbRoot();
    }

    protected boolean disableRoot() throws Exception {
        return getDevice().disableAdbRoot();
    }

    protected String runCommand(String cmd) throws Exception {
        return getDevice().executeShellCommand(cmd);
    }

    protected int runCommandGetExitCode(String cmd) throws Exception {
        String output = runCommand("(" + cmd + ") > /dev/null 2>&1; echo $?").trim();

        try {
            return Integer.parseInt(output);
        } catch (NumberFormatException e) {
            throw new IllegalArgumentException(String.format(
                    "Could not get the exit status (%s) for '%s'.", output, cmd));
        }
    }

    private void extractResource(String resFilePath, File file) throws Exception {
        try (InputStream in = this.getClass().getResourceAsStream(resFilePath);
            OutputStream out = new BufferedOutputStream(new FileOutputStream(file))) {
            if (in == null) {
                throw new IllegalArgumentException("Resource not found: " + resFilePath);
            }
            byte[] buf = new byte[65536];
            int chunkSize;
            while ((chunkSize = in.read(buf)) != -1) {
                out.write(buf, 0, chunkSize);
            }
        }
    }

    private void pushResource(String resFilePath, String deviceFilePath, boolean decompress)
            throws Exception {
        File resFile = File.createTempFile("VirtHostTestResource", "");
        try {
            extractResource(resFilePath, resFile);
            getDevice().pushFile(resFile, deviceFilePath);
        } finally {
            resFile.delete();
        }
    }

    private void pushResource(String resFilePath, String deviceFilePath) throws Exception {
        pushResource(resFilePath, deviceFilePath, /* decompress */ false);
    }

    private boolean deleteResource(String deviceFilePath) throws Exception {
        return runCommandGetExitCode(String.format("rm %s", deviceFilePath)) == EXIT_SUCCESS;
    }

    private boolean makeExecutable(String deviceFilePath) throws Exception {
        return runCommandGetExitCode(String.format("chmod +x %s", deviceFilePath)) == EXIT_SUCCESS;
    }

    protected boolean isVirtEnabled() throws Exception {
        return (runCommandGetExitCode("ls /dev/kvm") == EXIT_SUCCESS) &&
               (runCommandGetExitCode("which crosvm") == EXIT_SUCCESS);
    }

    private String getKernelResourcePath() {
        return String.format("/kernel_%s", mAbi.getName());
    }

    private String getRamdiskResourcePath() {
        return String.format("/initramfs_%s.cpio", mAbi.getName());
    }

    private String getHostBinaryResourcePath() {
        return String.format("/host_binary");
    }

    protected void prepareGuestVM() throws Exception {
        if (!isVirtEnabled()) {
            throw new IllegalStateException("Virtualization is not enabled on the target device");
        }

        if (!enableRoot()) {
            throw new IllegalStateException("Cannot enable root on the target device");
        }

        pushResource(getKernelResourcePath(), KERNEL_DEVICE_PATH);
        pushResource(getRamdiskResourcePath(), RAMDISK_DEVICE_PATH);
        pushResource(getHostBinaryResourcePath(), HOST_BINARY_DEVICE_PATH);

        if (!makeExecutable(HOST_BINARY_DEVICE_PATH)) {
            throw new IllegalStateException("Could not make host binary executable");
        }
    }

    protected void finishGuestVM() throws Exception {
        deleteResource(KERNEL_DEVICE_PATH);
        deleteResource(RAMDISK_DEVICE_PATH);
        deleteResource(HOST_BINARY_DEVICE_PATH);
        disableRoot();
    }

    protected void runGuestVM(String testName, Integer cid) throws Exception {
        ArrayList<String> cmd = new ArrayList<>();
        String params = String.format("%s=%s", CMDLINE_TEST_NAME_PARAM, testName);

        cmd.add("crosvm");
        cmd.add("run");

        cmd.add("--disable-sandbox");

        if (cid != null) {
            if (cid > CID_RESERVED) {
                cmd.add("--cid");
                cmd.add(cid.toString());
            } else {
                throw new IllegalArgumentException("Invalid CID");
            }
        }

        cmd.add("--params");
        cmd.add("\"" + params.replace("\"", "\\\"") + "\"");

        cmd.add("--initrd");
        cmd.add(RAMDISK_DEVICE_PATH);

        cmd.add(KERNEL_DEVICE_PATH);

        runCommand(String.join(" ", cmd));
    }

    protected boolean runHostBinary(String testName) throws Exception {
        ArrayList<String> cmd = new ArrayList<>();

        cmd.add(HOST_BINARY_DEVICE_PATH);
        cmd.add(testName);

        return runCommandGetExitCode(String.join(" ", cmd)) == EXIT_SUCCESS;
    }

    @Override
    public void setBuild(IBuildInfo buildInfo) {
        // Get the build, this is used to access the APK.
        mBuildInfo = buildInfo;
    }

    @Override
    public void setAbi(IAbi abi) {
        mAbi = abi;
    }
}