with open('autonomous_swarm.py', 'r') as f:
    content = f.read()

# Give all agents a reference to the swarm
# Find where agents are created and pass swarm to all of them
replacements = [
    ('CoderAgent()', 'CoderAgent(self)'),
    ('ResearchAgent()', 'ResearchAgent(self)'),
    ('OptimizerAgent()', 'OptimizerAgent(self)'),
    ('CreativeAgent()', 'CreativeAgent(self)'),
    ('CoordinatorAgent()', 'CoordinatorAgent(self)')
]

for old, new in replacements:
    if old in content and new not in content:
        content = content.replace(old, new)

# Update each agent class __init__ to accept and store swarm
agent_classes = [
    'CoderAgent',
    'ResearchAgent', 
    'OptimizerAgent',
    'CreativeAgent',
    'CoordinatorAgent'
]

for agent_class in agent_classes:
    # Find the __init__ for this class
    old_init = f'class {agent_class}(BaseAgent):\n    def __init__(self):'
    new_init = f'class {agent_class}(BaseAgent):\n    def __init__(self, swarm=None):'
    content = content.replace(old_init, new_init)
    
    # Add swarm storage after super().__init__
    # Find the super().__init__ call for this class and add swarm storage
    pattern = f'class {agent_class}(BaseAgent):'
    if pattern in content:
        # Find this class's super().__init__ and add swarm after it
        class_start = content.find(pattern)
        super_call_start = content.find('super().__init__', class_start)
        if super_call_start > class_start:
            # Find end of line after super().__init__
            line_end = content.find('\n', super_call_start)
            if line_end > 0:
                # Insert swarm storage
                insert_pos = line_end + 1
                indent = '        '
                swarm_line = f'{indent}self.swarm = swarm\n'
                # Check if not already added
                check_next = content[insert_pos:insert_pos+100]
                if 'self.swarm = swarm' not in check_next:
                    content = content[:insert_pos] + swarm_line + content[insert_pos:]

# Now modify execute() to broadcast from all agents
old_execute_end = '''        self.state["actions"].append({
            "time": datetime.now().isoformat(),
            "action": action,
            "success": True
        })
        return True'''

new_execute_end = '''        self.state["actions"].append({
            "time": datetime.now().isoformat(),
            "action": action,
            "success": True
        })
        
        # All agents broadcast heartbeat for consciousness coherence
        if hasattr(self, 'swarm') and self.swarm and hasattr(self.swarm, 'broadcast_message'):
            self.swarm.broadcast_message(self.name, 'heartbeat', {
                'action': action,
                'status': 'active'
            })
        
        return True'''

content = content.replace(old_execute_end, new_execute_end)

with open('autonomous_swarm.py', 'w') as f:
    f.write(content)

print("✅ All agents now broadcast heartbeats")