import os
import sys
import django
import re

# Set up Django
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings")
django.setup()

with open('core/models.py', 'r') as f:
    content = f.read()

# 1. New pre_save signal if not exists
new_pre_save = """
@receiver(pre_save, sender=Voter)
def handle_voter_status_on_voted_pre_save(sender, instance, **kwargs):
    \"\"\"
    If a voter has voted, ensure they are not targets for door visits or calls.
    \"\"\"
    if instance.voted:
        instance.target_door_visit = False
        instance.call_queue_status = 'no_call_required'
"""

if '@receiver(pre_save, sender=Voter)' not in content:
    # This should have been added by previous run, but let's be safe
    # If it was added, it's already there.
    pass

# 2. Update update_voter_call_queue_status_on_voter_save
new_voter_signal = """def update_voter_call_queue_status_on_voter_save(sender, instance, **kwargs):
    \"\"\"
    Sync call_queue_status when is_targeted, candidate_support or voted changes.
    \"\"\"
    if getattr(instance, '_skip_signals', False):
        return

    orig = getattr(instance, '_orig_obj', None)
    if orig and instance.call_queue_status != orig.call_queue_status:
        # If call_queue_status was manually changed, don't auto-override in this save
        return

    update_fields = kwargs.get('update_fields')
    if update_fields:
        relevant = {'is_targeted', 'candidate_support', 'voted'}
        if not relevant.intersection(update_fields):
            return

    # PRIORITY 1: If they voted, no call required and cancel pending calls
    if instance.voted:
        # Cancel any pending calls
        ScheduledCall.objects.filter(voter=instance, status='pending').update(status='cancelled')
        
        if instance.call_queue_status != 'no_call_required':
            Voter.objects.filter(pk=instance.pk).update(call_queue_status='no_call_required')
        return

    # PRIORITY 2: Check if in queue (pending scheduled call)
    if ScheduledCall.objects.filter(voter=instance, status='pending').exists():
        if instance.call_queue_status != 'in_call_queue':
            Voter.objects.filter(pk=instance.pk).update(call_queue_status='in_call_queue')
        return

    # PRIORITY 3: If support is 'supporting', then 'no_call_required'
    if instance.candidate_support == 'supporting':
        if instance.call_queue_status != 'no_call_required':
            Voter.objects.filter(pk=instance.pk).update(call_queue_status='no_call_required')
        return

    # PRIORITY 4: If un-targeted, set to no_call_required
    if not instance.is_targeted:
        if instance.call_queue_status != 'no_call_required':
            Voter.objects.filter(pk=instance.pk).update(call_queue_status='no_call_required')
    else:
        # If targeted, and currently no_call_required, set to to_be_called
        if instance.call_queue_status == 'no_call_required':
             Voter.objects.filter(pk=instance.pk).update(call_queue_status='to_be_called')"""

# Match the old function precisely based on what we saw in read_file
# We'll use the line numbers or the function signature and the next function
content = re.sub(r'def update_voter_call_queue_status_on_voter_save\(sender, instance, \*\*kwargs\):.*?(@receiver\(post_save, sender=ScheduledCall\))', new_voter_signal + '\n\n\\1', content, flags=re.DOTALL)

# 3. Update update_voter_call_queue_status_on_call_save
new_call_signal = """def update_voter_call_queue_status_on_call_save(sender, instance, **kwargs):
    \"\"\"
    Sync Voter.call_queue_status when a ScheduledCall is saved.
    \"\"\"
    voter = instance.voter
    
    # PRIORITY 0: If they voted, always no_call_required
    if voter.voted:
        if voter.call_queue_status != 'no_call_required':
            voter.call_queue_status = 'no_call_required'
            voter.save(update_fields=['call_queue_status'])
        return"""

# Match old call save
content = re.sub(r'def update_voter_call_queue_status_on_call_save\(sender, instance, \*\*kwargs\):.*?voter = instance.voter', new_call_signal, content, flags=re.DOTALL)

with open('core/models.py', 'w') as f:
    f.write(content)

print("Updated core/models.py with more robust signals (v4).")