/*
 * Copyright OpenSearch Contributors
 * SPDX-License-Identifier: Apache-2.0
 */


package org.opensearch.jdbc.protocol.http;

import org.opensearch.jdbc.protocol.exceptions.MalformedResponseException;
import org.opensearch.jdbc.protocol.exceptions.ResponseException;
import com.fasterxml.jackson.core.JsonProcessingException;
import org.apache.http.HttpResponse;
import org.apache.http.HttpStatus;
import org.apache.http.util.EntityUtils;

import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

public class JsonHttpResponseHandler {

    protected JsonHttpProtocol protocol;

    public static final Set<Integer> DEFAULT_ACCEPTABLE_HTTP_CODES =
            Collections.unmodifiableSet(new HashSet<>(Arrays.asList(HttpStatus.SC_OK)));

    public JsonHttpResponseHandler(JsonHttpProtocol protocol) {
        this.protocol = protocol;
    }

    public <R> R handleResponse(HttpResponse response, JsonContentParser<InputStream, R> contentParser) throws ResponseException {
        return handleResponse(response, null, true, contentParser);
    }

    public <R> R handleResponse(HttpResponse response, Set<Integer> acceptableHttpStatusCodes,
                                boolean expectResponseBody, JsonContentParser<InputStream, R> contentParser) throws ResponseException {

        try {
            checkResponseForErrors(response, acceptableHttpStatusCodes, expectResponseBody);

            try (InputStream responseContent = response.getEntity().getContent()) {
                return contentParser.apply(responseContent);
            }

        } catch (JsonProcessingException jpe) {
            throw new MalformedResponseException("Could not process server response", jpe);
        } catch (IOException ioe) {
            throw new ResponseException("Error reading server response", ioe);
        }
    }

    /**
     * Checks if an HttpResponse meets the requirements to be accepted.
     *
     * @param response HttpResponse to check
     *
     * @throws ResponseException if the HttpResponse fails any of the checks
     * @throws IOException if there is an I/O exception when reading the
     *         HttpResponse
     */
    protected void checkResponseForErrors(HttpResponse response, Set<Integer> acceptableHttpStatusCodes,
                                          boolean expectResponseBodyContent) throws ResponseException, IOException {
        int statusCode = response.getStatusLine().getStatusCode();

        if (!isHttpStatusCodeAcceptable(statusCode, acceptableHttpStatusCodes)) {

            String responseBody = extractResponseBody(response);

            HttpException httpException = new HttpException(statusCode,
                    response.getStatusLine().getReasonPhrase(), null, responseBody);

            if (statusCode == HttpStatus.SC_METHOD_NOT_ALLOWED) {
                throw new ResponseException("Could not submit SQL request to the target server. " +
                        "Make sure the SQL plugin is installed on the server and responding on the " +
                        "\"/" + protocol.getSqlContextPath() + "\" context path.", httpException);
            } else {
                throw httpException;
            }
        }

        if (expectResponseBodyContent &&
                (response.getEntity() == null || response.getEntity().getContent() == null)) {
            throw new MalformedResponseException("Empty response.");
        }
    }


    /**
     * Reads and returns the entire response body content present in an
     * HttpResponse as a String.
     * <p>
     * This is meant to be used only in cases where a request to the
     * server has failed. In such cases we check the response
     * body contains any text or json content which could potentially
     * be helpful for a user to understand why the request failed.
     * <p>
     * We expect such content to only contain error messages of a limited
     * size, so it's ok to read in the entire response body as a String.
     *
     * @param response HttpResponse to extract the response body from
     *
     * @return the response body as a String or null if no response body
     *         is present
     */
    protected String extractResponseBody(HttpResponse response) {
        // TODO - limit the amount read from response.getEntity() ?
        String responseBody = null;
        try {
            if (response.getEntity() != null && response.getEntity().getContent() != null &&
                    response.getEntity().getContentType() != null &&
                    response.getEntity().getContentType().getValue() != null &&
                    (response.getEntity().getContentType().getValue().contains("application/json") ||
                            response.getEntity().getContentType().getValue().contains("text/plain"))) {
                responseBody = EntityUtils.toString(response.getEntity());
            }

        } catch (IOException ioe) {
            // ignore
        }
        return responseBody;
    }

    /**
     * HTTP Status codes that indicate success for this response
     * handler
     *
     * @return Set of HTTP Status codes that indicate successful requests.
     */
    protected Set<Integer> getAcceptableHttpStatusCodes() {
        return DEFAULT_ACCEPTABLE_HTTP_CODES;
    }

    private boolean isHttpStatusCodeAcceptable(int statusCode, Set<Integer> acceptableHttpStatusCodes) {

        if (acceptableHttpStatusCodes == null)
            acceptableHttpStatusCodes = DEFAULT_ACCEPTABLE_HTTP_CODES;

        return acceptableHttpStatusCodes.contains(statusCode);
    }

    @FunctionalInterface
    public interface JsonContentParser<T, R> {
        R apply(T t) throws IOException;
    }
}