# Generated by Django 5.2 on 2025-07-28 00:27

from django.db import migrations
from django.db.models.expressions import F

from api.models import TaskCredit
from billing.services import BillingService


def update_additional_task_grant_and_expiration(apps, schema_editor):
    """
    Update the grant_date and expiration_date of additional task credits to match the user's billing cycle anchor.
    Note this only impacts one user - Andy - right now. haha. Going forward they're set correctly from beginning.
    """
    TaskCredit = apps.get_model("api", "TaskCredit")

    # Find task credits that are marked as additional tasks with granted_date equal to expiration_date
    task_credits = TaskCredit.objects.filter(
        additional_task=True,
        expiration_date=F('granted_date'),
    )

    anchor_day: dict[int, int] = {}

    for task_credit in task_credits:
        # If the task credit's user is not in the anchor_day dictionary, set it to the user's billing cycle anchor
        if task_credit.user.id not in anchor_day:
            # Get the user's billing cycle anchor day
            UserBilling = apps.get_model("api", "UserBilling")
            user_billing = UserBilling.objects.filter(user_id=task_credit.user.id).first()
            anchor_day[task_credit.user.id] = user_billing.billing_cycle_anchor

        # Update the granted_date and expiration_date to the user's billing cycle anchor
        start, end = BillingService.get_current_billing_period_from_day(anchor_day[task_credit.user.id])

        task_credit.granted_date = start
        task_credit.expiration_date = end

        # Save the changes
        task_credit.save(update_fields=['granted_date', 'expiration_date'])

class Migration(migrations.Migration):

    dependencies = [
        ('api', '0075_uniq_free_plan_blobackfill_billing_cycle_anchor'),
    ]

    operations = [
        migrations.RunPython(update_additional_task_grant_and_expiration, migrations.RunPython.noop),
    ]
