/*
 * SPDX-License-Identifier: Apache-2.0
 *
 * The OpenSearch Contributors require contributions made to
 * this file be licensed under the Apache-2.0 license or a
 * compatible open source license.
 *
 * Modifications Copyright OpenSearch Contributors. See
 * GitHub history for details.
 */

package org.opensearch.ad.util;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.action.DocWriteRequest;
import org.opensearch.action.bulk.BulkItemResponse;
import org.opensearch.action.bulk.BulkRequest;
import org.opensearch.action.bulk.BulkResponse;
import org.opensearch.action.index.IndexRequest;
import org.opensearch.timeseries.util.ExceptionUtil;

public class BulkUtil {
    private static final Logger logger = LogManager.getLogger(BulkUtil.class);

    public static List<IndexRequest> getFailedIndexRequest(BulkRequest bulkRequest, BulkResponse bulkResponse) {
        List<IndexRequest> res = new ArrayList<>();

        if (bulkResponse == null || bulkRequest == null) {
            return res;
        }

        Set<String> failedId = new HashSet<>();
        for (BulkItemResponse response : bulkResponse.getItems()) {
            if (response.isFailed() && ExceptionUtil.isRetryAble(response.getFailure().getStatus())) {
                failedId.add(response.getId());
            }
        }

        for (DocWriteRequest<?> request : bulkRequest.requests()) {
            try {
                if (failedId.contains(request.id())) {
                    res.add((IndexRequest) request);
                }
            } catch (ClassCastException e) {
                logger.error("We only support IndexRequest");
                throw e;
            }

        }
        return res;
    }

    /**
     * Copy original request's source without other information like autoGeneratedTimestamp.
     * otherwise, an exception will be thrown indicating autoGeneratedTimestamp should not be set
     * while request id is already set (id is set because we have already sent the request before).
     * @param indexRequest request to be cloned
     * @return cloned Request
     */
    public static IndexRequest cloneIndexRequest(IndexRequest indexRequest) {
        IndexRequest newRequest = new IndexRequest(indexRequest.index());
        newRequest.source(indexRequest.source(), indexRequest.getContentType());
        return newRequest;
    }
}