# Generated by Django 5.2 on 2025-10-04 14:58

from django.db import migrations


VISION_DEFAULT_PROVIDER_KEYS = ("openai", "anthropic", "google")


def forward_set_vision(apps, schema_editor):
    """Enable vision support on default Anthropic and OpenAI endpoints."""

    Provider = apps.get_model("api", "LLMProvider")
    PersistentEndpoint = apps.get_model("api", "PersistentModelEndpoint")
    BrowserEndpoint = apps.get_model("api", "BrowserModelEndpoint")

    provider_ids = list(
        Provider.objects.filter(key__in=VISION_DEFAULT_PROVIDER_KEYS).values_list("id", flat=True)
    )
    if not provider_ids:
        return

    PersistentEndpoint.objects.filter(provider_id__in=provider_ids).update(supports_vision=True)
    BrowserEndpoint.objects.filter(provider_id__in=provider_ids).update(supports_vision=True)


def reverse_clear_vision(apps, schema_editor):
    """Reset vision support to False on affected endpoints."""

    Provider = apps.get_model("api", "LLMProvider")
    PersistentEndpoint = apps.get_model("api", "PersistentModelEndpoint")
    BrowserEndpoint = apps.get_model("api", "BrowserModelEndpoint")

    provider_ids = list(
        Provider.objects.filter(key__in=VISION_DEFAULT_PROVIDER_KEYS).values_list("id", flat=True)
    )
    if not provider_ids:
        return

    PersistentEndpoint.objects.filter(provider_id__in=provider_ids).update(supports_vision=False)
    BrowserEndpoint.objects.filter(provider_id__in=provider_ids).update(supports_vision=False)


class Migration(migrations.Migration):

    dependencies = [
        ('api', '0153_browsermodelendpoint_supports_vision_and_more'),
    ]

    operations = [
        migrations.RunPython(forward_set_vision, reverse_clear_vision),
    ]
