package com.amazonaws.services.sqs.util;

import java.util.Collections;
import java.util.concurrent.ThreadLocalRandom;

import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider;
import com.amazonaws.auth.policy.Policy;
import com.amazonaws.auth.policy.Principal;
import com.amazonaws.auth.policy.Resource;
import com.amazonaws.auth.policy.Statement;
import com.amazonaws.auth.policy.actions.SQSActions;
import com.amazonaws.auth.profile.ProfileCredentialsProvider;
import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder;
import com.amazonaws.services.securitytoken.model.AssumeRoleRequest;
import com.amazonaws.services.securitytoken.model.AssumeRoleResult;
import com.amazonaws.services.securitytoken.model.Credentials;
import com.amazonaws.services.securitytoken.model.GetCallerIdentityRequest;
import com.amazonaws.services.sqs.model.AmazonSQSException;
import org.junit.After;
import org.junit.Before;

import com.amazonaws.services.sqs.AmazonSQS;
import com.amazonaws.services.sqs.AmazonSQSClientBuilder;
import com.amazonaws.services.sqs.model.QueueDoesNotExistException;

import static org.hamcrest.CoreMatchers.equalTo;
import static org.junit.Assume.assumeNoException;
import static org.junit.Assume.assumeThat;
import static org.junit.Assume.assumeTrue;

/**
 * Base class for integration tests
 */
public abstract class IntegrationTest {
    
    protected AmazonSQS sqs;
    // UUIDs are too long for this
    protected String queueNamePrefix = "__" + testSuiteName() + "-" + ThreadLocalRandom.current().nextInt(1000000);
    protected ExceptionAsserter exceptionHandler = new ExceptionAsserter();

    /**
     * Customizable to ensure queue names stay under 80 characters
     */
    protected String testSuiteName() {
        return getClass().getSimpleName();
    }

    @Before
    public void setupSQSClient() {
        sqs = AmazonSQSClientBuilder.defaultClient();
    }
    
    @After
    public void teardownSQSClient() {
        if (sqs != null) {
            // Best effort cleanup of queues. To be complete, we'd have to wait a minute
            // for the eventual consistency of listQueues()
            sqs.listQueues(queueNamePrefix).getQueueUrls().forEach(queueUrl -> {
                try {
                    sqs.deleteQueue(queueUrl);
                } catch (QueueDoesNotExistException e) {
                    // Ignore
                }
            });
            sqs.shutdown();
        }
        exceptionHandler.assertNothingThrown();
    }

    protected String getBuddyRoleARN() {
        String roleARN = System.getenv("BUDDY_ROLE_ARN");
        if (roleARN == null) {
            assumeTrue("This test requires a second 'buddy' AWS role, provided with the BUDDY_ROLE_ARN environment variable.", false);
        }
        return roleARN;
    }

    protected AWSCredentialsProvider getBuddyCredentials() {
        return new STSAssumeRoleSessionCredentialsProvider.Builder(getBuddyRoleARN(), testSuiteName()).build();
    }

    protected AmazonSQS getBuddyPrincipalClient() {
        AWSCredentialsProvider credentialsProvider = getBuddyCredentials();
        AmazonSQS client = AmazonSQSClientBuilder.standard()
                .withRegion("us-west-2")
                .withCredentials(credentialsProvider)
                .build();

        // Assume that the principal is not able to send messages to arbitrary queues
        String queueUrl = sqs.createQueue(queueNamePrefix + "TestQueue").getQueueUrl();
        try {
            client.sendMessage(queueUrl, "Haxxors!!");
            assumeTrue("The buddy credentials should not authorize sending to arbitrary queues", false);
        } catch (AmazonSQSException e) {
            // Access Denied
            assumeThat(e.getStatusCode(), equalTo(403));
        } finally {
            sqs.deleteQueue(queueUrl);
        }

        return client;
    }

    protected Policy allowSendMessagePolicy(String roleARN) {
        Policy policy = new Policy();
        Statement statement = new Statement(Statement.Effect.Allow);
        statement.setActions(Collections.singletonList(SQSActions.SendMessage));
        statement.setPrincipals(new Principal(roleARN));
        statement.setResources(Collections.singletonList(new Resource("arn:aws:sqs:*:*:*")));
        policy.setStatements(Collections.singletonList(statement));
        return policy;
    }
}