package com.airbnb.billow;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import com.amazonaws.ClientConfiguration;
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.retry.RetryPolicy;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
import com.amazonaws.services.ec2.AmazonEC2Client;
import com.amazonaws.services.ec2.AmazonEC2;
import com.amazonaws.services.ec2.AmazonEC2ClientBuilder;
import com.amazonaws.services.ec2.model.Region;
import com.amazonaws.services.elasticache.AmazonElastiCacheClient;
import com.amazonaws.services.elasticsearch.AWSElasticsearchClient;
import com.amazonaws.services.identitymanagement.AmazonIdentityManagement;
import com.amazonaws.services.identitymanagement.AmazonIdentityManagementClientBuilder;
import com.amazonaws.services.rds.AmazonRDSClient;
import com.amazonaws.services.sqs.AmazonSQSClient;
import com.codahale.metrics.health.HealthCheck;
import com.google.common.collect.Maps;
import com.typesafe.config.Config;

@Slf4j
public class AWSDatabaseHolder {
    private final Map<String, AmazonEC2Client> ec2Clients;
    private final Map<String, AmazonRDSClient> rdsClients;
    private final Map<String, AmazonDynamoDBClient> dynamoDBClients;
    private final Map<String, AmazonSQSClient> sqsClients;
    private final Map<String, AmazonElastiCacheClient> elasticacheClients;
    private final Map<String, AWSElasticsearchClient> elasticsearchClients;
    private final AmazonIdentityManagement iamClient;
    @Getter
    private AWSDatabase current;
    private final long maxAgeInMs;
    private final String awsAccountNumber;
    private final String awsARNPartition;

    public AWSDatabaseHolder(Config config) {
        maxAgeInMs = config.getDuration("maxAge", TimeUnit.MILLISECONDS);

        final DefaultAWSCredentialsProviderChain awsCredentialsProviderChain = new DefaultAWSCredentialsProviderChain();

        final ClientConfiguration clientConfig = new ClientConfiguration();
        clientConfig.setRetryPolicy(new RetryPolicy(null, null, config.getInt("maxErrorRetry"), true));
        clientConfig.setSocketTimeout(config.getInt("socketTimeout") * 1000);

        final AmazonEC2 bootstrapEC2Client = AmazonEC2ClientBuilder.standard().withCredentials(awsCredentialsProviderChain).build();

        ec2Clients = Maps.newHashMap();
        rdsClients = Maps.newHashMap();
        sqsClients = Maps.newHashMap();
        dynamoDBClients = Maps.newHashMap();
        elasticacheClients = Maps.newHashMap();
        elasticsearchClients = Maps.newHashMap();

        final List<Region> ec2Regions = bootstrapEC2Client.describeRegions().getRegions();
        for (Region region : ec2Regions) {
            final String regionName = region.getRegionName();
            final String endpoint = region.getEndpoint();
            log.debug("Adding ec2 region {}", region);

            if (config.getBoolean("ec2Enabled")) {
                final AmazonEC2Client ec2Client = new AmazonEC2Client(awsCredentialsProviderChain, clientConfig);
                ec2Client.setEndpoint(endpoint);
                ec2Clients.put(regionName, ec2Client);
            }

            if (config.getBoolean("rdsEnabled")) {
                final AmazonRDSClient rdsClient = new AmazonRDSClient(awsCredentialsProviderChain, clientConfig);
                rdsClient.setEndpoint(endpoint.replaceFirst("ec2\\.", "rds."));
                rdsClients.put(regionName, rdsClient);
            }

            if (config.getBoolean("dynamodbEnabled")) {
                final AmazonDynamoDBClient dynamoDBClient =
                    new AmazonDynamoDBClient(awsCredentialsProviderChain, clientConfig);
                dynamoDBClient.setEndpoint(endpoint.replaceFirst("ec2\\.", "dynamodb."));
                dynamoDBClients.put(regionName, dynamoDBClient);
            }

            if (config.getBoolean("sqsEnabled")) {
                final AmazonSQSClient sqsClient = new AmazonSQSClient(awsCredentialsProviderChain, clientConfig);
                sqsClient.setEndpoint(endpoint.replaceFirst("ec2\\.", "sqs."));
                sqsClients.put(regionName, sqsClient);
            }

            if (config.getBoolean("elasticacheEnabled")) {
                final AmazonElastiCacheClient elastiCacheClient = new AmazonElastiCacheClient
                    (awsCredentialsProviderChain, clientConfig);
                elastiCacheClient.setEndpoint(endpoint.replaceFirst("ec2\\.", "elasticache."));
                elasticacheClients.put(regionName, elastiCacheClient);
            }

            if (config.getBoolean("elasticsearchEnabled")) {
                final AWSElasticsearchClient elasticsearchClient = new AWSElasticsearchClient
                    (awsCredentialsProviderChain, clientConfig);
                elasticsearchClient.setEndpoint(endpoint.replaceFirst("ec2\\.", "es."));
                elasticsearchClients.put(regionName, elasticsearchClient);
            }
        }

        this.iamClient = AmazonIdentityManagementClientBuilder.standard()
            .withCredentials(awsCredentialsProviderChain)
            .withClientConfiguration(clientConfig)
            .build();

        if (config.hasPath("accountNumber")) {
            this.awsAccountNumber = config.getString("accountNumber");
        } else {
            this.awsAccountNumber = null;
        }

        if (config.hasPath("arnPartition")) {
            this.awsARNPartition = config.getString("arnPartition");
        } else {
            this.awsARNPartition = "aws";
        }

        rebuild();
    }

    public void rebuild() {
        current = new AWSDatabase(
            ec2Clients,
            rdsClients,
            dynamoDBClients,
            sqsClients,
            elasticacheClients,
            elasticsearchClients,
            iamClient,
            awsAccountNumber,
            awsARNPartition);
    }

    public HealthCheck.Result healthy() {
        final long ageInMs = current.getAgeInMs();
        if (ageInMs < maxAgeInMs)
            return HealthCheck.Result.healthy();
        else
            return HealthCheck.Result.unhealthy("DB too old: " + ageInMs + " ms");
    }

    public long getCacheTimeInMs() {
        return maxAgeInMs - current.getAgeInMs();
    }
}