/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License;
 * you may not use this file except in compliance with the Elastic License.
 */

package org.elasticsearch.xpack.sql.qa.rest;

import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.common.xcontent.XContentHelper;
import org.elasticsearch.common.xcontent.json.JsonXContent;
import org.elasticsearch.test.rest.ESRestTestCase;
import org.elasticsearch.xpack.sql.proto.Mode;
import org.elasticsearch.xpack.sql.proto.StringUtils;
import org.elasticsearch.xpack.sql.qa.FeatureMetric;
import org.junit.Before;

import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;

import static org.elasticsearch.xpack.sql.qa.rest.RestSqlTestCase.mode;

public abstract class RestSqlUsageTestCase extends ESRestTestCase {
    private List<IndexDocument> testData = Arrays.asList(
            new IndexDocument("used", "Don Quixote",     1072),
            new IndexDocument("used", "Vacuum Diagrams", 335),
            new IndexDocument("new",  "Darwin's Radio",  270),
            new IndexDocument("used", "The Martian",     387),
            new IndexDocument("new",  "Moving Mars",     495)
    );
    
    private enum ClientType {
        CANVAS, CLI, JDBC, ODBC, ODBC32, ODBC64, REST;

        @Override
        public String toString() {
            return this.name().toLowerCase(Locale.ROOT);
        }
        
    }

    private Map<String,Integer> baseMetrics = new HashMap<String,Integer>();
    private Integer baseClientTypeTotalQueries = 0;
    private Integer baseClientTypeFailedQueries = 0;
    private Integer baseAllTotalQueries = 0;
    private Integer baseAllFailedQueries = 0;
    private Integer baseTranslateRequests = 0;
    private String clientType;
    private boolean ignoreClientType;

    /**
     * This method gets the metrics' values before the test runs, in case these values
     * were changed by other tests running in the same REST test cluster. The test itself
     * will count the new metrics' values starting from the base values initialized here. 
     */
    @SuppressWarnings({ "unchecked", "rawtypes" })
    @Before
    private void getBaseMetrics() throws UnsupportedOperationException, IOException {
        Map<String, Object> baseStats = getStats();
        List<Map<String, Map<String, Map>>> nodesListStats = (List) baseStats.get("stats");
        
        // used for "client_id" request parameter value, but also for getting the stats from ES
        clientType = randomFrom(ClientType.values()).toString();
        ignoreClientType = randomBoolean();
        
        // "client_id" parameter will not be sent in the requests
        // and "clientType" will only be used for getting the stats back from ES
        if (ignoreClientType) {
            clientType = ClientType.REST.toString();
        }
        
        for (Map perNodeStats : nodesListStats) {
            Map featuresMetrics = (Map) ((Map) perNodeStats.get("stats")).get("features");
            Map queriesMetrics = (Map) ((Map) perNodeStats.get("stats")).get("queries");
            for (FeatureMetric metric : FeatureMetric.values()) {
                baseMetrics.put(metric.toString(), (Integer) featuresMetrics.get(metric.toString()));
            }
            
            // initialize the "base" metric values with whatever values are already recorded on ES
            baseClientTypeTotalQueries = ((Map<String,Integer>) queriesMetrics.get(clientType)).get("total");
            baseClientTypeFailedQueries = ((Map<String,Integer>) queriesMetrics.get(clientType)).get("failed");
            baseAllTotalQueries = ((Map<String,Integer>) queriesMetrics.get("_all")).get("total");
            baseAllFailedQueries = ((Map<String,Integer>) queriesMetrics.get("_all")).get("failed");
            baseTranslateRequests = ((Map<String,Integer>) queriesMetrics.get("translate")).get("count");
        }
    }
    
    public void testSqlRestUsage() throws IOException {
        index(testData);
        
        //
        // random WHERE and ORDER BY queries
        //
        int randomWhereExecutions = randomIntBetween(1, 15);
        int clientTypeTotalQueries = baseClientTypeTotalQueries + randomWhereExecutions;
        int allTotalQueries = baseAllTotalQueries + randomWhereExecutions;
        
        for (int i = 0; i < randomWhereExecutions; i++) {
            runSql("SELECT name FROM library WHERE page_count > 100 ORDER BY page_count");
        }
        
        Map<String, Object> responseAsMap = getStats();
        assertFeatureMetric(baseMetrics.get("where") + randomWhereExecutions, responseAsMap, "where");
        assertFeatureMetric(baseMetrics.get("orderby") + randomWhereExecutions, responseAsMap, "orderby");
        assertClientTypeAndAllQueryMetrics(clientTypeTotalQueries, allTotalQueries, responseAsMap);
        
        //
        // random HAVING and GROUP BY queries
        //
        int randomHavingExecutions = randomIntBetween(1, 15);
        clientTypeTotalQueries += randomHavingExecutions;
        allTotalQueries += randomHavingExecutions;
        for (int i = 0; i < randomHavingExecutions; i++) {
            runSql("SELECT condition FROM library GROUP BY condition HAVING MAX(page_count) > 1000");
        }
        responseAsMap = getStats();
        assertFeatureMetric(baseMetrics.get("having") + randomHavingExecutions, responseAsMap, "having");
        assertFeatureMetric(baseMetrics.get("groupby") + randomHavingExecutions, responseAsMap, "groupby");
        assertClientTypeAndAllQueryMetrics(clientTypeTotalQueries, allTotalQueries, responseAsMap);
        
        //
        // random LIMIT queries
        //
        int randomLimitExecutions = randomIntBetween(1, 15);
        clientTypeTotalQueries += randomLimitExecutions;
        allTotalQueries += randomLimitExecutions;
        for (int i = 0; i < randomLimitExecutions; i++) {
            runSql("SELECT * FROM library LIMIT " + testData.size());
        }
        responseAsMap = getStats();
        assertFeatureMetric(baseMetrics.get("limit") + randomLimitExecutions, responseAsMap, "limit");
        assertClientTypeAndAllQueryMetrics(clientTypeTotalQueries, allTotalQueries, responseAsMap);
        
        //
        // random LOCALly executed queries
        //
        int randomLocalExecutions = randomIntBetween(1, 15);
        clientTypeTotalQueries += randomLocalExecutions;
        allTotalQueries += randomLocalExecutions;
        for (int i = 0; i < randomLocalExecutions; i++) {
            runSql("SELECT 1+2");
        }
        responseAsMap = getStats();
        assertFeatureMetric(baseMetrics.get("local") + randomLocalExecutions, responseAsMap, "local");
        assertClientTypeAndAllQueryMetrics(clientTypeTotalQueries, allTotalQueries, responseAsMap);
        
        //
        // random COMMANDs
        //
        int randomCommandExecutions = randomIntBetween(1, 15);
        clientTypeTotalQueries += randomCommandExecutions;
        allTotalQueries += randomCommandExecutions;
        for (int i = 0; i < randomCommandExecutions; i++) {
            runSql(randomFrom("SHOW FUNCTIONS", "SHOW COLUMNS FROM library", "SHOW SCHEMAS",
                                 "SHOW TABLES", "SYS CATALOGS", "SYS COLUMNS LIKE '%name'",
                                 "SYS TABLES", "SYS TYPES"));
        }
        responseAsMap = getStats();
        assertFeatureMetric(baseMetrics.get("command") + randomCommandExecutions, responseAsMap, "command");
        assertClientTypeAndAllQueryMetrics(clientTypeTotalQueries, allTotalQueries, responseAsMap);
        
        //
        // random TRANSLATE requests
        //
        int randomTranslateExecutions = randomIntBetween(1, 15);
        for (int i = 0; i < randomTranslateExecutions; i++) {
            runTranslate("SELECT name FROM library WHERE page_count > 100 ORDER BY page_count");
        }
        responseAsMap = getStats();
        assertTranslateQueryMetric(baseTranslateRequests + randomTranslateExecutions, responseAsMap);
        
        //
        // random failed queries
        //
        int randomFailedExecutions = randomIntBetween(1, 15);
        int clientTypeFailedQueries = baseClientTypeFailedQueries + randomFailedExecutions;
        int allFailedQueries = baseAllFailedQueries + randomFailedExecutions;
        allTotalQueries += randomFailedExecutions;
        clientTypeTotalQueries += randomFailedExecutions;
        for (int i = 0; i < randomFailedExecutions; i++) {
            // not interested in the exception type, but in the fact that the metrics are incremented
            // when an exception is thrown
            expectThrows(Exception.class, () -> {
                runSql(randomFrom("SELECT missing_field FROM library",
                                  "SELECT * FROM missing_index",
                                  "SELECTT wrong_command",
                                  "SELECT name + page_count AS not_allowed FROM library",
                                  "SELECT incomplete_query FROM"));
            });
        }
        responseAsMap = getStats();
        assertClientTypeAndAllFailedQueryMetrics(clientTypeFailedQueries, allFailedQueries, responseAsMap);
        assertClientTypeAndAllQueryMetrics(clientTypeTotalQueries, allTotalQueries, responseAsMap);
    }

    private void assertClientTypeAndAllQueryMetrics(int clientTypeTotalQueries, int allTotalQueries, Map<String, Object> responseAsMap)
            throws IOException {
        assertClientTypeQueryMetric(clientTypeTotalQueries, responseAsMap, "total");
        assertAllQueryMetric(allTotalQueries, responseAsMap, "total");
    }
    
    private void assertClientTypeAndAllFailedQueryMetrics(int clientTypeFailedQueries, int allFailedQueries,
            Map<String, Object> responseAsMap) throws IOException {
        assertClientTypeQueryMetric(clientTypeFailedQueries, responseAsMap, "failed");
        assertAllQueryMetric(allFailedQueries, responseAsMap, "failed");
    }
    
    private void index(List<IndexDocument> docs) throws IOException {
        Request request = new Request("POST", "/library/_bulk");
        request.addParameter("refresh", "true");
        StringBuilder bulk = new StringBuilder();
        for (IndexDocument doc : docs) {
            bulk.append("{\"index\":{}}\n");
            bulk.append("{\"condition\":\"" + doc.condition + "\",\"name\":\"" + doc.name + "\",\"page_count\":" + doc.pageCount + "}\n");
        }
        request.setJsonEntity(bulk.toString());
        client().performRequest(request);
    }
    
    private Map<String, Object> getStats() throws UnsupportedOperationException, IOException {
        Request request = new Request("GET", "/_sql/stats");
        Map<String, Object> responseAsMap;
        try (InputStream content = client().performRequest(request).getEntity().getContent()) {
            responseAsMap = XContentHelper.convertToMap(JsonXContent.jsonXContent, content, false);
        }
        
        return responseAsMap;
    }
    
    private void runTranslate(String sql) throws IOException {
        Request request = new Request("POST", "/_sql/translate");
        if (randomBoolean()) {
            // We default to JSON but we force it randomly for extra coverage
            request.addParameter("format", "json");
        }
        if (randomBoolean()) {
            // JSON is the default but randomly set it sometime for extra coverage
            RequestOptions.Builder options = request.getOptions().toBuilder();
            options.addHeader("Accept", randomFrom("*/*", "application/json"));
            request.setOptions(options);
        }
        request.setEntity(new StringEntity("{\"query\":\"" + sql + "\"}", ContentType.APPLICATION_JSON));
        client().performRequest(request);
    }
    
    private void runSql(String sql) throws IOException {
        String mode = Mode.PLAIN.toString();
        if (clientType.equals(ClientType.JDBC.toString())) {
            mode = Mode.JDBC.toString();
        }
        if (clientType.startsWith(ClientType.ODBC.toString())) {
            mode = Mode.ODBC.toString();
        }

        runSql(mode, clientType, sql);
    }
    
    @SuppressWarnings({ "unchecked", "rawtypes" })
    private void assertTranslateQueryMetric(int expected, Map<String, Object> responseAsMap) throws IOException {
        List<Map<String, Map<String, Map>>> nodesListStats = (List) responseAsMap.get("stats");
        int actualMetricValue = 0;
        for (Map perNodeStats : nodesListStats) {
            Map queriesMetrics = (Map) ((Map) perNodeStats.get("stats")).get("queries");
            Map perTypeQueriesMetrics = (Map) queriesMetrics.get("translate");
            actualMetricValue += (int) perTypeQueriesMetrics.get("count");
        }
        assertEquals(expected, actualMetricValue);
    }
    
    private void runSql(String mode, String restClient, String sql) throws IOException {
        Request request = new Request("POST", "/_sql");
        request.addParameter("error_trace", "true");   // Helps with debugging in case something crazy happens on the server.
        request.addParameter("pretty", "true");        // Improves error reporting readability
        if (randomBoolean()) {
            // We default to JSON but we force it randomly for extra coverage
            request.addParameter("format", "json");
        }
        if (randomBoolean()) {
            // JSON is the default but randomly set it sometime for extra coverage
            RequestOptions.Builder options = request.getOptions().toBuilder();
            options.addHeader("Accept", randomFrom("*/*", "application/json"));
            request.setOptions(options);
        }
        request.setEntity(new StringEntity("{\"query\":\"" + sql + "\"" + mode(mode) +
                (ignoreClientType ? StringUtils.EMPTY : ",\"client_id\":\"" + restClient + "\"") + "}",
                ContentType.APPLICATION_JSON));
        client().performRequest(request);
    }
    
    @SuppressWarnings({ "unchecked", "rawtypes" })
    private void assertFeatureMetric(int expected, Map<String, Object> responseAsMap, String feature) throws IOException {
        List<Map<String, ?>> nodesListStats = (List<Map<String, ?>>) responseAsMap.get("stats");
        int actualMetricValue = 0;
        for (Map perNodeStats : nodesListStats) {
            Map featuresMetrics = (Map) ((Map) perNodeStats.get("stats")).get("features");
            actualMetricValue += (int) featuresMetrics.get(feature);
        }
        assertEquals(expected, actualMetricValue);
    }

    @SuppressWarnings({ "unchecked", "rawtypes" })
    private void assertQueryMetric(int expected, Map<String, Object> responseAsMap, String queryType, String metric) throws IOException {
        List<Map<String, Map<String, Map>>> nodesListStats = (List) responseAsMap.get("stats");
        int actualMetricValue = 0;
        for (Map perNodeStats : nodesListStats) {
            Map queriesMetrics = (Map) ((Map) perNodeStats.get("stats")).get("queries");
            Map perTypeQueriesMetrics = (Map) queriesMetrics.get(queryType);
            actualMetricValue += (int) perTypeQueriesMetrics.get(metric);
        }
        assertEquals(expected, actualMetricValue);
    }
    
    private void assertClientTypeQueryMetric(int expected, Map<String, Object> responseAsMap, String metric) throws IOException {
        assertQueryMetric(expected, responseAsMap, clientType, metric);
    }
    
    private void assertAllQueryMetric(int expected, Map<String, Object> responseAsMap, String metric) throws IOException {
        assertQueryMetric(expected, responseAsMap, "_all", metric);
    }
    
    private class IndexDocument {
        private String condition;
        private String name;
        private int pageCount;
        
        IndexDocument(String condition, String name, int pageCount) {
            this.condition = condition;
            this.name = name;
            this.pageCount = pageCount;
        }
    }
}
