# Generated by Django 4.2.9 on 2024-02-21 11:53

from datetime import datetime

from django.db import migrations
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps


def add_months(dt: datetime, months: int) -> datetime:
    assert months >= 0
    # It's fine that the max day in Feb is 28 for leap years.
    MAX_DAY_FOR_MONTH = {
        1: 31,
        2: 28,
        3: 31,
        4: 30,
        5: 31,
        6: 30,
        7: 31,
        8: 31,
        9: 30,
        10: 31,
        11: 30,
        12: 31,
    }
    year = dt.year
    month = dt.month + months
    while month > 12:
        year += 1
        month -= 12
    day = min(dt.day, MAX_DAY_FOR_MONTH[month])
    # datetimes don't support leap seconds, so don't need to worry about those
    return dt.replace(year=year, month=month, day=day)


def backfill_end_date_for_fixed_price_plans(
    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
) -> None:
    CustomerPlan = apps.get_model("corporate", "CustomerPlan")
    CustomerPlan.ACTIVE = 1

    fixed_price_plans = CustomerPlan.objects.filter(
        status=CustomerPlan.ACTIVE, fixed_price__isnull=False
    )
    for plan in fixed_price_plans:
        assert plan.end_date is None
        assert plan.billing_cycle_anchor is not None
        plan.end_date = add_months(plan.billing_cycle_anchor, 12)
        plan.save(update_fields=["end_date"])


class Migration(migrations.Migration):
    dependencies = [
        ("corporate", "0038_customerplanoffer_sent_invoice_id_invoice"),
    ]

    operations = [
        migrations.RunPython(
            backfill_end_date_for_fixed_price_plans,
            elidable=True,
        ),
    ]
