from django.db import migrations, models
from django.utils import timezone


LEGACY_LEVELS = {f"level-{level}": f"level-{level + 1}" for level in range(5)}


def _tier_level(tier):
    prefix = "level-"
    suffix = tier[len(prefix):] if tier.startswith(prefix) else ""
    return int(suffix) if suffix.isdigit() else 0


def migrate_membership_levels(apps, schema_editor, reverse=False):
    Account = apps.get_model("members", "Account")
    Subscription = apps.get_model("members", "Subscription")
    CryptoInvoice = apps.get_model("members", "CryptoInvoice")
    using = schema_editor.connection.alias
    levels = {old: new for old, new in LEGACY_LEVELS.items()}
    if reverse:
        levels = {new: old for old, new in levels.items()}

    for model in (Subscription, CryptoInvoice):
        for old, new in reversed(list(levels.items())):
            model.objects.using(using).filter(tier=old).update(tier=new)

    now = timezone.now()
    account_levels = {}
    active_subscriptions = Subscription.objects.using(using).filter(
        status__in=("active", "past_due"), paid_until__gt=now
    ).values_list("account_id", "tier")
    for account_id, tier in active_subscriptions.iterator():
        account_levels[account_id] = max(account_levels.get(account_id, 0), _tier_level(tier))

    active_invoices = CryptoInvoice.objects.using(using).filter(
        status__in=("PAID", "OVERPAID"), access_start__lte=now, access_end__gt=now
    ).values_list("account_id", "tier")
    for account_id, tier in active_invoices.iterator():
        account_levels[account_id] = max(account_levels.get(account_id, 0), _tier_level(tier))

    batch = []
    for account in Account.objects.using(using).all().iterator(chunk_size=500):
        account.membership_level = account_levels.get(account.pk, 0)
        batch.append(account)
        if len(batch) == 500:
            Account.objects.using(using).bulk_update(batch, ["membership_level"])
            batch.clear()
    if batch:
        Account.objects.using(using).bulk_update(batch, ["membership_level"])


class Migration(migrations.Migration):
    dependencies = [("members", "0002_crypto_invoices")]
    operations = [
        migrations.AddField(
            model_name="account",
            name="membership_level",
            field=models.PositiveSmallIntegerField(default=0),
        ),
        migrations.RunPython(migrate_membership_levels, migrate_membership_levels),
    ]