# Generated by Django 3.2.25 on 2024-09-12 21:59

from django.db import migrations, models, transaction
import django.db.models.deletion
import django_migration_linter as linter
from core.redis import start_job_async_or_sync
from ml_models.models import ThirdPartyModelVersion
from ml_model_providers.models import ModelProviderConnection, ModelProviders


def _fill_model_version_model_provider_connection(db_alias: str):
    for provider in [ModelProviders.OPENAI, ModelProviders.AZURE_OPENAI]:
        this_provider_model_versions = (
            ThirdPartyModelVersion.objects.using(db_alias)
            .filter(provider=provider)
            .values('id', 'organization_id', 'provider_model_id')
        )
        for provider_model_version in this_provider_model_versions:
            connection_ids = ModelProviderConnection.objects.using(db_alias).filter(
                organization_id=provider_model_version['organization_id'],
                provider=provider,
                **({'deployment_name': provider_model_version['provider_model_id']} if provider == ModelProviders.AZURE_OPENAI else {}),
            ).values_list('id', flat=True)[:1]
            connection_id = connection_ids[0] if connection_ids else None
            ThirdPartyModelVersion.objects.using(db_alias).filter(id=provider_model_version['id']).update(model_provider_connection_id=connection_id)

def forwards(apps, schema_editor):
    db_alias = schema_editor.connection.alias
    start_job_async_or_sync(_fill_model_version_model_provider_connection, db_alias=db_alias)


def backwards(apps, schema_editor):
    pass


class Migration(migrations.Migration):
    atomic = False

    dependencies = [
        ('ml_model_providers', '0003_modelproviderconnection_cached_available_models'),
        ('ml_models', '0010_modelinterface_skill_name'),
    ]

    operations = [
        linter.IgnoreMigration(),
        migrations.AddField(
            model_name='thirdpartymodelversion',
            name='model_provider_connection',
            field=models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='model_versions', to='ml_model_providers.modelproviderconnection'),
        ),
        migrations.RunPython(forwards, backwards)
    ]