from django.db import migrations, models


def forwards(apps, schema_editor):
    Condition = apps.get_model("atlas", "Condition")
    # Build connected components using existing M2M 'synonym' links
    visited = set()

    for cond in Condition.objects.all():
        if cond.pk in visited:
            continue
        # BFS/DFS to collect component
        stack = [cond]
        comp = []
        while stack:
            c = stack.pop()
            if c.pk in visited:
                continue
            visited.add(c.pk)
            comp.append(c)
            # access existing M2M relations
            for s in c.synonym.all():
                if s.pk not in visited:
                    stack.append(s)

        # choose canonical: prefer primary==True else lowest pk
        canon = None
        for c in comp:
            if getattr(c, "primary", False):
                canon = c
                break
        if not canon:
            canon = min(comp, key=lambda x: x.pk)

        # set canonical for all non-canonical members
        for c in comp:
            if c.pk != canon.pk:
                c.canonical_id = canon.pk
                c.save()


def backwards(apps, schema_editor):
    Condition = apps.get_model("atlas", "Condition")
    for c in Condition.objects.all():
        if c.canonical_id:
            c.canonical_id = None
            c.save()


class Migration(migrations.Migration):

    dependencies = [
        ("atlas", "0085_rename_case_detail_caseprior_casedetail_and_more"),
    ]

    operations = [
        migrations.AddField(
            model_name="condition",
            name="canonical",
            field=models.ForeignKey(
                related_name="aliases",
                null=True,
                on_delete=models.SET_NULL,
                to="atlas.condition",
                blank=True,
            ),
        ),
        migrations.RunPython(forwards, backwards),
    ]