package org.jmal98.metrics.collector;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import com.amazonaws.AmazonClientException;
import com.amazonaws.regions.DefaultAwsRegionProviderChain;
import com.amazonaws.services.sqs.AmazonSQS;
import com.amazonaws.services.sqs.AmazonSQSClientBuilder;
import com.amazonaws.services.sqs.model.GetQueueAttributesResult;
import com.amazonaws.services.sqs.model.ListQueuesResult;

import io.prometheus.client.Collector;
import io.prometheus.client.GaugeMetricFamily;

public class Sqs extends Collector {

	private final Logger logger = LogManager.getLogger(getClass());

	private AmazonSQS sqs = null;

	private List<String> attributeNames = new ArrayList<String>();
	
	private Map<String,String> attributeDescriptions = new HashMap<String,String>();

	@Override
	public List<MetricFamilySamples> collect() {
		List<MetricFamilySamples> mfs = new ArrayList<MetricFamilySamples>();

		try {

			if (sqs == null) {
				String region = new DefaultAwsRegionProviderChain().getRegion();
				sqs = AmazonSQSClientBuilder
								.standard()
								.withRegion(region)
							.build();
				logger.info("AmazonSQS client is connected to region: ({})", region);
			}

			ListQueuesResult queues = sqs.listQueues();
			for (String qUrl : queues.getQueueUrls()) {
				String[] tokens = qUrl.split("\\/");
				String queueName = tokens[tokens.length - 1];

				GetQueueAttributesResult attr = sqs.getQueueAttributes(qUrl, attributeNames);
				Map<String, String> qAttributes = attr.getAttributes();

				for (String key : qAttributes.keySet()) {
					GaugeMetricFamily labeledGauge = new GaugeMetricFamily(
							String.format("sqs_%s", key.toLowerCase().trim()),
							attributeDescriptions.get(key),
							Arrays.asList("queue"));
					
					labeledGauge.addMetric(Arrays.asList(queueName),
							Double.valueOf(qAttributes.get(key)));
					
					mfs.add(labeledGauge);
				}
			}

		} catch (AmazonClientException e) {
			logger.error(e.getMessage());
			if (sqs != null)
				sqs.shutdown();
			sqs = null; // force reconnect
		}

		return mfs;
	}

	public Sqs() {
		super();
		attributeNames.add("ApproximateNumberOfMessages");
		attributeDescriptions.put("ApproximateNumberOfMessages", "The approximate number of visible messages in a queue.");
		
		attributeNames.add("ApproximateNumberOfMessagesDelayed");
		attributeDescriptions.put("ApproximateNumberOfMessagesDelayed", "The approximate number of messages that are waiting to be added to the queue.");
		
		attributeNames.add("ApproximateNumberOfMessagesNotVisible");
		attributeDescriptions.put("ApproximateNumberOfMessagesNotVisible", "The approximate number of messages that have not timed-out and aren't deleted.");
	}

}

@@ -1365,63 +1365,316 @@
 List
-QueuesResult queues = sqs.listQueues
+%3CString%3E queueUrls;%0A%0A%09%09%09// check for manually-specified queue names%0A%09%09%09String queueNames = System.getenv(%22SQS_QUEUE_NAMES%22);%0A%09%09%09if (queueNames != null) %7B%0A%09%09%09    // find the URLs for the named queues%0A%09%09%09    String%5B%5D names = queueNames.split(%22,%22);%0A%09%09%09    queueUrls = new ArrayList%3CString%3E
 ();%0A%09%09%09
+    
 for
- 
 (String 
 qUrl
@@ -1673,36 +1673,289 @@
 ing 
-qUrl : queues.getQueueUrls()
+name : names) %7B%0A%09%09%09%09queueUrls.add(sqs.getQueueUrl(name).getQueueUrl());%0A%09%09%09    %7D%0A%09%09%09%7D else %7B%0A%09%09%09    // get URLs for all queues visible to this account%0A%09%09%09    ListQueuesResult queues = sqs.listQueues();%0A%09%09%09    queueUrls = queues.getQueueUrls();%0A%09%09%09%7D%0A%0A%09%09%09for (String qUrl : queueUrls
 ) %7B%0A

add SQS_QUEUE_NAMES env var to specify queues to query