/*
 * 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.dubbo.spring.boot.observability.autoconfigure.exporter.zipkin;

import org.springframework.http.HttpHeaders;
import org.springframework.util.unit.DataSize;
import zipkin2.Call;
import zipkin2.CheckResult;
import zipkin2.codec.Encoding;
import zipkin2.reporter.BytesMessageEncoder;
import zipkin2.reporter.ClosedSenderException;
import zipkin2.reporter.Sender;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.Collections;
import java.util.List;
import java.util.zip.GZIPOutputStream;

/**
 * A Zipkin {@link Sender} that uses an HTTP client to send JSON spans. Supports automatic compression with gzip.
 */
abstract class HttpSender extends Sender {
    private static final DataSize MESSAGE_MAX_SIZE = DataSize.ofKilobytes(512);

    private volatile boolean closed;

    @Override
    public Encoding encoding() {
        return Encoding.JSON;
    }

    @Override
    public int messageMaxBytes() {
        return (int) MESSAGE_MAX_SIZE.toBytes();
    }

    @Override
    public int messageSizeInBytes(List<byte[]> encodedSpans) {
        return encoding().listSizeInBytes(encodedSpans);
    }

    @Override
    public int messageSizeInBytes(int encodedSizeInBytes) {
        return encoding().listSizeInBytes(encodedSizeInBytes);
    }

    @Override
    public CheckResult check() {
        try {
            sendSpans(Collections.emptyList()).execute();
            return CheckResult.OK;
        } catch (IOException | RuntimeException ex) {
            return CheckResult.failed(ex);
        }
    }

    @Override
    public void close() throws IOException {
        this.closed = true;
    }

    /**
     * The returned {@link HttpPostCall} will send span(s) as a POST to a zipkin endpoint
     * when executed.
     *
     * @param batchedEncodedSpans list of encoded spans as a byte array
     * @return an instance of a Zipkin {@link Call} which can be executed
     */
    protected abstract HttpPostCall sendSpans(byte[] batchedEncodedSpans);

    @Override
    public Call<Void> sendSpans(List<byte[]> encodedSpans) {
        if (this.closed) {
            throw new ClosedSenderException();
        }
        return sendSpans(BytesMessageEncoder.JSON.encode(encodedSpans));
    }

    abstract static class HttpPostCall extends Call.Base<Void> {

        /**
         * Only use gzip compression on data which is bigger than this in bytes.
         */
        private static final DataSize COMPRESSION_THRESHOLD = DataSize.ofKilobytes(1);

        private final byte[] body;

        HttpPostCall(byte[] body) {
            this.body = body;
        }

        protected byte[] getBody() {
            if (needsCompression()) {
                return compress(this.body);
            }
            return this.body;
        }

        protected byte[] getUncompressedBody() {
            return this.body;
        }

        protected HttpHeaders getDefaultHeaders() {
            HttpHeaders headers = new HttpHeaders();
            headers.set("b3", "0");
            headers.set("Content-Type", "application/json");
            if (needsCompression()) {
                headers.set("Content-Encoding", "gzip");
            }
            return headers;
        }

        private boolean needsCompression() {
            return this.body.length > COMPRESSION_THRESHOLD.toBytes();
        }

        private byte[] compress(byte[] input) {
            ByteArrayOutputStream result = new ByteArrayOutputStream();
            try (GZIPOutputStream gzip = new GZIPOutputStream(result)) {
                gzip.write(input);
            } catch (IOException ex) {
                throw new UncheckedIOException(ex);
            }
            return result.toByteArray();
        }

    }
}
