from dimagi.utils.couch.database import get_db, iter_docs
from django.core.management import BaseCommand
from corehq.apps.domain.models import Domain

def add_to_case(case):
    forms = {}
    def get_xform(xid):
        xform = forms.get(xid)
        if not xform:
            xform = get_db().get(xid)
            forms[xid] = xform
        return xform

    should_save = False

    def get_user_id(form_id):
        if form_id:
            user_id = get_xform(form_id).get('form', {}).get("meta", {}).get("userID") or \
                get_xform(form_id).get('form', {}).get("case", {}).get("@user_id")
            if not user_id:
                print "Could not find a user_id for form %s" % form_id
            return user_id
            
        print "Case %s has actions without a proper xform_id" % case.get('_id')

    for action in case.get('actions', []):
        if not action.get('user_id'):
            user_id = get_user_id(action.get('xform_id'))
            if user_id:
                action['user_id'] = user_id
                should_save = True

        if not case.get('opened_by') and action.get('action_type') == 'create':
            opened_by = get_user_id(action.get('xform_id'))
            if opened_by:
                case['opened_by'] = opened_by
                should_save = True

        if case.get('closed', False) and not case.get('closed_by') and action.get('action_type') == 'close':
            closed_by = get_user_id(action.get('xform_id'))
            if closed_by:
                case['closed_by'] = closed_by
                should_save = True

    return should_save

class Command(BaseCommand):
    args = '<domain_name domain_name ...>'
    help = """
        Loops through every case and adds the opened_by and closed_by property to each case where applicable.
        It also adds user_id to each case action in order to fix a bug.
    """

    def handle(self, *args, **options):
        domain_names = args or [d["key"] for d in Domain.get_all(include_docs=False)]
        for d in domain_names:
            print "Migrating cases in project space: %s" % d
            key = ["all", d]
            results = get_db().view('case/all_cases',
                startkey=key,
                endkey=key + [{}],
                reduce=False,
                include_docs=False,
            ).all()

            cases_to_save = []
            for raw_case in iter_docs(get_db(), [r['id'] for r in results]):
                if add_to_case(raw_case):
                    cases_to_save.append(raw_case)

                if len(cases_to_save) > 100:
                    get_db().bulk_save(cases_to_save)
                    cases_to_save = []
            get_db().bulk_save(cases_to_save)

@@ -1,12 +1,85 @@
+from optparse import make_option%0Afrom couchdbkit import ResourceNotFound%0A
 from dimagi.
@@ -332,16 +332,37 @@
  xform:%0A
+            try:%0A    
         
@@ -383,32 +383,160 @@
 t_db().get(xid)%0A
+            except ResourceNotFound:%0A                print %22Resource not found for xform: %25s%22 %25 xid%0A                return None%0A
             form
@@ -655,31 +655,28 @@
             
-user_id
+form
  = get_xform
@@ -676,32 +676,112 @@
 t_xform(form_id)
+%0A            if not form:%0A                return None%0A            user_id = form
 .get('form', %7B%7D)
@@ -832,34 +832,20 @@
         
-get_xform(form_id)
+form
 .get('fo
@@ -2182,129 +2182,479 @@
     
-def handle(self, *args, **options):%0A        domain_names = args or %5Bd%5B%22key%22%5D for d in Domain.get_all(include_docs=False)%5D
+option_list = BaseCommand.option_list + (%0A        make_option('-s', '--startat',%0A            help=%22Begin migration at this domain. (Will ignore if domain isn't found)%22),%0A        )%0A%0A    def handle(self, *args, **options):%0A        domain_names = args or %5Bd%5B%22key%22%5D for d in Domain.get_all(include_docs=False)%5D%0A%0A        start_at = options.get(%22startat%22)%0A        if start_at and start_at in domain_names:%0A            domain_names = domain_names%5Bdomain_names.index(start_at):%5D%0A
 %0A   

add option for starting at a specific domain