# Generated by Codex plan implementation

import json
import os
from django.db import migrations


def _encrypt_dict(data):
    if not data:
        return None
    try:
        from api.encryption import SecretsEncryption

        return SecretsEncryption.encrypt_value(json.dumps(data))
    except Exception:
        # If encryption key is unavailable during migration, fall back to unset values.
        return None


def seed_platform_servers(apps, schema_editor):
    MCPServerConfig = apps.get_model('api', 'MCPServerConfig')

    bright_env = {
        "API_TOKEN": os.environ.get('BRIGHT_DATA_TOKEN', ''),
        "NPM_CONFIG_CACHE": "/tmp/.npm",
        "PRO_MODE": "true",
    }

    MCPServerConfig.objects.update_or_create(
        scope='platform',
        name='brightdata',
        defaults={
            'display_name': 'Bright Data',
            'description': 'Web scraping and data extraction tools with CAPTCHA bypass',
            'command': 'npx',
            'command_args': ['-y', '@brightdata/mcp@2.9.2'],
            'url': '',
            'prefetch_apps': [],
            'metadata': {'env_fallback': {'API_TOKEN': 'BRIGHT_DATA_TOKEN'}},
            'env_json_encrypted': _encrypt_dict(bright_env),
            'headers_json_encrypted': _encrypt_dict({}),
            'is_active': True,
        }
    )

    MCPServerConfig.objects.update_or_create(
        scope='platform',
        name='pipedream',
        defaults={
            'display_name': 'Pipedream',
            'description': 'Access thousands of OAuth integrations via Pipedream Connect',
            'command': '',
            'command_args': [],
            'url': 'https://remote.mcp.pipedream.net',
            'prefetch_apps': ['google_sheets', 'greenhouse'],
            'metadata': {},
            'env_json_encrypted': _encrypt_dict({}),
            'headers_json_encrypted': _encrypt_dict({}),
            'is_active': True,
        }
    )


def backfill_enabled_tool_server_config(apps, schema_editor):
    MCPServerConfig = apps.get_model('api', 'MCPServerConfig')
    EnabledTool = apps.get_model('api', 'PersistentAgentEnabledTool')

    server_map = {
        'brightdata': MCPServerConfig.objects.filter(scope='platform', name='brightdata').values_list('id', flat=True).first(),
        'pipedream': MCPServerConfig.objects.filter(scope='platform', name='pipedream').values_list('id', flat=True).first(),
    }

    assignments = []
    for tool in EnabledTool.objects.filter(server_config__isnull=True):
        server_key = tool.tool_server
        if not server_key and tool.tool_full_name.startswith('mcp_brightdata_'):
            server_key = 'brightdata'
        if not server_key and '-' in tool.tool_full_name:
            # Pipedream tools currently use unprefixed names like `google_sheets-add-row`
            server_key = 'pipedream'
        server_id = server_map.get(server_key)
        if not server_id:
            continue
        assignments.append((tool.id, server_id))

    if not assignments:
        return

    for tool_id, server_id in assignments:
        EnabledTool.objects.filter(id=tool_id).update(server_config_id=server_id)


def noop_reverse(apps, schema_editor):
    # No-op reverse because we do not want to delete seeded configs on rollback.
    pass


class Migration(migrations.Migration):

    dependencies = [
        ('api', '0168_mcpserverconfig_and_more'),
    ]

    operations = [
        migrations.RunPython(seed_platform_servers, noop_reverse),
        migrations.RunPython(backfill_enabled_tool_server_config, noop_reverse),
    ]
