/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.ambari.server.orm.dao;

import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import jakarta.persistence.EntityManager;
import jakarta.persistence.TypedQuery;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.actionmanager.HostRoleStatus;
import org.apache.ambari.server.cleanup.TimeBasedCleanupPolicy;
import org.apache.ambari.server.orm.RequiresSession;
import org.apache.ambari.server.orm.entities.ExecutionCommandEntity;
import org.apache.ambari.server.orm.entities.HostRoleCommandEntity;
import org.apache.ambari.server.orm.entities.RequestEntity;
import org.apache.ambari.server.orm.entities.RequestOperationLevelEntity;
import org.apache.ambari.server.orm.entities.RequestResourceFilterEntity;
import org.apache.ambari.server.orm.entities.RoleSuccessCriteriaEntity;
import org.apache.ambari.server.orm.entities.StageEntity;
import org.apache.ambari.server.orm.entities.TopologyHostRequestEntity;
import org.apache.ambari.server.orm.entities.TopologyHostTaskEntity;
import org.apache.ambari.server.orm.entities.TopologyLogicalTaskEntity;
import org.apache.ambari.server.orm.helpers.SQLConstants;
import org.apache.ambari.server.orm.helpers.SQLOperations;
import org.apache.ambari.server.state.Clusters;
import org.apache.commons.collections4.CollectionUtils;
import org.eclipse.persistence.config.HintValues;
import org.eclipse.persistence.config.QueryHints;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.google.inject.Inject;
import com.google.inject.Provider;
import com.google.inject.Singleton;
import com.google.inject.persist.Transactional;

@Singleton
public class RequestDAO implements Cleanable {

  private static final Logger LOG = LoggerFactory.getLogger(RequestDAO.class);

  /**
   * SQL template to retrieve all request IDs, sorted by the ID.
   */
  private final static String REQUEST_IDS_SORTED_SQL = "SELECT request.requestId FROM RequestEntity request ORDER BY request.requestId {0}";

  /**
   * Requests by cluster.  Cannot be a NamedQuery due to the ORDER BY clause.
   */
  private final static String REQUESTS_WITH_CLUSTER_SQL =
      "SELECT request.requestId FROM RequestEntity request WHERE request.clusterId = %s ORDER BY request.requestId %s";
  /**
   * Requests by cluster.  Cannot be a NamedQuery due to the ORDER BY clause.
   */
  private final static String REQUESTS_WITH_NO_CLUSTER_SQL =
      "SELECT request.requestId FROM RequestEntity request WHERE request.clusterId = -1 OR request.clusterId IS NULL ORDER BY request.requestId %s";

  @Inject
  Provider<EntityManager> entityManagerProvider;

  @Inject
  DaoUtils daoUtils;

  @Inject
  private Provider<Clusters> m_clusters;

  @Inject
  private HostRoleCommandDAO hostRoleCommandDAO;

  @Inject
  private TopologyLogicalTaskDAO topologyLogicalTaskDAO;

  @Inject
  private TopologyHostTaskDAO topologyHostTaskDAO;

  @Inject
  private TopologyLogicalRequestDAO topologyLogicalRequestDAO;

  @Inject
  private TopologyRequestDAO topologyRequestDAO;

  @RequiresSession
  public RequestEntity findByPK(Long requestId) {
    return entityManagerProvider.get().find(RequestEntity.class, requestId);
  }

  @RequiresSession
  public List<RequestEntity> findByPks(Collection<Long> requestIds) {
    return findByPks(requestIds, false);
  }

  /**
   * Given a collection of request ids, load the corresponding entities
   * @param requestIds  the collection of request ids
   * @param refreshHint {@code true} to hint JPA that the list should be refreshed
   * @return the list entities. An empty list if the requestIds are not provided
   */
  @RequiresSession
  public List<RequestEntity> findByPks(Collection<Long> requestIds, boolean refreshHint) {
    TypedQuery<RequestEntity> query = entityManagerProvider.get().createQuery("SELECT request FROM RequestEntity request " +
        "WHERE request.requestId IN ?1", RequestEntity.class);

    // !!! https://bugs.eclipse.org/bugs/show_bug.cgi?id=398067
    // ensure that an associated entity with a JOIN is not stale
    if (refreshHint) {
      query.setHint(QueryHints.REFRESH, HintValues.TRUE);
    }

    List<RequestEntity> result = new ArrayList<>();
    SQLOperations.batch(requestIds, SQLConstants.IN_ARGUMENT_MAX_SIZE, (chunk, currentBatch, totalBatches, totalSize) -> {
      result.addAll(daoUtils.selectList(query, chunk));
      return 0;
    });

    return Lists.newArrayList(result);
  }

  @RequiresSession
  public List<RequestEntity> findAll() {
    return daoUtils.selectAll(entityManagerProvider.get(), RequestEntity.class);
  }

  @RequiresSession
  public List<Long> findAllRequestIds(int limit, boolean ascending) {
    String sort = "ASC";
    if (!ascending) {
      sort = "DESC";
    }

    String sql = MessageFormat.format(REQUEST_IDS_SORTED_SQL, sort);
    TypedQuery<Long> query = entityManagerProvider.get().createQuery(sql,
        Long.class);

    query.setMaxResults(limit);

    return daoUtils.selectList(query);
  }

  @RequiresSession
  public List<RequestResourceFilterEntity> findAllResourceFilters() {
    return daoUtils.selectAll(entityManagerProvider.get(), RequestResourceFilterEntity.class);
  }

  @RequiresSession
  public boolean isAllTasksCompleted(long requestId) {
    TypedQuery<Long> query = entityManagerProvider.get().createQuery(
        "SELECT task.taskId FROM HostRoleCommandEntity task WHERE task.requestId = ?1 AND " +
          "task.stageId=(select max(stage.stageId) FROM StageEntity stage WHERE stage.requestId=?1) " +
          "AND task.status NOT IN ?2",
        Long.class
    );
    query.setMaxResults(1); //we don't need all
    return daoUtils.selectList(query, requestId, HostRoleStatus.getCompletedStates()).isEmpty();
  }

  @RequiresSession
  public Long getLastStageId(long requestId) {
    TypedQuery<Long> query = entityManagerProvider.get().createQuery("SELECT max(stage.stageId) " +
      "FROM StageEntity stage WHERE stage.requestId=?1", Long.class);
    return daoUtils.selectSingle(query, requestId);
  }

  @Transactional
  public RequestEntity updateStatus(long requestId, HostRoleStatus status, HostRoleStatus displayStatus) {
    RequestEntity requestEntity = findByPK(requestId);
    requestEntity.setStatus(status);
    requestEntity.setDisplayStatus(displayStatus);
    return merge(requestEntity);
  }

  @Transactional
  public void create(RequestEntity requestEntity) {
    entityManagerProvider.get().persist(requestEntity);
  }

  @Transactional
  public RequestEntity merge(RequestEntity requestEntity) {
    return entityManagerProvider.get().merge(requestEntity);
  }

  @Transactional
  public void remove(RequestEntity requestEntity) {
    entityManagerProvider.get().remove(merge(requestEntity));
  }

  @Transactional
  public void removeByPK(Long requestId) {
    remove(findByPK(requestId));
  }

  /**
   * Retrieves from the database for a cluster, or specifically for non-cluster requests.
   * This method should be considered temporary until Request/Stage/Task cleanup is achieved.
   *
   * @param limit the max number to return
   * @param sortAscending {@code true} to sort by requestId ascending, {@code false} for descending
   * @param clusterId   the cluster to find, or {@code null} to search for requests without cluster
   */
  @RequiresSession
  public List<Long> findAllRequestIds(int limit, boolean sortAscending, Long clusterId) {

    final String sql;

    if (null == clusterId) {
      sql = String.format(REQUESTS_WITH_NO_CLUSTER_SQL, sortAscending ? "ASC" : "DESC");
    } else {
      sql = String.format(REQUESTS_WITH_CLUSTER_SQL, clusterId, sortAscending ? "ASC" : "DESC");
    }

    TypedQuery<Long> query = entityManagerProvider.get().createQuery(sql,
        Long.class);

    query.setMaxResults(limit);

    return daoUtils.selectList(query);
  }

  public static final class StageEntityPK {
    private Long requestId;
    private Long stageId;

    public StageEntityPK(Long requestId, Long stageId) {
      this.requestId = requestId;
      this.stageId = stageId;
    }

    public Long getStageId() {
      return stageId;
    }

    public void setStageId(Long stageId) {
      this.stageId = stageId;
    }

    public Long getRequestId() {
      return requestId;
    }

    public void setRequestId(Long requestId) {
      this.requestId = requestId;
    }
  }

  /**
   * Search for all request ids in Upgrade table
   * @return the list of request ids
   */
  private Set<Long> findAllRequestIdsFromUpgrade() {
    EntityManager entityManager = entityManagerProvider.get();
    TypedQuery<Long> upgradeQuery =
            entityManager.createNamedQuery("UpgradeEntity.findAllRequestIds", Long.class);

    return Sets.newHashSet(daoUtils.selectList(upgradeQuery));
  }

  /**
   * Search for all request and stage ids in Request and Stage tables
   * @return the list of request/stage ids
   */
  public List<StageEntityPK> findRequestAndStageIdsInClusterBeforeDate(Long clusterId, long  beforeDateMillis) {
    EntityManager entityManager = entityManagerProvider.get();
    TypedQuery<StageEntityPK> requestQuery =
            entityManager.createNamedQuery("RequestEntity.findRequestStageIdsInClusterBeforeDate", StageEntityPK.class);

    requestQuery.setParameter("clusterId", clusterId);
    requestQuery.setParameter("beforeDate", beforeDateMillis);

    return daoUtils.selectList(requestQuery);
  }

  /**
   * In this method we are removing entities using passed ids,
   * To prevent issues we are using batch request to remove limited
   * count of entities.
   * @param ids              list of ids that we are using to remove rows from table
   * @param paramName        name of parameter that we are using in sql query (taskIds, stageIds)
   * @param entityName       name of entity which we will remove
   * @param beforeDateMillis timestamp which was set by user (remove all entities that were created before),
   *                         we are using it only for logging
   * @param entityQuery      name of NamedQuery which we will use to remove needed entities
   * @param type             type of entity class which we will use for casting query result
   * @return                 rows count that were removed
   */
  @Transactional
  protected <T> int cleanTableByIds(Set<Long> ids, String paramName, String entityName, Long beforeDateMillis,
                                  String entityQuery, Class<T> type) {
    LOG.info(String.format("Purging %s entity records before date %s", entityName, new Date(beforeDateMillis)));

    if (CollectionUtils.isEmpty(ids)) {
      return 0;
    }

    final EntityManager entityManager = entityManagerProvider.get();
    final TypedQuery<T> query = entityManager.createNamedQuery(entityQuery, type);

    // Batch delete
    return SQLOperations.batch(ids, SQLConstants.IN_ARGUMENT_MAX_SIZE, (chunk, currentBatch, totalBatches, totalSize) -> {
      LOG.info(String.format("Purging %s entity, batch %s/%s.", entityName, currentBatch, totalBatches));
      query.setParameter(paramName, chunk);
      return query.executeUpdate();
    });
  }

  /**
   * In this method we are removing entities using passed few ids,
   * To prevent issues we are using batch request to remove limited
   * count of entities.
   * @param ids              list of ids pairs that we are using to remove rows from table
   * @param paramNames       list of two names of parameters that we are using in sql query (taskIds, stageIds)
   * @param entityName       name of entity which we will remove
   * @param beforeDateMillis timestamp which was set by user (remove all entities that were created before),
   *                         we are using it only for logging
   * @param entityQuery      name of NamedQuery which we will use to remove needed entities
   * @param type             type of entity class which we will use for casting query result
   * @return                 rows count that were removed
   */
  @Transactional
  protected <T> int cleanTableByStageEntityPK(List<StageEntityPK> ids, LinkedList<String> paramNames, String entityName, Long beforeDateMillis,
                                  String entityQuery, Class<T> type) {
    LOG.info(String.format("Purging %s entity records before date %s", entityName, new Date(beforeDateMillis)));

    if (CollectionUtils.isEmpty(ids)) {
      return 0;
    }

    final EntityManager entityManager = entityManagerProvider.get();
    final TypedQuery<T> query = entityManager.createNamedQuery(entityQuery, type);

    // Batch delete
    return SQLOperations.batch(ids, SQLConstants.IN_ARGUMENT_MAX_SIZE, (chunk, currentBatch, totalBatches, totalSize) -> {
      int affectedRows = 0;

      for (StageEntityPK requestIds : chunk) {
        query.setParameter(paramNames.get(0), requestIds.getStageId());
        query.setParameter(paramNames.get(1), requestIds.getRequestId());
        affectedRows += query.executeUpdate();
      }
      return affectedRows;
    });
  }

  @Transactional
  @Override
  public long cleanup(TimeBasedCleanupPolicy policy) {
    long affectedRows = 0;
    try {
      Long clusterId = m_clusters.get().getCluster(policy.getClusterName()).getClusterId();
      // find request and stage ids that were created before date populated by user.
      List<StageEntityPK> requestStageIds = findRequestAndStageIdsInClusterBeforeDate(clusterId, policy.getToDateInMillis());

      // find request ids from Upgrade table and exclude these ids from
      // request ids set that we already have. We don't want to make any changes for upgrade
      Set<Long> requestIdsFromUpgrade = findAllRequestIdsFromUpgrade();
      Iterator<StageEntityPK> requestStageIdsIterator =  requestStageIds.iterator();
      Set<Long> requestIds = new HashSet<>();
      while (requestStageIdsIterator.hasNext()) {
        StageEntityPK nextRequestStageIds = requestStageIdsIterator.next();
        if (requestIdsFromUpgrade.contains(nextRequestStageIds.getRequestId())) {
          requestStageIdsIterator.remove();
          continue;
        }
        requestIds.add(nextRequestStageIds.getRequestId());
      }

      // find task ids using request stage ids
      Set<Long> taskIds = hostRoleCommandDAO.findTaskIdsByRequestStageIds(requestStageIds);
      LinkedList<String> params = new LinkedList<>();
      params.add("stageId");
      params.add("requestId");

      // find host task ids, to find related host requests and also to remove needed host tasks
      Set<Long> hostTaskIds = topologyLogicalTaskDAO.findHostTaskIdsByPhysicalTaskIds(taskIds);

      // find host request ids by host task ids to remove later needed host requests
      Set<Long> hostRequestIds = topologyHostTaskDAO.findHostRequestIdsByHostTaskIds(hostTaskIds);
      Set<Long> topologyRequestIds = topologyLogicalRequestDAO.findRequestIdsByIds(hostRequestIds);

      //removing all entities one by one according to their relations using stage, task and request ids
      affectedRows += cleanTableByIds(taskIds, "taskIds", "ExecutionCommand", policy.getToDateInMillis(),
        "ExecutionCommandEntity.removeByTaskIds", ExecutionCommandEntity.class);
      affectedRows += cleanTableByIds(taskIds, "taskIds", "TopologyLogicalTask", policy.getToDateInMillis(),
        "TopologyLogicalTaskEntity.removeByPhysicalTaskIds", TopologyLogicalTaskEntity.class);
      affectedRows += cleanTableByIds(hostTaskIds, "hostTaskIds", "TopologyHostTask", policy.getToDateInMillis(),
        "TopologyHostTaskEntity.removeByTaskIds", TopologyHostTaskEntity.class);
      affectedRows += cleanTableByIds(hostRequestIds, "hostRequestIds", "TopologyHostRequest", policy.getToDateInMillis(),
        "TopologyHostRequestEntity.removeByIds", TopologyHostRequestEntity.class);
      for (Long topologyRequestId : topologyRequestIds) {
        topologyRequestDAO.removeByPK(topologyRequestId);
      }
      affectedRows += cleanTableByIds(taskIds, "taskIds", "HostRoleCommand", policy.getToDateInMillis(),
        "HostRoleCommandEntity.removeByTaskIds", HostRoleCommandEntity.class);
      affectedRows += cleanTableByStageEntityPK(requestStageIds, params, "RoleSuccessCriteria", policy.getToDateInMillis(),
        "RoleSuccessCriteriaEntity.removeByRequestStageIds", RoleSuccessCriteriaEntity.class);
      affectedRows += cleanTableByStageEntityPK(requestStageIds, params, "Stage", policy.getToDateInMillis(),
        "StageEntity.removeByRequestStageIds", StageEntity.class);
      affectedRows += cleanTableByIds(requestIds, "requestIds", "RequestResourceFilter", policy.getToDateInMillis(),
        "RequestResourceFilterEntity.removeByRequestIds", RequestResourceFilterEntity.class);
      affectedRows += cleanTableByIds(requestIds, "requestIds", "RequestOperationLevel", policy.getToDateInMillis(),
        "RequestOperationLevelEntity.removeByRequestIds", RequestOperationLevelEntity.class);
      affectedRows += cleanTableByIds(requestIds, "requestIds", "Request", policy.getToDateInMillis(),
        "RequestEntity.removeByRequestIds", RequestEntity.class);

    } catch (AmbariException e) {
      LOG.error("Error while looking up cluster with name: {}", policy.getClusterName(), e);
      throw new IllegalStateException(e);
    }

    return affectedRows;
  }
}
