from django.db import migrations

def update_rwi_codes(apps, schema_editor):
    NatureRWIAsset = apps.get_model('assets', 'NatureRWIAsset')
    
    # We try-catch the GuidedOperationData import because it's in another app,
    # and we want to ensure robust execution.
    try:
        GuidedOperationData = apps.get_model('interventions', 'GuidedOperationData')
    except LookupError:
        GuidedOperationData = None

    # Mappings from structure type to acronym
    mappings = {
        'avaloir': 'AV',
        'bordure sequencee': 'BS',
        'bouche injection': 'BI',
        'clapet anti-retour': 'CA',
        'empierrement sedimentation': 'ES',
        'kerb cell': 'KC',
        'puits infiltration': 'PI',
        'surverse': 'SU',
        'vanne isolation': 'VI',
        'vanne regulation': 'VR',
        'gargouille de trottoir': 'GT',
        'filet eau': 'FE',
        'avaloir de rail': 'AR',
        'bordure arasee': 'BA',
        'caniveau infiltrant': 'CI',
        'drain': 'D',
        'trop plein': 'TP',
        'reseau egout separatif': 'RES',
        'bassin sec/infiltration': 'OED',
        'bassin en eau': 'OED',
        'fosse': 'EVC',
        'jardin de pluie': 'EVC',
        'noue': 'EVC',
        'massif drainant': 'OED',
        'tranchee drainante': 'OED',
        'toiture stockante': 'OSI',
        'toiture verte': 'OSI',
        'bassin orage': 'OED',
        'SAUL-vide': 'OED',
        'SAUL-terre': 'TA',
        'citerne': 'OED',
        'espace vert sur dalle': 'EVC',
        'bande filtrante': 'EVC',
        'collecteur en beton poreux': 'OED',
        'chambre de visite': 'CV',
    }

    code_updates = {}

    for asset in NatureRWIAsset.objects.all():
        old_code = asset.code
        
        # Calculate acronym
        struct_type_lower = asset.structure_type.strip().lower() if asset.structure_type else None
        
        acronym = None
        if asset.code_name:
            acronym = asset.code_name.strip().upper()
        elif struct_type_lower in mappings:
            acronym = mappings[struct_type_lower]
        elif asset.geometry_type:
            acronym = asset.geometry_type.upper()
        else:
            acronym = 'UNK'
            
        fid = asset.external_id or asset.id
        new_code = f"RWI_{acronym}_{fid:03d}"
        
        if old_code != new_code:
            asset.code = new_code
            if asset.name_fr == old_code:
                asset.name_fr = new_code
            if asset.name_nl == old_code:
                asset.name_nl = new_code
            asset.save(update_fields=['code', 'name_fr', 'name_nl'])
            code_updates[old_code] = new_code

    # Update GuidedOperationData JSON references
    if GuidedOperationData and code_updates:
        for gd in GuidedOperationData.objects.all():
            modified = False
            data = gd.data
            if not isinstance(data, dict):
                continue
                
            def replace_codes(obj):
                nonlocal modified
                if isinstance(obj, dict):
                    for k, v in list(obj.items()):
                        if k == 'rwi_asset_code' and v in code_updates:
                            obj[k] = code_updates[v]
                            modified = True
                        else:
                            replace_codes(v)
                elif isinstance(obj, list):
                    for item in obj:
                        replace_codes(item)
                        
            replace_codes(data)
            if modified:
                gd.data = data
                gd.save(update_fields=['data'])

def reverse_rwi_codes(apps, schema_editor):
    NatureRWIAsset = apps.get_model('assets', 'NatureRWIAsset')
    try:
        GuidedOperationData = apps.get_model('interventions', 'GuidedOperationData')
    except LookupError:
        GuidedOperationData = None

    code_updates_reverse = {}

    for asset in NatureRWIAsset.objects.all():
        old_code = asset.code
        
        # Calculate old format: RWI_{geometry_type.upper()}_{external_id:03d}
        geom_type_upper = asset.geometry_type.upper() if asset.geometry_type else 'POLYGON'
        fid = asset.external_id or asset.id
        original_code = f"RWI_{geom_type_upper}_{fid:03d}"
        
        if old_code != original_code:
            asset.code = original_code
            if asset.name_fr == old_code:
                asset.name_fr = original_code
            if asset.name_nl == old_code:
                asset.name_nl = original_code
            asset.save(update_fields=['code', 'name_fr', 'name_nl'])
            code_updates_reverse[old_code] = original_code

    # Revert GuidedOperationData JSON references
    if GuidedOperationData and code_updates_reverse:
        for gd in GuidedOperationData.objects.all():
            modified = False
            data = gd.data
            if not isinstance(data, dict):
                continue
                
            def replace_codes(obj):
                nonlocal modified
                if isinstance(obj, dict):
                    for k, v in list(obj.items()):
                        if k == 'rwi_asset_code' and v in code_updates_reverse:
                            obj[k] = code_updates_reverse[v]
                            modified = True
                        else:
                            replace_codes(v)
                elif isinstance(obj, list):
                    for item in obj:
                        replace_codes(item)
                        
            replace_codes(data)
            if modified:
                gd.data = data
                gd.save(update_fields=['data'])


class Migration(migrations.Migration):

    dependencies = [
        ('assets', '0012_trafficlightintersection_is_radar'),
    ]

    operations = [
        migrations.RunPython(update_rwi_codes, reverse_rwi_codes),
    ]