import datetime
import logging

from corehq.apps.es import CaseSearchES
from django.core.management import BaseCommand
from casexml.apps.case.mock import CaseStructure, CaseFactory

logger = logging.getLogger('enikshay_2b_reason_for_test_fix')


class Command(BaseCommand):

    def add_arguments(self, parser):
        parser.add_argument(
            'domain',
            help="The domain to migrate."
        )
        parser.add_argument(
            "log_path",
            help="Path to write the log to"
        )
        parser.add_argument(
            '--commit',
            action='store_true',
            help="actually modifies the cases. Without this flag, it's a dry run."
        )

    def handle(self, domain, log_path, **options):
        commit = options['commit']
        factory = CaseFactory(domain)

        logger.info("Starting {} migration on {} at {}".format(
            "real" if commit else "fake", domain, datetime.datetime.utcnow()
        ))

        cases = (CaseSearchES()
                 .domain(domain)
                 .case_type("episode")
                 .case_property_query("updated_by_migration", "enikshay_2b_case_properties", "must")
                 .run().hits)

        with open(log_path, "w") as f:
            for case in cases:
                case_props = {prop['key']: prop['value'] for prop in case['case_properties']}
                treatment_status = None
                treatment_initiated = case_props.get('treatment_initiated')
                diagnosing_facility_id = case_props.get('diagnosing_facility_id')
                treatment_initiating_facility_id = case_props.get('treatment_initiating_facility_id')

                if diagnosing_facility_id and treatment_initiating_facility_id:
                    if treatment_initiated == 'yes_phi' and \
                                    diagnosing_facility_id != treatment_initiating_facility_id:
                        treatment_status = 'initiated_outside_facility'
                    elif treatment_initiated == 'yes_phi':
                        treatment_status = 'initiated_first_line_treatment'
                    elif treatment_initiated == 'yes_private':
                        treatment_status = 'initiated_outside_rntcp'

                    if treatment_status:
                        case_id = case['_id']
                        f.write(case_id + "\n")
                        logger.info(case_id)

                        case_structure = CaseStructure(
                            case_id=case_id,
                            walk_related=False,
                            attrs={
                                "create": False,
                                "update": {
                                    "treatment_status": treatment_status,
                                    "updated_by_migration": "enikshay_2b_treatment_status_fix",
                                },
                            },
                        )

                        if commit:
                            factory.create_or_update_case(case_structure)
        logger.info("Migration finished at {}".format(datetime.datetime.utcnow()))

@@ -1697,91 +1697,97 @@
  if 
-diagnosing_facility_id and treatment_initiating_facility_id:%0A                    if
+treatment_initiated == 'yes_phi' and %5C%0A                        diagnosing_facility_id and
  tre
@@ -1796,39 +1796,39 @@
 ment_initiat
-ed == 'yes_phi'
+ing_facility_id
  and %5C%0A     
@@ -1842,36 +1842,24 @@
             
-  
-          
 diagnosing_f
@@ -1918,36 +1918,32 @@
                 
-    
 treatment_status
@@ -1974,20 +1974,16 @@
 cility'%0A
-    
         
@@ -2027,22 +2027,108 @@
 yes_phi'
-:
+ and %5C
 %0A    
+                    diagnosing_facility_id and treatment_initiating_facility_id:%0A
         
@@ -2199,36 +2199,32 @@
                 
-   
- 
 elif treatment_i
@@ -2250,20 +2250,16 @@
 ivate':%0A
-    
         
@@ -2324,36 +2324,32 @@
                 
-    
 if treatment_sta
@@ -2369,28 +2369,24 @@
             
- 
-   
 case_id = ca
@@ -2395,20 +2395,16 @@
 %5B'_id'%5D%0A
-    
         
@@ -2455,28 +2455,24 @@
             
-  
-  
 logger.info(
@@ -2481,20 +2481,16 @@
 se_id)%0A%0A
-    
         
@@ -2553,28 +2553,24 @@
             
-   
- 
 case_id=case
@@ -2574,20 +2574,16 @@
 ase_id,%0A
-    
         
@@ -2614,28 +2614,24 @@
 ated=False,%0A
-    
             
@@ -2678,20 +2678,16 @@
         
-    
 %22create%22
@@ -2719,28 +2719,24 @@
             
-    
 %22update%22: %7B%0A
@@ -2767,20 +2767,16 @@
         
- 
-   
 %22treatme
@@ -2805,20 +2805,16 @@
 status,%0A
-    
         
@@ -2921,35 +2921,27 @@
             
-  
-  
 %7D,%0A
-    
             
@@ -2975,27 +2975,19 @@
         
-   
- 
 )%0A%0A
-    
         
@@ -3009,20 +3009,16 @@
 commit:%0A
-    
         

Fix treatment initiation bug