#!/usr/bin/env python
# s3funnel - Multithreaded tool for performing operations on Amazon's S3
# Copyright (c) 2008 Andrey Petrov
#
# This module is part of s3funnel and is released under
# the MIT license: http://www.opensource.org/licenses/mit-license.php

"""
s3funnel is a multithreaded tool for performing operations on Amazon's S3.

Key Operations:
    DELETE Delete key from the bucket
    GET    Get key from the bucket
    PUT    Put file into the bucket (key is the basename of the path)
    COPY   Copy keys to the bucket from another bucket.

Bucket Operations:
    CREATE Create a new bucket
    DROP   Delete an existing bucket (must be empty)
    LIST   List keys in the bucket. If no bucket is given, buckets will be listed.
"""

__VERSION__ = '0.6.3'

"""
TODO:
* Use callback to do an interactive progress bar
"""

import logging

log = logging.getLogger(__name__)
s3funnel_log = logging.getLogger('s3funnel')
boto_log = logging.getLogger('boto')

def set_log_level(level):
    for i in [log, s3funnel_log, boto_log]:
        i.setLevel(level)


# Official Python modules
import os
import sys
import signal
import threading
import socket
import httplib

from glob import glob
from optparse import OptionParser, SUPPRESS_HELP

# Third party modules
import workerpool

# Local imports
from s3funnel import S3Funnel
from s3funnel.exceptions import FunnelError

def main():
    # Parse the command line...
    usage="%prog BUCKET OPERATION [OPTIONS] [FILE]...\n" + __doc__
    parser = OptionParser(usage)
    parser.add_option("-a", "--aws_key",    dest="aws_key", type="string", help="Overrides AWS_ACCESS_KEY_ID environment variable")
    parser.add_option("-s", "--aws_secret_key", dest="aws_secret_key", type="string", help="Overrides AWS_SECRET_ACCESS_KEY environment variable")
    parser.add_option("-t", "--threads",    dest="numthreads", default=1, type="int", metavar="N", help="Number of threads to use [default: %default]")
    parser.add_option("-T", "--timeout",    dest="timeout", default=0, type="float", metavar="SECONDS", help="Socket timeout time, 0 is never [default: %default]")
    parser.add_option("--insecure",         dest="secure", action="store_false", default=True, help="Don't use secure (https) connection")
    parser.add_option("--list-marker",      dest="list_marker", type="string", default=None, metavar="KEY", help="(`list` only) Start key for list operation")
    parser.add_option("--list-prefix",      dest="list_prefix", type="string", default=None, metavar="STRING", help="(`list` only) Limit results to a specific prefix")
    parser.add_option("--list-delimiter",   dest="list_delimiter", type="string", default=None, metavar="CHAR", help="(`list` only) Treat value as a delimiter for hierarchical listing")
    parser.add_option("--put-acl",          dest="acl", type="string", default="public-read", help="(`put` only) Set the ACL permission for each file [default: %default]")
    parser.add_option("--put-full-path",    dest="put_full_path", action="store_true", help="(`put` only) Use the full given path as the key name, instead of just the basename")
    parser.add_option("--put-only-new",     dest="put_only_new", action="store_true", help="(`put` only) Only PUT keys which don't already exist in the bucket with the same md5 digest")
    parser.add_option("--put-header",       dest="headers", type="string", action="append", help="(`put` only) Add the specified header to the request")
    parser.add_option("--add-prefix",       dest="add_prefix", type="string", default="", help="(`put` and `copy` only) Add specified prefix to keys in destination bucket")
    parser.add_option("--del-prefix",       dest="del_prefix", type="string", default="", help="(`put` and `copy` only) Delete specified prefix from keys in destination bucket")
    parser.add_option("--source-bucket",    dest="source_bucket", type="string", help="(`copy` only) Source bucket for files")
    parser.add_option("--no-ignore-s3fs-dirs",    dest="ignore_s3fs_dirs", action="store_false", default=True, help="(`get` only) Don't ignore s3fs directory objects ")
    parser.add_option("-i", "--input",      dest="input", type="string", metavar="FILE", help="Read one file per line from a FILE manifest")
    parser.add_option("-v", "--verbose",    dest="verbose", action="count", default=None, help="Enable verbose output. Use twice to enable debug output")
    parser.add_option("--version",          dest="version", action="store_true", help="Output version information and exit")

    # Deprecated options (for backwards compatibility)
    parser.add_option("--start_key",        dest="list_marker", type="string", default=None, help=SUPPRESS_HELP)
    parser.add_option("--acl",              dest="acl", type="string", default="public-read", help=SUPPRESS_HELP)

    # Hidden options for debugging, development, and lols
    # TODO: ...

    options, args = parser.parse_args()

    # Version?
    if options.version:
        print "s3funnel %s" % __VERSION__
        return 0

    # Check input
    aws_key = os.environ.get("AWS_ACCESS_KEY_ID")
    aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")

    ## AWS
    if options.aws_key:
        aws_key = options.aws_key
    if options.aws_secret_key:
        aws_secret_key = options.aws_secret_key
    if None in [aws_key, aws_secret_key]:
        parser.error("Missing required arguments `aws_key` or `aws_secret_key`")

    ## Threads
    if options.numthreads < 1:
        parser.error("`theads` must be at least 1")

    ## Misc. options
    if options.timeout:
        try:
            socket.setdefaulttimeout(options.timeout)
        except TypeError, e:
            parser.error("`timeout` error: %s" % e.message)

    ## Parse put headers
    headers = {}
    if options.headers:
        for header in options.headers:
            if ':' not in header:
                parser.error("Header must be use ':' separator")
            key, value = header.split(':', 2)
            headers[key.strip()] = value.strip()

    # Arguments
    if len(args) < 1:
        parser.error("BUCKET not specified")
    bucket = args[0]
    if len(args) < 2:
        # Exception for single-argument operations
        if bucket.lower() in ['list']:
            operation = bucket.lower()
            bucket = None
        else:
            parser.error("OPERATION not specified")
    else:
        operation = args[1].lower()

    if operation == 'list' and not bucket:
        operation = 'show'
        
    if options.add_prefix and options.del_prefix:
        parser.error("--add-prefix and --del-prefix options are mutually exclusive")
        return -1
    
    if operation == 'copy' and not options.source_bucket:
        parser.error("Source bucket not specified")
        return -1

    # Setup logging
    if options.verbose > 1:
        set_log_level(logging.DEBUG)
    elif options.verbose > 0:
        set_log_level(logging.INFO)

    # Setup operation configuration
    config = {'acl': options.acl,
              'list_marker': options.list_marker or '',
              'list_prefix': options.list_prefix or '',
              'list_delimiter': options.list_delimiter or '',
              'aws_key': aws_key,
              'aws_secret_key': aws_secret_key,
              'secure': options.secure,
              'bucket': bucket,
              'source_bucket': options.source_bucket, 
              'put_full_path': options.put_full_path,
              'put_only_new': options.put_only_new,
              'add_prefix': options.add_prefix,
              'del_prefix': options.del_prefix,
              'headers': headers,
              'numthreads': options.numthreads,
              'ignore_s3fs_dirs': options.ignore_s3fs_dirs
              }

    funnel = S3Funnel(**config)

    # Setup operation mapping
    methods_keys = {
       'get':    funnel.get,
       'put':    funnel.put,
       'delete': funnel.delete,
       'copy':   funnel.copy,
    }

    methods_bucket = {
       'list':   funnel.list_bucket,
       'drop':   funnel.drop_bucket,
       'create': funnel.create_bucket,
    }

    methods_global = {
       'show':   funnel.show_buckets,
    }

    valid_operations = methods_keys.keys() + methods_bucket.keys() + methods_global.keys()
    valid_operations.sort()
    if operation not in valid_operations:
        parser.error("OPERATION must be one of: %s" % ', '.join(valid_operations))

    # Get data source
    input_src = None
    if options.input:
        # Get source from manifest or stdin (via -i flag)
        if options.input == '-':
            input_src = "stdin"
            options.input = sys.stdin
        try:
            data = open(glob(options.input)[0])
        except (IOError, IndexError), e:
            log.error("%s: File not found" % options.input)
            return -1
        input_src = "`%s'" % options.input
    elif len(args) < 3:
        # Get source from stdin
        input_src = "stdin"
        data = sys.stdin
    else:
        if operation == 'put':
            # Get source from glob-expanded arguments
            data = []
            for arg in args[2:]:
                found = glob(arg)
                if not found:
                    log.error("%s: No such file." % arg)
                    continue
                data += found
        else:
            data = args[2:]
        input_src = "arguments: %s" % ', '.join(data)

    # Setup interrupt handling
    def shutdown(signum, stack):
        log.warning("Interrupted, shutting down...")
        funnel.shutdown()
        sys.exit()
    signal.signal(signal.SIGINT, shutdown)

    # Setup output logger
    output = logging.getLogger('output')
    output_handler = logging.StreamHandler(sys.stdout)
    output_handler.setFormatter(logging.Formatter('%(message)s'))
    output.addHandler(output_handler)
    output.setLevel(logging.INFO)

    # Feed input into the appropriate method
    log.info("Using input from %s" % input_src)

    try:
        # TODO: Rewrite this into something fancier
        if operation in methods_global:
            m = methods_global[operation]
            for i in m():
                print i
        elif operation in methods_bucket:
            m = methods_bucket[operation]
            for i in m(bucket, **config):
                print i
        elif operation in methods_keys:
            m = methods_keys[operation]
            r = m(ikeys=(d.strip() for d in data), **config)
            if r:
                log.critical("%d keys failed." % len(r))
                return len(r)
    except TypeError, e:
        pass
    except FunnelError, e:
        log.critical(e)
        return -1

    funnel.shutdown()

if __name__ == "__main__":
    log_handler = logging.StreamHandler()
    log_handler.setFormatter(logging.Formatter('%(levelname)-8s %(message)s'))

    log.addHandler(log_handler)
    s3funnel_log.addHandler(log_handler)
    boto_log.addHandler(log_handler)

    r = main()
    if r:
        sys.exit(r)