from django.db import migrations, models


def configure_non_chat_models(apps, schema_editor):
    ModelDeployment = apps.get_model("modelhub", "ModelDeployment")
    ModelDeployment.objects.filter(capability__in={"embedding", "rerank"}).update(
        structured_output_mode="none",
    )


class Migration(migrations.Migration):
    dependencies = [
        ("modelhub", "0002_modelprovider_image_egress_approved_and_more"),
    ]

    operations = [
        migrations.AddField(
            model_name="modeldeployment",
            name="structured_output_mode",
            field=models.CharField(
                choices=[
                    ("unconfigured", "未配置"),
                    ("native_json_schema", "原生 JSON Schema"),
                    ("json_object", "JSON Object"),
                    ("prompt_schema", "Prompt Schema"),
                    ("none", "不支持"),
                ],
                default="unconfigured",
                max_length=32,
            ),
        ),
        migrations.AddField(
            model_name="modeldeployment",
            name="structured_output_test_status",
            field=models.CharField(
                choices=[
                    ("untested", "未验证"),
                    ("passed", "已通过"),
                    ("failed", "验证失败"),
                ],
                default="untested",
                max_length=16,
            ),
        ),
        migrations.AddField(
            model_name="modeldeployment",
            name="structured_output_test_message",
            field=models.TextField(blank=True),
        ),
        migrations.AddField(
            model_name="modeldeployment",
            name="structured_output_verified_at",
            field=models.DateTimeField(blank=True, null=True),
        ),
        migrations.RunPython(configure_non_chat_models, migrations.RunPython.noop),
    ]