/*
 * Copyright 2019 Mitsunori Komatsu (komamitsu)
 *
 * Licensed 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.komamitsu.fluency.aws.s3.ingester.sender;

import com.fasterxml.jackson.databind.util.ByteBufferBackedInputStream;
import net.jodah.failsafe.Failsafe;
import net.jodah.failsafe.RetryPolicy;
import org.komamitsu.fluency.NonRetryableException;
import org.komamitsu.fluency.RetryableException;
import org.komamitsu.fluency.ingester.sender.ErrorHandler;
import org.komamitsu.fluency.ingester.sender.Sender;
import org.komamitsu.fluency.validation.Validatable;
import org.komamitsu.fluency.validation.annotation.DecimalMin;
import org.komamitsu.fluency.validation.annotation.Min;
import org.msgpack.core.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.core.sync.RequestBody;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.S3ClientBuilder;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.StandardOpenOption;
import java.time.temporal.ChronoUnit;
import java.util.zip.GZIPOutputStream;

public class AwsS3Sender
        implements Closeable, Sender
{
    private static final Logger LOG = LoggerFactory.getLogger(AwsS3Sender.class);
    private final Config config;
    private final RetryPolicy retryPolicy;
    private final S3Client client;

    public AwsS3Sender(S3ClientBuilder s3ClientBuilder)
    {
        this(s3ClientBuilder, new Config());
    }

    public AwsS3Sender(S3ClientBuilder s3ClientBuilder, Config config)
    {
        config.validateValues();
        this.config = config;
        this.retryPolicy =
                new RetryPolicy<Void>().
                        handleIf(ex -> {
                            if (ex == null) {
                                // Success. Shouldn't retry.
                                return false;
                            }

                            ErrorHandler errorHandler = config.getErrorHandler();

                            if (errorHandler != null) {
                                errorHandler.handle(ex);
                            }

                            if (ex instanceof InterruptedException || ex instanceof NonRetryableException) {
                                return false;
                            }

                            return true;
                        }).
                        withBackoff(
                                getRetryInternalMs(),
                                getMaxRetryInternalMs(),
                                ChronoUnit.MILLIS,
                                getRetryFactor()).
                        withMaxRetries(getRetryMax());

        this.client = buildClient(s3ClientBuilder);
    }

    @VisibleForTesting
    protected S3Client buildClient(S3ClientBuilder s3ClientBuilder)
    {
        if (config.getEndpoint() != null) {
            try {
                URI uri = new URI(config.getEndpoint());
                s3ClientBuilder.endpointOverride(uri);
            }
            catch (URISyntaxException e) {
                throw new NonRetryableException(
                        String.format("Invalid endpoint. %s", config.getEndpoint()), e);
            }
        }

        if (config.getRegion() != null) {
            s3ClientBuilder.region(Region.of(config.getRegion()));
        }

        if (config.getAwsAccessKeyId() != null && config.getAwsSecretAccessKey() != null) {
            AwsBasicCredentials credentials =
                    AwsBasicCredentials.create(config.getAwsAccessKeyId(), config.getAwsSecretAccessKey());
            s3ClientBuilder.credentialsProvider(StaticCredentialsProvider.create(credentials));
        }

        return s3ClientBuilder.build();
    }

    public int getRetryInternalMs()
    {
        return config.getRetryIntervalMs();
    }

    public int getMaxRetryInternalMs()
    {
        return config.getMaxRetryIntervalMs();
    }

    public float getRetryFactor()
    {
        return config.getRetryFactor();
    }

    public int getRetryMax()
    {
        return config.getRetryMax();
    }

    public int getWorkBufSize()
    {
        return config.getWorkBufSize();
    }

    private void copyStreams(InputStream in, OutputStream out)
            throws IOException
    {
        byte[] buf = new byte[getWorkBufSize()];
        while (true) {
            int readLen = in.read(buf);
            if (readLen < 0) {
                break;
            }
            out.write(buf, 0, readLen);
        }
    }

    private void uploadData(String bucket, String key, File file)
    {
        LOG.debug("Upload data to S3: bucket={}, key={}, fileSize={}",
                bucket, key, file.length());

        try {
            PutObjectRequest.Builder builder = PutObjectRequest.builder()
                    .bucket(bucket)
                    .key(key);
            client.putObject(builder.build(), RequestBody.fromFile(file));
        }
        catch (NonRetryableException e) {
            throw e;
        }
        catch (Throwable e) {
            throw new RetryableException(
                    String.format("Failed to upload data. bucket=%s, key=%s", bucket, key), e);
        }
    }

    public void send(String bucket, String key, ByteBuffer dataBuffer)
            throws IOException
    {
        File file = File.createTempFile("tmp-fluency-", ".tmp");
        try {
            try (InputStream in = new ByteBufferBackedInputStream(dataBuffer);
                    OutputStream fout = Files.newOutputStream(file.toPath(), StandardOpenOption.WRITE);
                    OutputStream out = config.isCompressionEnabled() ? new GZIPOutputStream(fout) : fout) {
                copyStreams(in, out);
            }

            Failsafe.with(retryPolicy).run(() -> uploadData(bucket, key, file));
        }
        finally {
            if (!file.delete()) {
                LOG.warn("Failed to delete a temp file: {}", file.getAbsolutePath());
            }
        }
    }

    @Override
    public void close()
    {
        client.close();
    }

    public static class Config
            extends Sender.Config
            implements Validatable
    {
        private String endpoint;
        private String region;
        private String awsAccessKeyId;
        private String awsSecretAccessKey;
        private boolean isCompressionEnabled = true;

        @Min(10)
        private int retryIntervalMs = 1000;
        @Min(10)
        private int maxRetryIntervalMs = 30000;
        @DecimalMin("1.0")
        private float retryFactor = 2;
        @Min(0)
        private int retryMax = 10;
        @Min(1024)
        private int workBufSize = 8192;

        public String getEndpoint()
        {
            return endpoint;
        }

        public void setEndpoint(String endpoint)
        {
            this.endpoint = endpoint;
        }

        public String getRegion()
        {
            return region;
        }

        public void setRegion(String region)
        {
            this.region = region;
        }

        public String getAwsAccessKeyId()
        {
            return awsAccessKeyId;
        }

        public void setAwsAccessKeyId(String awsAccessKeyId)
        {
            this.awsAccessKeyId = awsAccessKeyId;
        }

        public String getAwsSecretAccessKey()
        {
            return awsSecretAccessKey;
        }

        public void setAwsSecretAccessKey(String awsSecretAccessKey)
        {
            this.awsSecretAccessKey = awsSecretAccessKey;
        }

        public boolean isCompressionEnabled()
        {
            return isCompressionEnabled;
        }

        public void setCompressionEnabled(boolean isCompressionEnabled)
        {
            this.isCompressionEnabled = isCompressionEnabled;
        }

        public int getRetryIntervalMs()
        {
            return retryIntervalMs;
        }

        public void setRetryIntervalMs(int retryIntervalMs)
        {
            this.retryIntervalMs = retryIntervalMs;
        }

        public int getMaxRetryIntervalMs()
        {
            return maxRetryIntervalMs;
        }

        public void setMaxRetryIntervalMs(int maxRetryIntervalMs)
        {
            this.maxRetryIntervalMs = maxRetryIntervalMs;
        }

        public float getRetryFactor()
        {
            return retryFactor;
        }

        public void setRetryFactor(float retryFactor)
        {
            this.retryFactor = retryFactor;
        }

        public int getRetryMax()
        {
            return retryMax;
        }

        public void setRetryMax(int retryMax)
        {
            this.retryMax = retryMax;
        }

        public int getWorkBufSize()
        {
            return workBufSize;
        }

        public void setWorkBufSize(int workBufSize)
        {
            this.workBufSize = workBufSize;
        }

        @Override
        public String toString()
        {
            return "Config{" +
                    "endpoint='" + endpoint + '\'' +
                    ", region='" + region + '\'' +
                    ", isCompressionEnabled=" + isCompressionEnabled +
                    ", retryIntervalMs=" + retryIntervalMs +
                    ", maxRetryIntervalMs=" + maxRetryIntervalMs +
                    ", retryFactor=" + retryFactor +
                    ", retryMax=" + retryMax +
                    ", workBufSize=" + workBufSize +
                    "} " + super.toString();
        }

        void validateValues()
        {
            validate();
        }

    }
}
