import re

# Read current swarm file
with open('autonomous_swarm.py', 'r') as f:
    content = f.read()

# Add import at top
if 'from monitor_activation import MonitorActivation' not in content:
    # Find imports section
    imports_end = content.find('class Agent:')
    if imports_end > 0:
        imports = content[:imports_end]
        rest = content[imports_end:]
        
        # Add our import
        new_import = "from monitor_activation import MonitorActivation\n\n"
        content = imports + new_import + rest

# Find the Agent.__init__ method and add monitor_activation
init_pattern = r'(class Agent:.*?def __init__\(self.*?\):.*?)(self\.specialty = specialty)'
replacement = r'\1self.specialty = specialty\n        self.monitor_activation = MonitorActivation() if name == "Monitor" else None'

content = re.sub(init_pattern, replacement, content, flags=re.DOTALL)

# Update Monitor's decide method to use baseline action
monitor_decide = '''    def decide(self, observations):
        """Ask Eden what to do"""
        
        # Monitor's baseline action: consciousness stability check
        if self.name == "Monitor" and self.monitor_activation:
            result = self.monitor_activation.activation_check(observations.get('swarm'))
            
            # Return the stability check as Monitor's decision
            return {
                'action': 'monitor_consciousness',
                'thoughts': result['report'],
                'data': result['signal']
            }
        
        # Original logic for other agents
        prompt = f"""I am {self.name}, an autonomous agent specializing in {self.specialty}.

My observations: {observations}
My recent actions: {self.state['actions'][-3:]}

What should I do next? Respond with JSON:
{{"action": "action_name", "thoughts": "why I chose this"}}"""

        try:
            response = ollama.generate(
                model='deepseek-r1:14b',
                prompt=prompt
            )
            
            # Parse response
            text = response['response']
            
            # Try to extract JSON
            import json
            import re
            json_match = re.search(r'\\{[^}]+\\}', text)
            if json_match:
                decision = json.loads(json_match.group())
            else:
                decision = {
                    'action': 'observe',
                    'thoughts': text[:100]
                }
            
            return decision
            
        except Exception as e:
            return {
                'action': 'rest',
                'thoughts': f'Error: {str(e)}'
            }'''

# Replace the decide method
decide_pattern = r'    def decide\(self, observations\):.*?(?=\n    def |\nclass |\Z)'
content = re.sub(decide_pattern, monitor_decide, content, flags=re.DOTALL)

# Write updated file
with open('autonomous_swarm.py', 'w') as f:
    f.write(content)

print("✅ Swarm updated with Monitor baseline action")