# coding=utf-8

"""
The S3BucketCollector collects bucket size using boto

#### Dependencies

  * boto (https://github.com/boto/boto)
"""

import diamond.collector
import diamond.convertor

try:
    import boto

    boto
    from boto.s3.connection import S3Connection
except ImportError:
    boto = None


class S3BucketCollector(diamond.collector.Collector):
    def get_default_config_help(self):
        config_help = super(S3BucketCollector, self).get_default_config_help()
        config_help.update({})

        return config_help

    def get_default_config(self):
        """
        Returns the default collector settings
        """
        config = super(S3BucketCollector, self).get_default_config()
        config.update({"path": "aws.s3", "byte_unit": "byte"})

        return config

    def getBucket(self, aws_access, aws_secret, bucket_name):
        self.log.info(
            "S3: Open Bucket, %s, %s, %s" % (bucket_name, aws_access, aws_secret)
        )
        s3 = S3Connection(aws_access, aws_secret)

        return s3.lookup(bucket_name)

    def getBucketSize(self, bucket):
        total_bytes = 0

        for key in bucket:
            total_bytes += key.size

        return total_bytes

    def collect(self):
        """
        Collect s3 bucket stats
        """
        if boto is None:
            self.log.error("Unable to import boto python module")
            return {}

        for s3instance in self.config["s3"]:
            self.log.info("S3: byte_unit: %s" % self.config["byte_unit"])
            aws_access = self.config["s3"][s3instance]["aws_access_key"]
            aws_secret = self.config["s3"][s3instance]["aws_secret_key"]

            for bucket_name in self.config["s3"][s3instance]["buckets"]:
                bucket = self.getBucket(aws_access, aws_secret, bucket_name)

                # collect bucket size
                total_size = self.getBucketSize(bucket)

                for byte_unit in self.config["byte_unit"]:
                    new_size = diamond.convertor.binary.convert(
                        value=total_size, old_unit="byte", new_unit=byte_unit
                    )
                    self.publish("%s.size.%s" % (bucket_name, byte_unit), new_size)
