/*
 * Copyright The Original Author or Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.jenkins.plugins.opentelemetry.backend.elastic;

import com.cloudbees.plugins.credentials.SystemCredentialsProvider;
import com.cloudbees.plugins.credentials.common.IdCredentials;
import com.cloudbees.plugins.credentials.common.UsernamePasswordCredentials;
import groovy.text.Template;
import hudson.util.FormValidation;
import io.jenkins.plugins.opentelemetry.TemplateBindingsProvider;
import io.jenkins.plugins.opentelemetry.backend.ElasticBackend;
import io.jenkins.plugins.opentelemetry.job.log.ConsoleNotes;
import io.jenkins.plugins.opentelemetry.job.log.LogStorageRetriever;
import io.jenkins.plugins.opentelemetry.job.log.LogsQueryResult;
import io.jenkins.plugins.opentelemetry.job.log.LogsViewHeader;
import io.jenkins.plugins.opentelemetry.semconv.JenkinsOtelSemanticAttributes;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Scope;
import net.sf.json.JSONArray;
import org.apache.commons.lang.StringUtils;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.BasicUserPrincipal;
import org.apache.http.auth.Credentials;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.IndexLifecycleClient;
import org.elasticsearch.client.IndicesClient;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indexlifecycle.GetLifecyclePolicyRequest;
import org.elasticsearch.client.indexlifecycle.GetLifecyclePolicyResponse;
import org.elasticsearch.client.indexlifecycle.LifecyclePolicyMetadata;
import org.elasticsearch.client.indexlifecycle.Phase;
import org.elasticsearch.client.indices.ComposableIndexTemplateExistRequest;
import org.elasticsearch.client.indices.GetComposableIndexTemplateRequest;
import org.elasticsearch.client.indices.GetComposableIndexTemplatesResponse;
import org.elasticsearch.cluster.metadata.ComposableIndexTemplate;
import org.elasticsearch.common.collect.ImmutableOpenMap;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.kohsuke.stapler.framework.io.ByteBuffer;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.Principal;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Optional;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;

/**
 * TODO implement streaming of the logs results. As long as we don't know how to stream the results, we just grab the first lines of logs.
 * <p>
 * Use the old `org.elasticsearch.client:elasticsearch-rest-high-level-client` waiting for
 * `co.elastic.clients:elasticsearch-java` to fix https://github.com/elastic/elasticsearch-java/issues/163
 */
public class ElasticsearchLogStorageRetriever implements LogStorageRetriever<ElasticsearchLogsQueryContext>, Closeable {
    /**
     * Field used by the Elastic-Otel mapping to store the {@link io.opentelemetry.sdk.logs.LogBuilder#setBody(String)}
     */
    public static final String FIELD_MESSAGE = "message";
    public static final String FIELD_LABELS = "labels";
    /**
     * Mapping for `SpanContext#getTraceId()`
     */
    public static final String FIELD_TRACE_ID = "trace.id";
    public static final String FIELD_TIMESTAMP = "@timestamp";

    public static final int PAGE_SIZE = 100;
    public static final String INDEX_TEMPLATE_PATTERNS = "logs-apm.app-*";
    public static final String INDEX_TEMPLATE_NAME = "logs-apm.app";

    private final static Logger logger = Logger.getLogger(ElasticsearchLogStorageRetriever.class.getName());

    @Nonnull
    private final Template buildLogsVisualizationUrlTemplate;

    private final TemplateBindingsProvider templateBindingsProvider;

    @Nonnull
    final String elasticsearchUrl;

    @Nonnull
    private final RestHighLevelClient esClient;

    private final Tracer tracer;

    /**
     * TODO verify unsername:password auth vs apiKey auth
     */
    public ElasticsearchLogStorageRetriever(
        String elasticsearchUrl, Credentials elasticsearchCredentials,
        Template buildLogsVisualizationUrlTemplate, TemplateBindingsProvider templateBindingsProvider,
        Tracer tracer) {
        if (StringUtils.isBlank(elasticsearchUrl)) {
            throw new IllegalArgumentException("Elasticsearch url cannot be blank");
        }

        BasicCredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(AuthScope.ANY, elasticsearchCredentials);

        this.elasticsearchUrl = elasticsearchUrl;
        this.esClient = new RestHighLevelClient(RestClient
            .builder(HttpHost.create(elasticsearchUrl))
            .setHttpClientConfigCallback(httpClient -> httpClient.setDefaultCredentialsProvider(credentialsProvider)));
        this.tracer = tracer;

        this.buildLogsVisualizationUrlTemplate = buildLogsVisualizationUrlTemplate;
        this.templateBindingsProvider = templateBindingsProvider;
    }

    @Nonnull
    @Override
    public LogsQueryResult overallLog(@Nonnull String traceId, @Nonnull String spanId, boolean complete, @Nullable ElasticsearchLogsQueryContext logsQueryContext) throws IOException {
        Span span = tracer.spanBuilder("elasticsearch.search")
            .startSpan();
        try (Scope scope = span.makeCurrent()) {
            Charset charset = StandardCharsets.UTF_8;

            SearchRequest searchRequest = new SearchRequest(ElasticsearchLogStorageRetriever.INDEX_TEMPLATE_PATTERNS)
                .source(new SearchSourceBuilder()
                    .size(PAGE_SIZE)
                    .sort(new FieldSortBuilder(FIELD_TIMESTAMP).order(SortOrder.ASC))
                    .query(QueryBuilders.matchQuery(FIELD_TRACE_ID, traceId)));

            SearchResponse searchResponse = esClient.search(searchRequest, RequestOptions.DEFAULT);
            SearchHit[] hits = searchResponse.getHits().getHits();
            boolean completed = hits.length != PAGE_SIZE; // TODO find better algorithm

            ByteBuffer byteBuffer = new ByteBuffer();

            try (Writer w = new OutputStreamWriter(byteBuffer, charset)) {
                writeOutput(w, hits);
                span.setAttribute("results", hits.length);
                span.setAttribute("completed", completed);
            }

            Map<String, String> localBindings = new HashMap<>();
            localBindings.put("traceId", traceId);
            localBindings.put("spanId", spanId);

            Map<String, String> bindings = TemplateBindingsProvider.compose(this.templateBindingsProvider, localBindings).getBindings();
            String logsVisualizationUrl = this.buildLogsVisualizationUrlTemplate.make(bindings).toString();

            logger.log(Level.FINE, () -> "overallLog(written.length: " + byteBuffer.length() + ")");
            return new LogsQueryResult(
                byteBuffer,
                new LogsViewHeader(bindings.get(ElasticBackend.TemplateBindings.BACKEND_NAME), logsVisualizationUrl, bindings.get(ElasticBackend.TemplateBindings.BACKEND_24_24_ICON_URL)),
                charset, completed,
                new ElasticsearchLogsQueryContext()
            );
        } catch (IOException | RuntimeException e) {
            span.recordException(e);
            throw e;
        } finally {
            span.end();
        }
    }

    @Nonnull
    @Override
    public LogsQueryResult stepLog(@Nonnull String traceId, @Nonnull String spanId, @Nullable ElasticsearchLogsQueryContext logsQueryContext) throws IOException {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    private void writeOutput(Writer writer, SearchHit[] hits) throws IOException {
        for (SearchHit hit : hits) {
            Map<String, Object> source = hit.getSourceAsMap();
            Map<String, Object> labels = (Map<String, Object>) source.get(FIELD_LABELS);
            //Retrieve the label message and annotations to show the formatted message in Jenkins.
            String message = Optional.ofNullable(source.get(FIELD_MESSAGE)).map(m -> m.toString()).orElse(null);
            if (message == null) {
                logger.log(Level.FINE, () -> "Skip log with no message (document id: " + hit.getId() + ")");
                continue;
            }
            JSONArray annotations;
            if (labels == null) {
                annotations = null;
            } else {
                annotations = Optional
                    .ofNullable(labels.get(JenkinsOtelSemanticAttributes.JENKINS_ANSI_ANNOTATIONS.getKey()))
                    .map(a -> JSONArray.fromObject(a.toString())).orElse(null);
            }
            logger.log(Level.FINER, () -> "Write: " + message + ", id: " + hit.getId());
            ConsoleNotes.write(writer, message, annotations);
        }
    }

    /**
     * Check if the configured index template exists.
     * It's a good way to verify the Elastic setup is capable of receiving OpenTelemetry Logs
     *
     * @return true if the index template exists.
     */
    public boolean indexTemplateExists() throws IOException {
        IndicesClient indicesClient = this.esClient.indices();
        return indicesClient.existsIndexTemplate(new ComposableIndexTemplateExistRequest(INDEX_TEMPLATE_NAME), RequestOptions.DEFAULT);
    }

    public List<FormValidation> checkStatus() throws IOException {
        List<FormValidation> validations = new ArrayList<>();
        IndicesClient indicesClient = this.esClient.indices();
        GetComposableIndexTemplatesResponse indexTemplatesResponse;
        try {
            indexTemplatesResponse = indicesClient.getIndexTemplate(new GetComposableIndexTemplateRequest(INDEX_TEMPLATE_NAME), RequestOptions.DEFAULT);
        } catch (IOException e) {
            e.printStackTrace();
            validations.add(FormValidation.error("Exception accessing Elasticsearch " + elasticsearchUrl, e));
            return validations;
        }
        validations.add(FormValidation.ok("Connected to Elasticsearch " + elasticsearchUrl));

        Map<String, ComposableIndexTemplate> indexTemplates = indexTemplatesResponse.getIndexTemplates();
        if (indexTemplates.size() == 0) {
            validations.add(FormValidation.error("Index Template '" + INDEX_TEMPLATE_NAME + " NOT found"));
            return validations;
        } else if (indexTemplates.size() > 1) {
            validations.add(FormValidation.error("More than 1 Index Template matching '" + INDEX_TEMPLATE_NAME + " found: " + indexTemplates.keySet().stream().collect(Collectors.joining(", "))));
        } else {
            validations.add(FormValidation.ok("Index Template '" + INDEX_TEMPLATE_NAME + "' found"));
        }
        ComposableIndexTemplate indexTemplate = indexTemplates.get(INDEX_TEMPLATE_NAME);
        String lifecyclePolicyName;
        try {
            lifecyclePolicyName = indexTemplate.template().settings().getAsSettings("index.lifecycle").get("name");
        } catch (Exception e) {
            validations.add(FormValidation.error("Lifecycle Policy NOT found for Index Template '" + INDEX_TEMPLATE_NAME + "'"));
            return validations;
        }
        IndexLifecycleClient indexLifecycleClient = esClient.indexLifecycle();
        GetLifecyclePolicyResponse lifecyclePolicyResponse = indexLifecycleClient.getLifecyclePolicy(new GetLifecyclePolicyRequest(lifecyclePolicyName), RequestOptions.DEFAULT);
        ImmutableOpenMap<String, LifecyclePolicyMetadata> lifecyclePolicieMetadatas = lifecyclePolicyResponse.getPolicies();
        if (lifecyclePolicieMetadatas.size() == 0) {
            validations.add(FormValidation.error("Lifecycle Policy '" + lifecyclePolicyName + "' NOT found"));
            return validations;
        } else if (lifecyclePolicieMetadatas.size() > 1) {
            validations.add(FormValidation.error("More than 1 Lifecycle Policy found for name '" + lifecyclePolicyName + "': " + lifecyclePolicieMetadatas.stream().map(p -> p.getValue().getName())));
        }
        LifecyclePolicyMetadata lifecyclePolicyMetadata = lifecyclePolicieMetadatas.get(lifecyclePolicyName);
        Map<String, Phase> phases = lifecyclePolicyMetadata.getPolicy().getPhases();
        return validations;
    }

    @Override
    public void close() throws IOException {
        logger.log(Level.INFO, () -> "Shutdown Elasticsearch client...");
        this.esClient.close();
    }

    @Override
    public String toString() {
        return "ElasticsearchLogStorageRetriever{" +
            "buildLogsVisualizationUrlTemplate=" + buildLogsVisualizationUrlTemplate +
            ", templateBindingsProvider=" + templateBindingsProvider +
            '}';
    }

    /**
     * FIXME optimize search
     */
    public static Credentials getCredentials(String jenkinsCredentialsId) throws NoSuchElementException {
        final UsernamePasswordCredentials usernamePasswordCredentials = (UsernamePasswordCredentials) SystemCredentialsProvider.getInstance().getCredentials().stream()
            .filter(credentials ->
                (credentials instanceof UsernamePasswordCredentials)
                    && ((IdCredentials) credentials)
                    .getId().equals(jenkinsCredentialsId))
            .findAny().get();

        return new Credentials() {
            @Override
            public Principal getUserPrincipal() {
                return new BasicUserPrincipal(usernamePasswordCredentials.getUsername());
            }

            @Override
            public String getPassword() {
                return usernamePasswordCredentials.getPassword().getPlainText();
            }
        };
    }

}
