#!/usr/bin/env python
# -*-python-*-

import pump_transfer
import pump
from cbqueue import PumpQueue

import base64
import optparse
import os
import platform
import re
import simplejson as json
import socket
import subprocess
import sys
import threading
import time
import urllib2

"""Written by Daniel Owen owend@couchbase.com on 27 June 2014
Version 1.4    Last updated 10 July 2014

The current implementation of cbbackup that comes with Couchbase Server 2.5.1
uses only one thead per node.  Therefore when using cbbackup with the single-node
parameter we are limited to one thread - this impacts performance.

This script provides a wrapper to invoke multiple cbbackup processes.
It automatically detects which buckets and vbuckets are
on the node.  It allow the user to specify how many vbuckets to backup in a single
cbbackup process and then invokes the necessary number of processes.
An example invocation is as follows:

python cbbackupwrapper.py http://127.0.0.1:8091 ../backup/ --single-node -n 4 \
-u Administrator -p myPassword --path /opt/couchbbase/bin/  -v

This will backup all the buckets on node 127.0.0.1 into ../backup
It will backup 4 vbuckets per cbbackup process
Access to the cluster is authenticated using username=Administrator and
password=myPassword.and cbbackup will be found in /opt/couchbase/bin

Run python cbbackupwrapper -h for more information.

See the cbrestorewrapper.py script for restoring backups made with this script."""

bucketList = []
vbucketList = []
process_queue = PumpQueue()
lock = threading.Lock()

def opt_extra_help(parser, extra_defaults):
    extra_help = "; ".join(["%s=%s (%s)" %
                           (k, extra_defaults[k][0], extra_defaults[k][1])
                           for k in sorted(extra_defaults.iterkeys())])

    group = optparse.OptionGroup(parser, "Available extra config parameters (-x)",
                        extra_help)
    parser.add_option_group(group)

def opt_extra_defaults():
    return {
        "batch_max_size":  (1000,   "Transfer this # of documents per batch"),
        "batch_max_bytes": (400000, "Transfer this # of bytes per batch"),
        "cbb_max_mb":      (100000, "Split backup file on destination cluster if it exceeds MB"),
        "max_retry":       (10,     "Max number of sequential retries if transfer fails"),
        "report":          (5,      "Number batches transferred before updating progress bar in console"),
        "report_full":     (2000,   "Number batches transferred before emitting progress information in console"),
        "recv_min_bytes":  (4096,   "Amount of bytes for every TCP/IP call transferred"),
        "rehash":          (0,      "For value 1, rehash the partition id's of each item; \
this is needed when transferring data between clusters with different number of partitions, \
such as when transferring data from an OSX server to a non-OSX cluster"),
        "data_only":       (0,      "For value 1, only transfer data from a backup file or cluster"),
        "design_doc_only": (0,      "For value 1, transfer design documents only from a backup file or cluster"),
        "seqno":           (0,      "By default, start seqno from beginning."),
        "backoff_cap":     (10,     "Max backoff time during rebalance period"),
    }

def opt_parse_extra(extra, extra_defaults):
    """Convert an extra string (comma-separated key=val pairs) into
       a dict, using default values from extra_defaults dict."""
    extra_in = dict([(x[0], x[1]) for x in
                     [(kv + '=').split('=') for kv in
                      (extra or "").split(',')]])
    for k, v in extra_in.iteritems():
        if k and not extra_defaults.get(k):
            sys.exit("error: unknown extra option: " + k)
    return dict([(k, float(extra_in.get(k, extra_defaults[k][0])))
                 for k in extra_defaults.iterkeys()])

def argumentParsing():
    usage = "usage: %prog CLUSTER BACKUPDIR OPTIONS"
    parser = optparse.OptionParser(usage)
    opt_extra_help(parser, opt_extra_defaults())

    parser.add_option('-b', '--bucket-source', default='',
                        help='Specify the bucket to backup.  Defaults to all buckets')
    parser.add_option('--single-node', action='store_true',
                        default=False, help='use a single server node from the source only')
    parser.add_option('-u', '--username', default='Administrator',
                        help='REST username for source cluster or server node. Default is Administrator')
    parser.add_option('-p', '--password', default='PASSWORD',
                        help='REST password for source cluster or server node. Defaults to PASSWORD')
    parser.add_option("-s", "--ssl",
                     action="store_true", default=False,
                     help="Transfer data with SSL enabled")
    parser.add_option('-v', '--verbose', action='store_true',
                        default=False, help='Enable verbose messaging')
    parser.add_option('--path', default='.',
                        help='Specify the path to cbbackup. Defaults to current directory')
    parser.add_option('--port', default='11210',
                        help='Specify the bucket port.  Defaults to 11210')
    parser.add_option('-n', '--number', default='100',
                        help='Specify the number of vbuckets per process. Defaults to 100')
    parser.add_option('-P', '--parallelism', default='1',
                        help='Number of vbucket backup jobs to run at a time. Defaults to 1')
    parser.add_option('-x', '--extra', default=None,
                        help="""Provide extra, uncommon config parameters;
                        comma-separated key=val(,key=val)* pairs""")
    try:
        import pump_bfd2
        parser.add_option("-m", "--mode",
                        action="store", type="string", default="diff",
                        help="backup mode: full, diff or accu [default:%default]")
    except ImportError:
        parser.add_option("-m", "--mode",
                        action="store", type="string", default="full",
                        help="backup mode: full")
    options, rest = parser.parse_args()
    if len(rest) != 2:
        parser.print_help()
        sys.exit("\nError: please provide both cluster IP and backup directory path.")

    opt_parse_extra(options.extra, opt_extra_defaults())

    return options, rest[0], rest[1]

def findAllVbucketsForBucket(node, bucket, path, port, restport, username, password, single_node):
    localvbucketlist = []
    request = urllib2.Request(
        'http://' + node + ':' + restport + '/pools/default/buckets/' + bucket)
    base64string = base64.encodestring(
        '%s:%s' % (username, password)).replace('\n', '')
    request.add_header('Authorization', 'Basic %s' % base64string)
    try:
        response = urllib2.urlopen(request)
    except:
        print('Authorization failed.  Please check username and password.')
        sys.exit(1)
    data = json.loads(response.read())
    vbucketserverdata = data['vBucketServerMap']
    vbucketdata = vbucketserverdata['vBucketMap']
    serverlist = vbucketserverdata['serverList']
    # all possibles names / ipaddress for the node
    aliases = []
    # check to see if node was given as ip addess
    matchObj = re.match(r'^\d+.\d+.\d+.\d+$', node, re.I)
    if matchObj:
        # node was entered as its IP address
        nodeip = node
        aliases.append(nodeip)
        try:
            (node, other_names, other_ips) = socket.gethostbyaddr(nodeip)
            aliases.append(node)
            aliases + other_names
        except:
            pass
    else:
        aliases.append(node)
        nodeip = socket.gethostbyname(node)
        aliases.append(nodeip)

    aliases = [alias + ":" + port for alias in aliases]

    if args.verbose:
        print("aliases list is ")
        for x in aliases:
            print(str(x))
        print("server list is")
        for x in serverlist:
            print(str(x))

    # find out the index in the serverlist for this node
    serverindex = -1
    for i in range(len(serverlist)):
        for nodewithport in aliases:
            if nodewithport == serverlist[i]:
                serverindex = i
    if serverindex == -1:
        print serverindex
        print 'Could not find node:port in server list.'
        sys.exit(1)

    if single_node:
        # iterate through all vbuckets and see which are active on this node
        for i in range(len(vbucketdata)):
            if vbucketdata[i][0] == serverindex:
                vbucket = i
                localvbucketlist.append(vbucket)
    else:
        # Just iterate through all vbuckets
        for i in range(len(vbucketdata)):
            vbucket = i
            localvbucketlist.append(vbucket)

    return localvbucketlist


# Get the buckets that exist on the cluster
def getBuckets(node, rest_port, username, password):
    request = urllib2.Request(
        'http://' + node + ':' + rest_port + '/pools/default/buckets')
    base64string = base64.encodestring(
        '%s:%s' % (username, password)).replace('\n', '')
    request.add_header('Authorization', 'Basic %s' % base64string)
    try:
        response = urllib2.urlopen(request)
    except:
        print('Authorization failed.  Please check username and password.')
        sys.exit(1)
    bucketsOnCluster = []
    data = json.loads(response.read())
    for item in data:
        bucket = item['name']
        bucketsOnCluster.append(bucket)
    return bucketsOnCluster

def consumer(id, verbose):
    while True:
        try:
            commandline = process_queue.get(block=True, timeout=10)
            if verbose:
                with lock:
                    print "T(%s): %s" %(id, commandline)
            p = subprocess.Popen(commandline, shell=True)
            p.wait()
            if p.returncode == 1:
                with lock:
                    print 'Error with backup for running %s' % commandline
            process_queue.task_done()
            time.sleep(1)
        except Exception:
            if process_queue.empty():
                return
            print 'Exception ' + str(e)

if __name__ == '__main__':
    # Parse the arguments given.
    args, cluster, backupDir = argumentParsing()

    backup_exe = 'cbbackup'
    if platform.system() == "Windows":
        backup_exe = 'cbbackup.exe'

    # Remove any white-spaces from start and end of strings
    backupDir = backupDir.strip()
    path = args.path.strip()
    if path == ".":
        path = os.path.abspath(path)

    # Check to see if root backup directory exists
    if not os.path.isdir(backupDir):
        try:
            os.makedirs(backupDir)
        except:
            sys.exit("Cannot create backup root directory:%s" % backupDir)

    # Check to see if path is correct
    if not os.path.isdir(path):
        print 'The path to cbbackup does not exist'
        print 'Please run with a different path'
        sys.exit(1)
    if not os.path.isfile(os.path.join(path, backup_exe)):
        print 'cbbackup could not be found in ' + path
        sys.exit(1)

    # Check to see if log directory exists if not create it
    dir = os.path.join(backupDir, 'logs')
    try:
        os.stat(dir)
    except:
        try:
            os.mkdir(dir)
        except:
            print('Error trying to create directory ' + dir)
            sys.exit(1)

    # Separate out node and REST port
    matchObj = re.match(r'^http://(.*):(\d+)$', cluster, re.I)
    if matchObj:
        node = matchObj.group(1)
        rest = matchObj.group(2)
    else:
        print("Please enter the source as http://hostname:port")
        print("For example http://localhost:8091 or http://127.0.0.1:8091")
        sys.exit(1)

    # Check to see if backing-up all buckets or just a specified bucket
    if args.bucket_source == '':
        bucketList = getBuckets(
            node, rest, args.username, args.password)
    else:
        # Check that the bucket exists
        for item in getBuckets(node, rest, args.username, args.password):
            if item == args.bucket_source:
                bucketList.append(args.bucket_source)

        if len(bucketList) == 0:
            print 'Bucket ' + args.bucket_source + ' does not exist'
            print 'Please enter a different bucket'
            sys.exit(1)

    # For each bucket
    for item in bucketList:
        perbucketvbucketlist = findAllVbucketsForBucket(
            node, item, path, args.port, rest, args.username, args.password, args.single_node)
        for item in perbucketvbucketlist:
            if item not in vbucketList:
                vbucketList.append(item)

    # Handle the case when path has spaces
    # i.e. /Applications/Couchbase Server.app/Contents/...
    if os.name == 'nt':
        path = re.sub(r' ', '^ ', path)
    else:
        path = re.sub(r' ', '\ ', path)

    # If a bucket was specfified then set-up the string to pass to cbbackup.
    specific_bucket = ''
    if len(bucketList) == 1:
        specific_bucket = ' -b ' + bucketList[0]

    extra_options = ''
    if args.extra:
        extra_options = ' -x ' + args.extra

    mode_options = ''
    if args.mode:
        mode_options = ' -m ' + args.mode

    ssl_option = ''
    if args.ssl:
        ssl_option = ' -s '

    for i in range(int(args.parallelism)):
        t = threading.Thread(target=consumer, args=(i, args.verbose,))
        t.daemon = True
        t.start()

    # Group the number of vbuckets per process
    print 'Waiting for the backup to complete...'
    processes = []
    for i in range(0, len(vbucketList), int(args.number)):
        chunk = vbucketList[i:i + int(args.number)]
        vbucketsname = str(chunk[0]) + '-' + str(chunk[-1])
        command_line = os.path.join(path, backup_exe) + ' -v -t 1 --vbucket-list=' + ''.join(str(chunk).split()) + ' http://' \
            + node + ':' + rest + ' ' + os.path.join(backupDir, vbucketsname) + ' -u ' + args.username \
            + ' -p ' + args.password + extra_options + mode_options + ssl_option + specific_bucket + ' 2>' + \
            os.path.join(backupDir, 'logs', vbucketsname) + '.err'
        process_queue.put(command_line)

    process_queue.join()

    with lock:
        print 'SUCCESSFULLY COMPLETED!'
