/*
 * Copyright (C) 2025 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.tradefed.invoker;

import com.android.tradefed.build.BuildInfoKey.BuildInfoFileKey;
import com.android.tradefed.cache.ExecutableAction;
import com.android.tradefed.cache.ExecutableActionResult;
import com.android.tradefed.cache.ICacheClient;
import com.android.tradefed.config.ConfigurationException;
import com.android.tradefed.config.IConfiguration;
import com.android.tradefed.config.proxy.TradefedDelegator;
import com.android.tradefed.invoker.logger.CurrentInvocation;
import com.android.tradefed.invoker.logger.InvocationMetricLogger;
import com.android.tradefed.invoker.logger.InvocationMetricLogger.InvocationMetricKey;
import com.android.tradefed.invoker.tracing.CloseableTraceScope;
import com.android.tradefed.log.LogUtil.CLog;
import com.android.tradefed.result.proto.InvocationProtoResultReporter;
import com.android.tradefed.result.proto.ModuleProtoResultReporter;
import com.android.tradefed.util.CacheClientFactory;
import com.android.tradefed.util.FileUtil;
import com.android.tradefed.util.QuotationAwareTokenizer;

import build.bazel.remote.execution.v2.Digest;

import com.google.common.base.Strings;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

/** Utility to handle uploading and looking up invocation cache results. */
public class InvocationCacheHelper {

    private static final String CACHE_ERROR = "%%%CACHE_ERROR%%%";
    private static File sCacheDirectory;

    /** Describes the cache results. */
    public static class CacheInvocationResultDescriptor {
        private final boolean cacheHit;
        private final String cacheExplanation;

        public CacheInvocationResultDescriptor(boolean cacheHit, String explanation) {
            this.cacheHit = cacheHit;
            this.cacheExplanation = explanation;
        }

        public boolean isCacheHit() {
            return cacheHit;
        }

        public String getDetails() {
            return cacheExplanation;
        }
    }

    /**
     * Upload invocation results
     *
     * @param mainConfig
     * @param protoResultsReporter
     * @param testInfo
     */
    public static void uploadInvocationResults(
            IConfiguration mainConfig,
            InvocationProtoResultReporter protoResultsReporter,
            TestInformation testInfo) {
        if (testInfo.getDevices().size() > 1) {
            return;
        }
        if (protoResultsReporter == null) {
            return;
        }
        if (checkIterations(testInfo)) {
            return;
        }
        File protoResults = protoResultsReporter.getOutputFile();
        boolean emptyTestsDir = false;
        File invocationTestsDir = null;
        if (testInfo.getBuildInfo().getFile(BuildInfoFileKey.ROOT_DIRECTORY) == null
                || testInfo.getBuildInfo().getFile(BuildInfoFileKey.TESTDIR_IMAGE) == null) {
            emptyTestsDir = true;
        }
        if (!emptyTestsDir
                && (mainConfig.getSkipManager().getTestArtifactsToDigest().isEmpty()
                        || mainConfig
                                .getSkipManager()
                                .getTestArtifactsToDigest()
                                .containsValue(null))) {
            CLog.d("Cannot handle testsdir.");
            return;
        }
        File errorInCache = null;
        try (CloseableTraceScope ignored = new CloseableTraceScope("upload_invocation_results")) {
            String cacheInstance = mainConfig.getCommandOptions().getRemoteCacheInstanceName();
            ICacheClient cacheClient =
                    CacheClientFactory.createCacheClient(getCacheDirectory(), cacheInstance);
            invocationTestsDir =
                    FileUtil.createNamedTempDir(CurrentInvocation.getWorkFolder(), "invoc-cache");
            ExecutableAction action =
                    ExecutableAction.create(
                            invocationTestsDir,
                            getCommonCommandLine(mainConfig.getCommandLine()),
                            computeEnvironment(mainConfig),
                            60000L);
            if (protoResultsReporter.stopCaching()) {
                errorInCache = FileUtil.createTempFile("cache-error", ".txt");
                FileUtil.writeToFile(CACHE_ERROR, errorInCache);
            }
            ExecutableActionResult result =
                    ExecutableActionResult.create(0, protoResults, errorInCache);
            CLog.d("Uploading cache for %s and %s", action, protoResults);
            cacheClient.uploadCache(action, result);
        } catch (IOException | RuntimeException | InterruptedException e) {
            CLog.e(e);
        } finally {
            FileUtil.recursiveDelete(invocationTestsDir);
            FileUtil.deleteFile(errorInCache);
        }
    }

    public static CacheInvocationResultDescriptor lookupInvocationResults(
            IConfiguration mainConfig, TestInformation testInfo) {
        if (testInfo.getDevices().size() > 1) {
            return null;
        }
        if (InvocationContext.isOnDemand(testInfo.getContext())
                || InvocationContext.isApiTrigger(testInfo.getContext())) {
            return null;
        }
        if (mainConfig.getSkipManager().getImageToDigest().containsValue(null)) {
            CLog.d("No digest for device.");
            return new CacheInvocationResultDescriptor(false, null);
        }
        if (checkIterations(testInfo)) {
            CLog.d("No caching when using iterations.");
            return new CacheInvocationResultDescriptor(false, null);
        }
        boolean emptyTestsDir = false;
        File invocationTestsDir = null;
        if (testInfo.getBuildInfo().getFile(BuildInfoFileKey.ROOT_DIRECTORY) == null
                && testInfo.getBuildInfo().getFile(BuildInfoFileKey.TESTDIR_IMAGE) == null) {
            emptyTestsDir = true;
        }
        if (!emptyTestsDir
                && (mainConfig.getSkipManager().getTestArtifactsToDigest().isEmpty()
                        || mainConfig
                                .getSkipManager()
                                .getTestArtifactsToDigest()
                                .containsValue(null))) {
            CLog.d("Cannot handle testsdir.");
            return new CacheInvocationResultDescriptor(false, null);
        }
        ExecutableActionResult cachedResults = null;
        try (CloseableTraceScope ignored = new CloseableTraceScope("lookup_invocation_results")) {
            String cacheInstance = mainConfig.getCommandOptions().getRemoteCacheInstanceName();
            ICacheClient cacheClient =
                    CacheClientFactory.createCacheClient(getCacheDirectory(), cacheInstance);
            invocationTestsDir =
                    FileUtil.createNamedTempDir(CurrentInvocation.getWorkFolder(), "invoc-cache");
            ExecutableAction action =
                    ExecutableAction.create(
                            invocationTestsDir,
                            getCommonCommandLine(mainConfig.getCommandLine()),
                            computeEnvironment(mainConfig),
                            60000L);
            CLog.d("Looking up cache for %s", action);
            InvocationMetricLogger.addInvocationMetrics(
                    InvocationMetricKey.INVOCATION_RESULTS_CHECKING_CACHE, 1);
            cachedResults = cacheClient.lookupCache(action);
            if (cachedResults == null) {
                CLog.d("No cached results for the invocation.");
                return null;
            } else if (cachedResults.stdErr() != null &&
                           FileUtil.readStringFromFile(cachedResults.stdErr()).contains(CACHE_ERROR)) {
                InvocationMetricLogger.addInvocationMetrics(
                        InvocationMetricKey.INVOCATION_CACHE_HIT_WITH_FAILURE, 1);
                CLog.d("Cached invocation but with failures.");
                return null;
            } else {
                InvocationMetricLogger.addInvocationMetrics(
                        InvocationMetricKey.INVOCATION_CACHE_HIT, 1);
                String details = "Cached results.";
                Map<String, String> metadata =
                        ModuleProtoResultReporter.parseResultsMetadata(cachedResults.stdOut());
                if (metadata.containsKey(ModuleProtoResultReporter.INVOCATION_ID_KEY)) {
                    details +=
                            String.format(
                                    " origin of results: http://ab/%s",
                                    metadata.get(ModuleProtoResultReporter.INVOCATION_ID_KEY));
                    CLog.d(details);
                }
                return new CacheInvocationResultDescriptor(true, details);
            }
        } catch (IOException | RuntimeException | InterruptedException e) {
            CLog.e(e);
        } finally {
            FileUtil.recursiveDelete(invocationTestsDir);
            if (cachedResults != null) {
                FileUtil.deleteFile(cachedResults.stdOut());
                FileUtil.deleteFile(cachedResults.stdErr());
            }
        }
        return null;
    }

    private static Map<String, String> computeEnvironment(IConfiguration mainConfig) {
        Map<String, String> environment = new HashMap<>();
        for (Entry<String, Digest> entry :
                mainConfig.getSkipManager().getImageToDigest().entrySet()) {
            environment.put(entry.getKey(), entry.getValue().getHash());
        }
        for (Entry<String, Digest> entry :
                mainConfig.getSkipManager().getTestArtifactsToDigest().entrySet()) {
            environment.put(entry.getKey(), entry.getValue().getHash());
        }
        String atpTestName =
                mainConfig
                        .getCommandOptions()
                        .getInvocationData()
                        .getUniqueMap()
                        .get("atp_test_name");
        if (atpTestName != null) {
            environment.put("atp_test_name", atpTestName);
        }
        // add tradefed.jar version
        return environment;
    }

    private static List<String> getCommonCommandLine(String commandLine) {
        String[] commandArray = QuotationAwareTokenizer.tokenizeLine(commandLine, false);
        try {
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(commandArray, "invocation-data");
            commandArray = TradefedDelegator.clearCommandlineFromOneArg(commandArray, "build-id");
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(commandArray, "cts-build-id");
            commandArray = TradefedDelegator.clearCommandlineFromOneArg(commandArray, "serial");
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(
                            commandArray, "extra-target-downloads");
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(commandArray, "gce-driver-param");
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(
                            commandArray, "max-testcase-run-count");
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(
                            commandArray, "max-testrun-run-count");
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(
                            commandArray, "retry-strategy");
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(
                            commandArray, "reboot-at-last-retry", true);
            commandArray =
                    TradefedDelegator.clearCommandlineFromOneArg(commandArray, "cas-manifest-file");
        } catch (ConfigurationException e) {
            throw new RuntimeException(e);
        }
        return Arrays.asList(commandArray);
    }

    private static File getCacheDirectory() throws IOException {
        if (sCacheDirectory == null) {
            synchronized (InvocationCacheHelper.class) {
                if (sCacheDirectory == null) {
                    sCacheDirectory = FileUtil.createTempDir("invocation-cache-dir");
                }
            }
        }
        return sCacheDirectory;
    }

    private static boolean checkIterations(TestInformation testInfo) {
        IInvocationContext context = testInfo.getContext();
        String runCount = context.getAttribute("run_count");
        if (Strings.isNullOrEmpty(runCount)) {
            return false;
        }
        CLog.d("Evaluating run_count: %s", runCount);
        try {
            int runCountInt = Integer.parseInt(runCount);
            if (runCountInt > 1) {
                return true;
            }
        } catch (NumberFormatException e) {
            // Not a number
        }
        return false;
    }
}