# Generated by Django migration for secrets refactor

import logging
import json
import os
from django.db import migrations, models
import django.db.models.deletion
import uuid

logger = logging.getLogger(__name__)


def migrate_secrets_to_table(apps, schema_editor):
    """
    Migrate existing JSON-based secrets to the new PersistentAgentSecret table.
    """
    PersistentAgent = apps.get_model("api", "PersistentAgent")
    PersistentAgentSecret = apps.get_model("api", "PersistentAgentSecret")
    
    # We need to decrypt the existing secrets and re-encrypt them as individual values
    # This requires access to the encryption functions
    
    total_agents = 0
    total_secrets_migrated = 0
    failed_agents = []
    
    for agent in PersistentAgent.objects.filter(encrypted_secrets__isnull=False):
        total_agents += 1
        
        try:
            # Import encryption functions for migration
            from api.encryption import SecretsEncryption
            
            # Decrypt existing secrets
            decrypted_secrets = SecretsEncryption.decrypt_secrets(agent.encrypted_secrets)
            if not decrypted_secrets:
                continue
            
            # Handle both legacy flat format and domain-specific format
            if SecretsEncryption._is_domain_specific_format(decrypted_secrets):
                # Domain-specific format: {"https://example.com": {"key": "value"}}
                secrets_by_domain = decrypted_secrets
            else:
                # Legacy flat format: {"key": "value"}
                # Convert to domain-specific format with default domain
                default_domain = "https://*.legacy-migrated.local"
                secrets_by_domain = {default_domain: decrypted_secrets}
            
            # Create individual secret records
            agent_secrets_count = 0
            for domain_pattern, secrets in secrets_by_domain.items():
                for key, value in secrets.items():
                    # Encrypt the individual value
                    encrypted_value = SecretsEncryption.encrypt_value(value)
                    
                    # Create the secret record
                    PersistentAgentSecret.objects.create(
                        id=uuid.uuid4(),
                        agent=agent,
                        domain_pattern=domain_pattern,
                        key=key,
                        encrypted_value=encrypted_value
                    )
                    agent_secrets_count += 1
            
            total_secrets_migrated += agent_secrets_count
            logger.info(f"Migrated {agent_secrets_count} secrets for agent {agent.id}")
            
        except Exception as e:
            logger.error(f"Failed to migrate secrets for agent {agent.id}: {str(e)}")
            failed_agents.append(str(agent.id))
    
    logger.info(f"Migration complete: {total_agents} agents processed, {total_secrets_migrated} secrets migrated")
    if failed_agents:
        logger.warning(f"Failed to migrate {len(failed_agents)} agents: {failed_agents}")


def reverse_migrate_secrets_from_table(apps, schema_editor):
    """
    Reverse migration: convert PersistentAgentSecret records back to JSON format.
    This is for rollback purposes only.
    """
    PersistentAgent = apps.get_model("api", "PersistentAgent")
    PersistentAgentSecret = apps.get_model("api", "PersistentAgentSecret")
    
    for agent in PersistentAgent.objects.all():
        secrets_by_domain = {}
        
        # Get all secrets for this agent
        for secret in PersistentAgentSecret.objects.filter(agent=agent):
            try:
                from api.encryption import SecretsEncryption
                
                # Decrypt the value
                decrypted_value = SecretsEncryption.decrypt_value(secret.encrypted_value)
                
                # Group by domain
                if secret.domain_pattern not in secrets_by_domain:
                    secrets_by_domain[secret.domain_pattern] = {}
                
                secrets_by_domain[secret.domain_pattern][secret.key] = decrypted_value
                
            except Exception as e:
                logger.error(f"Failed to decrypt secret {secret.id} for agent {agent.id}: {str(e)}")
        
        if secrets_by_domain:
            try:
                from api.encryption import SecretsEncryption
                
                # Re-encrypt as JSON format
                agent.encrypted_secrets = SecretsEncryption.encrypt_secrets(secrets_by_domain, allow_legacy=True)
                agent.secret_keys = SecretsEncryption.get_secret_keys_for_audit(secrets_by_domain)
                agent.save(update_fields=['encrypted_secrets', 'secret_keys'])
                
            except Exception as e:
                logger.error(f"Failed to restore JSON secrets for agent {agent.id}: {str(e)}")


class Migration(migrations.Migration):

    dependencies = [
        ('api', '0066_uniq_free_plan_block_per_month'),
    ]

    operations = [
        # Create the new PersistentAgentSecret model
        migrations.CreateModel(
            name='PersistentAgentSecret',
            fields=[
                ('id', models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False)),
                ('domain_pattern', models.CharField(help_text="Domain pattern where this secret can be used (e.g., 'https://example.com', '*.google.com')", max_length=256)),
                ('key', models.CharField(help_text='Secret key name (alphanumeric with underscores only)', max_length=64)),
                ('encrypted_value', models.BinaryField(help_text='AES-256-GCM encrypted secret value')),
                ('created_at', models.DateTimeField(auto_now_add=True)),
                ('updated_at', models.DateTimeField(auto_now=True)),
                ('agent', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='secrets', to='api.persistentagent')),
            ],
            options={
                'ordering': ['domain_pattern', 'key'],
            },
        ),
        
        # Add constraints and indexes
        migrations.AddConstraint(
            model_name='persistentagentsecret',
            constraint=models.UniqueConstraint(fields=('agent', 'domain_pattern', 'key'), name='unique_agent_domain_secret_key'),
        ),
        migrations.AddIndex(
            model_name='persistentagentsecret',
            index=models.Index(fields=['agent', 'domain_pattern'], name='pa_secret_agent_domain_idx'),
        ),
        migrations.AddIndex(
            model_name='persistentagentsecret',
            index=models.Index(fields=['agent'], name='pa_secret_agent_idx'),
        ),
        
        # Migrate existing secrets from JSON to table
        migrations.RunPython(
            migrate_secrets_to_table,
            reverse_migrate_secrets_from_table,
        ),
        
        # Remove old fields from PersistentAgent
        migrations.RemoveField(
            model_name='persistentagent',
            name='encrypted_secrets',
        ),
        migrations.RemoveField(
            model_name='persistentagent',
            name='secret_keys',
        ),
    ] 