from django.core.management.base import BaseCommand
from couchforms.models import XFormInstance
from dimagi.utils.couch.database import iter_docs
from corehq.apps.commtrack.models import SupplyPointCase


class Command(BaseCommand):
    startkey = ['ipm-senegal', 'by_type', 'XFormInstance']
    endkey = startkey + [{}]

    ids = [row['id'] for row in XFormInstance.get_db().view(
        "couchforms/all_submissions_by_domain",
        startkey=startkey,
        endkey=endkey,
        reduce=False
    )]

    to_save = []

    for doc in iter_docs(XFormInstance.get_db(), ids):
        if 'location_id' in doc['form'] and not doc['form']['location_id']:
            case = SupplyPointCase.get(doc['form']['case']['@case_id'])
            if case.type == 'supply-point':
                print 'updating'
                print 'case location_id:', case.location_id
                print 'from:', doc['form']['location_id']
                doc['form']['location_id'] = case.location_id
                print 'to:', doc['form']['location_id']
                to_save.append(doc)

        if len(to_save) > 500:
            XFormInstance.get_db().bulk_save(to_save)
            to_save = []

    if to_save:
        XFormInstance.get_db().bulk_save(to_save)

@@ -139,16 +139,78 @@
 er_docs%0A
+from corehq.util.couch_helpers import CouchAttachmentsBuilder%0A
 from cor
@@ -258,16 +258,49 @@
 intCase%0A
+from lxml import etree%0Aimport re%0A
 %0A%0Aclass 
@@ -671,16 +671,33 @@
 , ids):%0A
+        try:%0A    
         
@@ -776,16 +776,20 @@
         
+    
 case = S
@@ -852,16 +852,20 @@
         
+ 
+   
 if case.
@@ -908,53 +908,133 @@
     
-print 'updating'%0A                print 'case 
+    instance = XFormInstance.get(doc%5B'_id'%5D)%0A%0A                    # fix the XFormInstance%0A                    instance.form%5B'
 loca
@@ -1040,19 +1040,20 @@
 ation_id
-:',
+'%5D =
  case.lo
@@ -1054,32 +1054,37 @@
 ase.location_id%0A
+%0A    
                 
@@ -1087,49 +1087,205 @@
     
-print 'from:', doc%5B'form'%5D%5B'location_id'%5D
+# fix the actual form.xml%0A                    xml_object = etree.fromstring(instance.get_xml())%0A                    location_id_node = xml_object.find(re.sub('%7D.*', '%7Dlocation_id', xml_object.tag))
 %0A   
@@ -1297,29 +1297,20 @@
         
-doc%5B'form'%5D%5B'
+    
 location
@@ -1304,34 +1304,42 @@
      location_id
-'%5D
+_node.text
  = case.location
@@ -1362,42 +1362,513 @@
     
-print 'to:', doc%5B'form'%5D%5B'location
+    updated_xml = etree.tostring(xml_object)%0A%0A                    attachment_builder = CouchAttachmentsBuilder(instance._attachments)%0A                    attachment_builder.add(%0A                        name='form.xml',%0A                        content=updated_xml,%0A                        content_type=instance._attachments%5B'form.xml'%5D%5B'content_type'%5D%0A                    )%0A                    instance._attachments = attachment_builder.to_json()%0A%0A                    print 'Updating XFormInstance:', doc%5B'
 _id'
@@ -1881,24 +1881,28 @@
             
+   
+ 
 to_save.appe
@@ -1908,12 +1908,95 @@
 end(
-doc)
+instance)%0A        except:%0A            print 'Failed to save XFormInstance:', doc%5B'_id'%5D
 %0A%0A  

Make management command update xml too