import sys
from django.core.management.base import BaseCommand, CommandError
from dimagi.utils.couch.database import get_db, iter_docs
from corehq.apps.domainsync.config import DocumentTransform, save
from couchdbkit.client import Database
from optparse import make_option
from datetime import datetime

# doctypes we want to be careful not to copy, which must be explicitly
# specified with --include
DEFAULT_EXCLUDE_TYPES = [
    'ReportNotification',
    'WeeklyNotification',
    'DailyNotification'
]

class Command(BaseCommand):
    help = "Copies the contents of a domain to another database."
    args = '<sourcedb> <domain>'
    option_list = BaseCommand.option_list + (
        make_option('--include',
                    action='store',
                    dest='doc_types',
                    default='',
                    help='Comma-separated list of Document Types to copy'),
        make_option('--exclude',
                    action='store',
                    dest='doc_types_exclude',
                    default='',
                    help='Comma-separated list of Document Types to NOT copy.'),
        make_option('--since',
                    action='store',
                    dest='since',
                    default='',
                    help='Only copy documents newer than this date. Format: yyyy-MM-dd. Only '),
        make_option('--list-types',
                    action='store_true',
                    dest='list_types',
                    default=False,
                    help='Don\'t copy anything, just list all the available document types.'),
        make_option('--simulate',
                    action='store_true',
                    dest='simulate',
                    default=False,
                    help='Don\'t copy anything, print what would be copied.'),
    )

    def handle(self, *args, **options):
        
        if len(args) != 2:
            raise CommandError('Usage is copy_domain %s' % self.args)

        sourcedb = Database(args[0])
        domain = args[1].strip()
        simulate = options['simulate']

        since = datetime.strptime(options['since'], '%Y-%m-%d').isoformat() if options['since'] else None

        if options['list_types']:
            self.list_types(sourcedb, domain)
            sys.exit(0)

        if simulate:
            print "\nSimulated run, no data will be copied.\n"

        if options['doc_types']:
            doc_types = options['doc_types'].split(',')
            for type in doc_types:
                startkey = [x for x in [domain, type, since] if x is not None]
                endkey = [x for x in [domain, type, {}] if x is not None]
                self.copy_docs(sourcedb, domain, startkey, endkey, simulate, type=type, since=since)
        else:
            startkey = [domain]
            endkey = [domain, {}]
            exclude_types = DEFAULT_EXCLUDE_TYPES + options['doc_types_exclude'].split(',')
            self.copy_docs(sourcedb, domain, startkey, endkey, simulate, exclude_types=exclude_types)

    def list_types(self, sourcedb, domain):
        doc_types = sourcedb.view("domain/docs", startkey=[domain],
                                  endkey=[domain, {}], reduce=True, group=True, group_level=2)
        for row in doc_types:
            print "{:<30}- {}".format(row['key'][1], row['value'])

    def copy_docs(self, sourcedb, domain, startkey, endkey, simulate, type=None, since=None, exclude_types=None):
        doc_ids = [result["id"] for result in sourcedb.view("domain/docs", startkey=startkey,
                             endkey=endkey, reduce=False)]

        total = len(doc_ids)
        count = 0
        targetdb = get_db()
        msg = "Found %s matching documents in domain: %s" % (total, domain)
        msg += " of type: %s" % (type) if type else ""
        msg += " since: %s" % (since) if since else ""
        print msg
        for doc in iter_docs(sourcedb, doc_ids):
            try:
                count += 1
                if exclude_types and doc["doc_type"] in exclude_types:
                    print "     SKIPPED (excluded type: %s). Synced %s/%s docs (%s: %s)" % \
                          (doc["doc_type"], count, total, doc["doc_type"], doc["_id"])
                else:
                    if not simulate:
                        dt = DocumentTransform(doc, sourcedb)
                        save(dt, targetdb)
                    print "     Synced %s/%s docs (%s: %s)" % (count, total, doc["doc_type"], doc["_id"])
            except Exception, e:
                print "     Document %s failed! Error is: %s" % (doc["_id"], e)

@@ -1,12 +1,55 @@
+from multiprocessing import Process, Queue%0A
 import sys%0Af
@@ -543,16 +543,36 @@
 ion'%0A%5D%0A%0A
+NUM_PROCESSES = 8%0A%0A%0A
 class Co
@@ -1932,17 +1932,8 @@
 s):%0A
-        %0A
     
@@ -3599,16 +3599,47 @@
 artkey,%0A
+                               
         
@@ -3973,120 +3973,839 @@
 msg%0A
-        for doc in iter_docs(sourcedb, doc_ids):%0A            try:%0A                count += 1%0A                if 
+%0A        queue = Queue(150)%0A        for i in range(NUM_PROCESSES):%0A            Worker(queue, sourcedb, targetdb, exclude_types, total, simulate).start()%0A%0A        for doc in iter_docs(sourcedb, doc_ids, chunksize=100):%0A            count += 1%0A            queue.put((doc, count))%0A%0A        # shutdown workers%0A        for i in range(NUM_PROCESSES):%0A            queue.put(None)%0A%0A%0Aclass Worker(Process):%0A%0A    def __init__(self, queue, sourcedb, targetdb, exclude_types, total, simulate):%0A        super(Worker, self).__init__()%0A        self.queue = queue%0A        self.sourcedb = sourcedb%0A        self.targetdb = targetdb%0A        self.exclude_types = exclude_types%0A        self.total = total%0A        self.simulate = simulate%0A%0A    def run(self):%0A        for doc, count in iter(self.queue.get, None):%0A            try:%0A                if self.
 excl
@@ -4837,16 +4837,21 @@
 pe%22%5D in 
+self.
 exclude_
@@ -4997,24 +4997,29 @@
 e%22%5D, count, 
+self.
 total, doc%5B%22
@@ -5091,16 +5091,21 @@
  if not 
+self.
 simulate
@@ -5158,16 +5158,21 @@
 rm(doc, 
+self.
 sourcedb
@@ -5206,16 +5206,21 @@
 ave(dt, 
+self.
 targetdb
@@ -5291,16 +5291,21 @@
 (count, 
+self.
 total, d

add multithreading to copy_domain