# Generated by Django 5.2 on 2025-10-08 16:49

from django.conf import settings
from django.db import migrations, models

from constants.plans import PlanNamesChoices


class Migration(migrations.Migration):

    dependencies = [
        ('api', '0154_set_vision_defaults'),
    ]

    def _set_default_daily_limits(apps, schema_editor):  # pragma: no cover - data migration
        PersistentAgent = apps.get_model('api', 'PersistentAgent')
        UserBilling = apps.get_model('api', 'UserBilling')
        OrganizationBilling = apps.get_model('api', 'OrganizationBilling')

        free_value = PlanNamesChoices.FREE.value

        default_limit = getattr(settings, "DEFAULT_AGENT_DAILY_CREDIT_LIMIT", 10)

        for agent in PersistentAgent.objects.filter(daily_credit_limit__isnull=True):
            plan_value = free_value

            if agent.organization_id:
                billing = OrganizationBilling.objects.filter(organization_id=agent.organization_id).only('subscription').first()
                if billing and billing.subscription:
                    plan_value = billing.subscription
            else:
                billing = UserBilling.objects.filter(user_id=agent.user_id).only('subscription').first()
                if billing and billing.subscription:
                    plan_value = billing.subscription

            if plan_value == free_value:
                agent.daily_credit_limit = default_limit
                agent.save(update_fields=['daily_credit_limit'])

    def _noop(apps, schema_editor):  # pragma: no cover - reverse noop
        return

    operations = [
        migrations.AddField(
            model_name='persistentagent',
            name='daily_credit_limit',
            field=models.PositiveIntegerField(blank=True, help_text='Maximum task credits this agent may consume per day. Null means unlimited.', null=True),
        ),
        migrations.RunPython(_set_default_daily_limits, _noop),
    ]
