/**
 * Copyright 2016 Google Inc. All Rights Reserved.
 *
 * <p>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
 *
 * <p>http://www.apache.org/licenses/LICENSE-2.0
 *
 * <p>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.vts.util;


import com.android.vts.entity.ApiCoverageEntity;
import com.android.vts.entity.BranchEntity;
import com.android.vts.entity.BuildTargetEntity;
import com.android.vts.entity.CodeCoverageEntity;
import com.android.vts.entity.CoverageEntity;
import com.android.vts.entity.DeviceInfoEntity;
import com.android.vts.entity.ProfilingPointRunEntity;
import com.android.vts.entity.TestCaseRunEntity;
import com.android.vts.entity.TestEntity;
import com.android.vts.entity.TestPlanEntity;
import com.android.vts.entity.TestPlanRunEntity;
import com.android.vts.entity.TestRunEntity;
import com.android.vts.entity.TestRunEntity.TestRunType;
import com.android.vts.job.VtsAlertJobServlet;
import com.android.vts.job.VtsCoverageAlertJobServlet;
import com.android.vts.job.VtsProfilingStatsJobServlet;
import com.android.vts.proto.VtsReportMessage.AndroidDeviceInfoMessage;
import com.android.vts.proto.VtsReportMessage.ApiCoverageReportMessage;
import com.android.vts.proto.VtsReportMessage.CoverageReportMessage;
import com.android.vts.proto.VtsReportMessage.HalInterfaceMessage;
import com.android.vts.proto.VtsReportMessage.LogMessage;
import com.android.vts.proto.VtsReportMessage.ProfilingReportMessage;
import com.android.vts.proto.VtsReportMessage.TestCaseReportMessage;
import com.android.vts.proto.VtsReportMessage.TestCaseResult;
import com.android.vts.proto.VtsReportMessage.TestPlanReportMessage;
import com.android.vts.proto.VtsReportMessage.TestReportMessage;
import com.android.vts.proto.VtsReportMessage.UrlResourceMessage;
import com.google.appengine.api.datastore.DatastoreFailureException;
import com.google.appengine.api.datastore.DatastoreService;
import com.google.appengine.api.datastore.DatastoreServiceFactory;
import com.google.appengine.api.datastore.DatastoreTimeoutException;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityNotFoundException;
import com.google.appengine.api.datastore.FetchOptions;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.datastore.KeyFactory;
import com.google.appengine.api.datastore.Query;
import com.google.appengine.api.datastore.Query.Filter;
import com.google.appengine.api.datastore.Query.FilterOperator;
import com.google.appengine.api.datastore.Query.FilterPredicate;
import com.google.appengine.api.datastore.Transaction;
import com.google.appengine.api.datastore.TransactionOptions;
import com.google.common.collect.Lists;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.ConcurrentModificationException;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;

/**
 * DatastoreHelper, a helper class for interacting with Cloud Datastore.
 */
public class DatastoreHelper {

  /**
   * The default kind name for datastore
   */
  public static final String NULL_ENTITY_KIND = "nullEntity";

  public static final int MAX_WRITE_RETRIES = 5;
  /**
   * This variable is for maximum number of entities per transaction You can find the detail here
   * (https://cloud.google.com/datastore/docs/concepts/limits)
   */
  public static final int MAX_ENTITY_SIZE_PER_TRANSACTION = 300;

  protected static final Logger logger = Logger.getLogger(DatastoreHelper.class.getName());
  private static final DatastoreService datastore = DatastoreServiceFactory.getDatastoreService();

  /**
   * Get query fetch options for large batches of entities.
   *
   * @return FetchOptions with a large chunk and prefetch size.
   */
  public static FetchOptions getLargeBatchOptions() {
    return FetchOptions.Builder.withChunkSize(1000).prefetchSize(1000);
  }

  /**
   * Returns true if there are data points newer than lowerBound in the results table.
   *
   * @param parentKey The parent key to use in the query.
   * @param kind The query entity kind.
   * @param lowerBound The (exclusive) lower time bound, long, microseconds.
   * @return boolean True if there are newer data points.
   */
  public static boolean hasNewer(Key parentKey, String kind, Long lowerBound) {
    if (lowerBound == null || lowerBound <= 0) {
      return false;
    }
    DatastoreService datastore = DatastoreServiceFactory.getDatastoreService();
    Key startKey = KeyFactory.createKey(parentKey, kind, lowerBound);
    Filter startFilter =
            new FilterPredicate(
                    Entity.KEY_RESERVED_PROPERTY, FilterOperator.GREATER_THAN, startKey);
    Query q = new Query(kind).setAncestor(parentKey).setFilter(startFilter).setKeysOnly();
    return datastore.prepare(q).countEntities(FetchOptions.Builder.withLimit(1)) > 0;
  }

  /**
   * Returns true if there are data points older than upperBound in the table.
   *
   * @param parentKey The parent key to use in the query.
   * @param kind The query entity kind.
   * @param upperBound The (exclusive) upper time bound, long, microseconds.
   * @return boolean True if there are older data points.
   */
  public static boolean hasOlder(Key parentKey, String kind, Long upperBound) {
    if (upperBound == null || upperBound <= 0) {
      return false;
    }
    Key endKey = KeyFactory.createKey(parentKey, kind, upperBound);
    Filter endFilter =
            new FilterPredicate(Entity.KEY_RESERVED_PROPERTY, FilterOperator.LESS_THAN, endKey);
    Query q = new Query(kind).setAncestor(parentKey).setFilter(endFilter).setKeysOnly();
    return datastore.prepare(q).countEntities(FetchOptions.Builder.withLimit(1)) > 0;
  }

  /**
   * Get all of the devices branches.
   *
   * @return a list of all branches.
   */
  public static List<String> getAllBranches() {
    Query query = new Query(BranchEntity.KIND).setKeysOnly();
    List<String> branches = new ArrayList<>();
    for (Entity e : datastore.prepare(query).asIterable(getLargeBatchOptions())) {
      branches.add(e.getKey().getName());
    }
    return branches;
  }

  /**
   * Get all of the device build flavors.
   *
   * @return a list of all device build flavors.
   */
  public static List<String> getAllBuildFlavors() {
    Query query = new Query(BuildTargetEntity.KIND).setKeysOnly();
    List<String> devices = new ArrayList<>();
    for (Entity e : datastore.prepare(query).asIterable(getLargeBatchOptions())) {
      devices.add(e.getKey().getName());
    }
    return devices;
  }

  /**
   * Upload data from a test report message
   *
   * @param report The test report containing data to upload.
   */
  public static void insertTestReport(TestReportMessage report) {

    List<Entity> testEntityList = new ArrayList<>();
    List<Entity> branchEntityList = new ArrayList<>();
    List<Entity> buildTargetEntityList = new ArrayList<>();
    List<Entity> coverageEntityList = new ArrayList<>();
    List<Entity> profilingPointRunEntityList = new ArrayList<>();

    if (!report.hasStartTimestamp()
            || !report.hasEndTimestamp()
            || !report.hasTest()
            || !report.hasHostInfo()
            || !report.hasBuildInfo()) {
      // missing information
      return;
    }
    long startTimestamp = report.getStartTimestamp();
    long endTimestamp = report.getEndTimestamp();
    String testName = report.getTest().toStringUtf8();
    String testBuildId = report.getBuildInfo().getId().toStringUtf8();
    String hostName = report.getHostInfo().getHostname().toStringUtf8();

    TestEntity testEntity = new TestEntity(testName);

    Key testRunKey =
            KeyFactory.createKey(
                    testEntity.getOldKey(), TestRunEntity.KIND, report.getStartTimestamp());

    long passCount = 0;
    long failCount = 0;
    long coveredLineCount = 0;
    long totalLineCount = 0;

    Set<Key> buildTargetKeys = new HashSet<>();
    Set<Key> branchKeys = new HashSet<>();
    List<TestCaseRunEntity> testCases = new ArrayList<>();
    List<Key> profilingPointKeys = new ArrayList<>();
    List<String> links = new ArrayList<>();

    // Process test cases
    for (TestCaseReportMessage testCase : report.getTestCaseList()) {
      String testCaseName = testCase.getName().toStringUtf8();
      TestCaseResult result = testCase.getTestResult();
      // Track global pass/fail counts
      if (result == TestCaseResult.TEST_CASE_RESULT_PASS) {
        ++passCount;
      } else if (result != TestCaseResult.TEST_CASE_RESULT_SKIP) {
        ++failCount;
      }
      if (testCase.getSystraceCount() > 0
              && testCase.getSystraceList().get(0).getUrlCount() > 0) {
        String systraceLink = testCase.getSystraceList().get(0).getUrl(0).toStringUtf8();
        links.add(systraceLink);
      }

      // Process coverage data for test case
      for (CoverageReportMessage coverage : testCase.getCoverageList()) {
        CoverageEntity coverageEntity =
                CoverageEntity.fromCoverageReport(testRunKey, testCaseName, coverage);
        if (coverageEntity == null) {
          logger.log(Level.WARNING, "Invalid coverage report in test run " + testRunKey);
        } else {
          coveredLineCount += coverageEntity.getCoveredCount();
          totalLineCount += coverageEntity.getTotalCount();
          coverageEntityList.add(coverageEntity.toEntity());
        }
      }

      // Process profiling data for test case
      for (ProfilingReportMessage profiling : testCase.getProfilingList()) {
        ProfilingPointRunEntity profilingPointRunEntity =
                ProfilingPointRunEntity.fromProfilingReport(testRunKey, profiling);
        if (profilingPointRunEntity == null) {
          logger.log(Level.WARNING, "Invalid profiling report in test run " + testRunKey);
        } else {
          profilingPointRunEntityList.add(profilingPointRunEntity.toEntity());
          profilingPointKeys.add(profilingPointRunEntity.getKey());
          testEntity.setHasProfilingData(true);
        }
      }

      int lastIndex = testCases.size() - 1;
      if (lastIndex < 0 || testCases.get(lastIndex).isFull()) {
        testCases.add(new TestCaseRunEntity());
        ++lastIndex;
      }
      TestCaseRunEntity testCaseEntity = testCases.get(lastIndex);
      testCaseEntity.addTestCase(testCaseName, result.getNumber());
    }

    List<Entity> testCasePuts = new ArrayList<>();
    for (TestCaseRunEntity testCaseEntity : testCases) {
      testCasePuts.add(testCaseEntity.toEntity());
    }
    List<Key> testCaseKeys = datastore.put(testCasePuts);

    List<Long> testCaseIds = new ArrayList<>();
    for (Key key : testCaseKeys) {
      testCaseIds.add(key.getId());
    }

    // Process device information
    long testRunType = 0;
    for (AndroidDeviceInfoMessage device : report.getDeviceInfoList()) {
      DeviceInfoEntity deviceInfoEntity =
              DeviceInfoEntity.fromDeviceInfoMessage(testRunKey, device);
      if (deviceInfoEntity == null) {
        logger.log(Level.WARNING, "Invalid device info in test run " + testRunKey);
      } else {
        // Run type on devices must be the same, else set to OTHER
        TestRunType runType = TestRunType.fromBuildId(deviceInfoEntity.getBuildId());
        if (runType == null) {
          testRunType = TestRunType.OTHER.getNumber();
        } else {
          testRunType = runType.getNumber();
        }
        testEntityList.add(deviceInfoEntity.toEntity());
        BuildTargetEntity target = new BuildTargetEntity(deviceInfoEntity.getBuildFlavor());
        if (buildTargetKeys.add(target.key)) {
          buildTargetEntityList.add(target.toEntity());
        }
        BranchEntity branch = new BranchEntity(deviceInfoEntity.getBranch());
        if (branchKeys.add(branch.key)) {
          branchEntityList.add(branch.toEntity());
        }
      }
    }

    // Overall run type should be determined by the device builds unless test build is OTHER
    if (testRunType == TestRunType.OTHER.getNumber()) {
      testRunType = TestRunType.fromBuildId(testBuildId).getNumber();
    } else if (TestRunType.fromBuildId(testBuildId) == TestRunType.OTHER) {
      testRunType = TestRunType.OTHER.getNumber();
    }

    // Process global coverage data
    for (CoverageReportMessage coverage : report.getCoverageList()) {
      CoverageEntity coverageEntity =
              CoverageEntity.fromCoverageReport(testRunKey, new String(), coverage);
      if (coverageEntity == null) {
        logger.log(Level.WARNING, "Invalid coverage report in test run " + testRunKey);
      } else {
        coveredLineCount += coverageEntity.getCoveredCount();
        totalLineCount += coverageEntity.getTotalCount();
        coverageEntityList.add(coverageEntity.toEntity());
      }
    }

    // Process global API coverage data
    for (ApiCoverageReportMessage apiCoverage : report.getApiCoverageList()) {
      HalInterfaceMessage halInterfaceMessage = apiCoverage.getHalInterface();
      List<String> halApiList = apiCoverage.getHalApiList().stream().map(h -> h.toStringUtf8())
              .collect(
                      Collectors.toList());
      List<String> coveredHalApiList = apiCoverage.getCoveredHalApiList().stream()
              .map(h -> h.toStringUtf8()).collect(
                      Collectors.toList());
      ApiCoverageEntity apiCoverageEntity = new ApiCoverageEntity(
              testRunKey,
              halInterfaceMessage.getHalPackageName().toStringUtf8(),
              halInterfaceMessage.getHalVersionMajor(),
              halInterfaceMessage.getHalVersionMinor(),
              halInterfaceMessage.getHalInterfaceName().toStringUtf8(),
              halApiList,
              coveredHalApiList
      );
      com.googlecode.objectify.Key apiCoverageEntityKey = apiCoverageEntity.save();
      if (apiCoverageEntityKey == null) {
        logger.log(Level.WARNING, "Invalid API coverage report in test run " + testRunKey);
      }
    }

    // Process global profiling data
    for (ProfilingReportMessage profiling : report.getProfilingList()) {
      ProfilingPointRunEntity profilingPointRunEntity =
              ProfilingPointRunEntity.fromProfilingReport(testRunKey, profiling);
      if (profilingPointRunEntity == null) {
        logger.log(Level.WARNING, "Invalid profiling report in test run " + testRunKey);
      } else {
        profilingPointRunEntityList.add(profilingPointRunEntity.toEntity());
        profilingPointKeys.add(profilingPointRunEntity.getKey());
        testEntity.setHasProfilingData(true);
      }
    }

    // Process log data
    for (LogMessage log : report.getLogList()) {
      if (log.hasUrl()) {
        links.add(log.getUrl().toStringUtf8());
      }
    }
    // Process url resource
    for (UrlResourceMessage resource : report.getLinkResourceList()) {
      if (resource.hasUrl()) {
        links.add(resource.getUrl().toStringUtf8());
      }
    }

    boolean hasCodeCoverage = totalLineCount > 0 && coveredLineCount >= 0;
    TestRunEntity testRunEntity =
            new TestRunEntity(
                    testEntity.getOldKey(),
                    testRunType,
                    startTimestamp,
                    endTimestamp,
                    testBuildId,
                    hostName,
                    passCount,
                    failCount,
                    hasCodeCoverage,
                    testCaseIds,
                    links);
    testEntityList.add(testRunEntity.toEntity());

    CodeCoverageEntity codeCoverageEntity = new CodeCoverageEntity(
            testRunEntity.getKey(),
            coveredLineCount,
            totalLineCount);
    testEntityList.add(codeCoverageEntity.toEntity());

    Entity test = testEntity.toEntity();

    if (datastoreTransactionalRetry(test, testEntityList)) {
      List<List<Entity>> auxiliaryEntityList =
              Arrays.asList(
                      profilingPointRunEntityList,
                      coverageEntityList,
                      branchEntityList,
                      buildTargetEntityList);
      int indexCount = 0;
      for (List<Entity> entityList : auxiliaryEntityList) {
        switch (indexCount) {
          case 0:
          case 1:
            if (entityList.size() > MAX_ENTITY_SIZE_PER_TRANSACTION) {
              List<List<Entity>> partitionedList =
                      Lists.partition(entityList, MAX_ENTITY_SIZE_PER_TRANSACTION);
              partitionedList.forEach(
                      subEntityList -> {
                        datastoreTransactionalRetry(
                                new Entity(NULL_ENTITY_KIND), subEntityList);
                      });
            } else {
              datastoreTransactionalRetry(new Entity(NULL_ENTITY_KIND), entityList);
            }
            break;
          case 2:
          case 3:
            datastoreTransactionalRetryWithXG(
                    new Entity(NULL_ENTITY_KIND), entityList, true);
            break;
          default:
            break;
        }
        indexCount++;
      }

      if (testRunEntity.getType() == TestRunType.POSTSUBMIT.getNumber()) {
        VtsAlertJobServlet.addTask(testRunKey);
        if (testRunEntity.getHasCodeCoverage()) {
          VtsCoverageAlertJobServlet.addTask(testRunKey);
        }
        if (profilingPointKeys.size() > 0) {
          VtsProfilingStatsJobServlet.addTasks(profilingPointKeys);
        }
      } else {
        logger.log(
                Level.WARNING,
                "The alert email was not sent as testRunEntity type is not POSTSUBMIT!" +
                        " \n " + " testRunEntity type => " + testRunEntity.getType());
      }
    }
  }

  /**
   * Upload data from a test plan report message
   *
   * @param report The test plan report containing data to upload.
   */
  public static void insertTestPlanReport(TestPlanReportMessage report) {
    List<Entity> testEntityList = new ArrayList<>();

    List<String> testModules = report.getTestModuleNameList();
    List<Long> testTimes = report.getTestModuleStartTimestampList();
    if (testModules.size() != testTimes.size() || !report.hasTestPlanName()) {
      logger.log(Level.WARNING, "TestPlanReportMessage is missing information.");
      return;
    }

    String testPlanName = report.getTestPlanName();
    Entity testPlanEntity = new TestPlanEntity(testPlanName).toEntity();
    List<Key> testRunKeys = new ArrayList<>();
    for (int i = 0; i < testModules.size(); i++) {
      String test = testModules.get(i);
      long time = testTimes.get(i);
      Key parentKey = KeyFactory.createKey(TestEntity.KIND, test);
      Key testRunKey = KeyFactory.createKey(parentKey, TestRunEntity.KIND, time);
      testRunKeys.add(testRunKey);
    }
    Map<Key, Entity> testRuns = datastore.get(testRunKeys);
    long passCount = 0;
    long failCount = 0;
    long startTimestamp = -1;
    long endTimestamp = -1;
    String testBuildId = null;
    long type = 0;
    Set<DeviceInfoEntity> deviceInfoEntitySet = new HashSet<>();
    for (Key testRunKey : testRuns.keySet()) {
      TestRunEntity testRun = TestRunEntity.fromEntity(testRuns.get(testRunKey));
      if (testRun == null) {
        continue; // not a valid test run
      }
      passCount += testRun.getPassCount();
      failCount += testRun.getFailCount();
      if (startTimestamp < 0 || testRunKey.getId() < startTimestamp) {
        startTimestamp = testRunKey.getId();
      }
      if (endTimestamp < 0 || testRun.getEndTimestamp() > endTimestamp) {
        endTimestamp = testRun.getEndTimestamp();
      }
      type = testRun.getType();
      testBuildId = testRun.getTestBuildId();
      Query deviceInfoQuery = new Query(DeviceInfoEntity.KIND).setAncestor(testRunKey);
      for (Entity deviceInfoEntity : datastore.prepare(deviceInfoQuery).asIterable()) {
        DeviceInfoEntity device = DeviceInfoEntity.fromEntity(deviceInfoEntity);
        if (device == null) {
          continue; // invalid entity
        }
        deviceInfoEntitySet.add(device);
      }
    }
    if (startTimestamp < 0 || testBuildId == null || type == 0) {
      logger.log(Level.WARNING, "Couldn't infer test run information from runs.");
      return;
    }
    TestPlanRunEntity testPlanRun =
            new TestPlanRunEntity(
                    testPlanEntity.getKey(),
                    testPlanName,
                    type,
                    startTimestamp,
                    endTimestamp,
                    testBuildId,
                    passCount,
                    failCount,
                    0L,
                    0L,
                    testRunKeys);

    // Create the device infos.
    for (DeviceInfoEntity device : deviceInfoEntitySet) {
      testEntityList.add(device.copyWithParent(testPlanRun.key).toEntity());
    }
    testEntityList.add(testPlanRun.toEntity());

    // Add the task to calculate total number API list.
    testPlanRun.addCoverageApiTask();

    datastoreTransactionalRetry(testPlanEntity, testEntityList);
  }

  /**
   * Datastore Transactional process for data insertion with MAX_WRITE_RETRIES times and withXG of
   * false value
   *
   * @param entity The entity that you want to insert to datastore.
   * @param entityList The list of entity for using datastore put method.
   */
  private static boolean datastoreTransactionalRetry(Entity entity, List<Entity> entityList) {
    return datastoreTransactionalRetryWithXG(entity, entityList, false);
  }

  /**
   * Datastore Transactional process for data insertion with MAX_WRITE_RETRIES times
   *
   * @param entity The entity that you want to insert to datastore.
   * @param entityList The list of entity for using datastore put method.
   */
  private static boolean datastoreTransactionalRetryWithXG(
          Entity entity, List<Entity> entityList, boolean withXG) {
    int retries = 0;
    while (true) {
      Transaction txn;
      if (withXG) {
        TransactionOptions options = TransactionOptions.Builder.withXG(withXG);
        txn = datastore.beginTransaction(options);
      } else {
        txn = datastore.beginTransaction();
      }

      try {
        // Check if test already exists in the database
        if (!entity.getKind().equalsIgnoreCase(NULL_ENTITY_KIND)) {
          try {
            if (entity.getKind().equalsIgnoreCase("Test")) {
              Entity datastoreEntity = datastore.get(entity.getKey());
              TestEntity datastoreTestEntity = TestEntity.fromEntity(datastoreEntity);
              if (datastoreTestEntity == null
                      || !datastoreTestEntity.equals(entity)) {
                entityList.add(entity);
              }
            } else if (entity.getKind().equalsIgnoreCase("TestPlan")) {
              datastore.get(entity.getKey());
            } else {
              datastore.get(entity.getKey());
            }
          } catch (EntityNotFoundException e) {
            entityList.add(entity);
          }
        }
        datastore.put(txn, entityList);
        txn.commit();
        break;
      } catch (ConcurrentModificationException
              | DatastoreFailureException
              | DatastoreTimeoutException e) {
        entityList.remove(entity);
        logger.log(
                Level.WARNING,
                "Retrying insert kind: " + entity.getKind() + " key: " + entity.getKey());
        if (retries++ >= MAX_WRITE_RETRIES) {
          logger.log(
                  Level.SEVERE,
                  "Exceeded maximum retries kind: "
                          + entity.getKind()
                          + " key: "
                          + entity.getKey());
          return false;
        }
      } finally {
        if (txn.isActive()) {
          logger.log(
                  Level.WARNING, "Transaction rollback forced for : " + entity.getKind());
          txn.rollback();
        }
      }
    }
    return true;
  }
}