/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch 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.elasticsearch.repositories.s3;

import com.amazonaws.AmazonClientException;
import com.amazonaws.AmazonServiceException;
import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.model.AmazonS3Exception;
import com.amazonaws.services.s3.model.ObjectMetadata;
import com.amazonaws.services.s3.model.PutObjectResult;
import com.amazonaws.services.s3.model.S3Object;
import com.amazonaws.services.s3.model.UploadPartRequest;
import com.amazonaws.services.s3.model.UploadPartResult;
import org.apache.logging.log4j.Logger;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.logging.Loggers;
import org.elasticsearch.common.settings.Settings;

import java.io.IOException;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;

import static com.carrotsearch.randomizedtesting.RandomizedTest.randomDouble;

public class TestAmazonS3 extends AmazonS3Wrapper {

    protected final Logger logger = Loggers.getLogger(getClass());

    private double writeFailureRate = 0.0;
    private double readFailureRate = 0.0;

    private String randomPrefix;

    ConcurrentMap<String, AtomicLong> accessCounts = new ConcurrentHashMap<String, AtomicLong>();

    private long incrementAndGet(String path) {
        AtomicLong value = accessCounts.get(path);
        if (value == null) {
            value = accessCounts.putIfAbsent(path, new AtomicLong(1));
        }
        if (value != null) {
            return value.incrementAndGet();
        }
        return 1;
    }

    public TestAmazonS3(AmazonS3 delegate, Settings settings) {
        super(delegate);
        randomPrefix = settings.get("cloud.aws.test.random");
        writeFailureRate = settings.getAsDouble("cloud.aws.test.write_failures", 0.0);
        readFailureRate = settings.getAsDouble("cloud.aws.test.read_failures", 0.0);
    }

    @Override
    public PutObjectResult putObject(String bucketName, String key, InputStream input, ObjectMetadata metadata) throws AmazonClientException, AmazonServiceException {
        if (shouldFail(bucketName, key, writeFailureRate)) {
            long length = metadata.getContentLength();
            long partToRead = (long) (length * randomDouble());
            byte[] buffer = new byte[1024];
            for (long cur = 0; cur < partToRead; cur += buffer.length) {
                try {
                    input.read(buffer, 0, (int) (partToRead - cur > buffer.length ? buffer.length : partToRead - cur));
                } catch (IOException ex) {
                    throw new ElasticsearchException("cannot read input stream", ex);
                }
            }
            logger.info("--> random write failure on putObject method: throwing an exception for [bucket={}, key={}]", bucketName, key);
            AmazonS3Exception ex = new AmazonS3Exception("Random S3 exception");
            ex.setStatusCode(400);
            ex.setErrorCode("RequestTimeout");
            throw ex;
        } else {
            return super.putObject(bucketName, key, input, metadata);
        }
    }

    @Override
    public UploadPartResult uploadPart(UploadPartRequest request) throws AmazonClientException, AmazonServiceException {
        if (shouldFail(request.getBucketName(), request.getKey(), writeFailureRate)) {
            long length = request.getPartSize();
            long partToRead = (long) (length * randomDouble());
            byte[] buffer = new byte[1024];
            for (long cur = 0; cur < partToRead; cur += buffer.length) {
                try (InputStream input = request.getInputStream()){
                    input.read(buffer, 0, (int) (partToRead - cur > buffer.length ? buffer.length : partToRead - cur));
                } catch (IOException ex) {
                    throw new ElasticsearchException("cannot read input stream", ex);
                }
            }
            logger.info("--> random write failure on uploadPart method: throwing an exception for [bucket={}, key={}]", request.getBucketName(), request.getKey());
            AmazonS3Exception ex = new AmazonS3Exception("Random S3 write exception");
            ex.setStatusCode(400);
            ex.setErrorCode("RequestTimeout");
            throw ex;
        } else {
            return super.uploadPart(request);
        }
    }

    @Override
    public S3Object getObject(String bucketName, String key) throws AmazonClientException, AmazonServiceException {
        if (shouldFail(bucketName, key, readFailureRate)) {
            logger.info("--> random read failure on getObject method: throwing an exception for [bucket={}, key={}]", bucketName, key);
            AmazonS3Exception ex = new AmazonS3Exception("Random S3 read exception");
            ex.setStatusCode(404);
            throw ex;
        } else {
            return super.getObject(bucketName, key);
        }
    }

    private boolean shouldFail(String bucketName, String key, double probability) {
        if (probability > 0.0) {
            String path = randomPrefix + "-" + bucketName + "+" + key;
            path += "/" + incrementAndGet(path);
            return Math.abs(hashCode(path)) < Integer.MAX_VALUE * probability;
        } else {
            return false;
        }
    }

    private int hashCode(String path) {
        try {
            MessageDigest digest = MessageDigest.getInstance("MD5");
            byte[] bytes = digest.digest(path.getBytes("UTF-8"));
            int i = 0;
            return ((bytes[i++] & 0xFF) << 24) | ((bytes[i++] & 0xFF) << 16)
                    | ((bytes[i++] & 0xFF) << 8) | (bytes[i++] & 0xFF);
        } catch (UnsupportedEncodingException ex) {
            throw new ElasticsearchException("cannot calculate hashcode", ex);
        } catch (NoSuchAlgorithmException ex) {
            throw new ElasticsearchException("cannot calculate hashcode", ex);
        }
    }
}
