from django.db import connection, migrations
from netbox.plugins import get_plugin_config

from netbox_branching.choices import BranchStatusChoices


def copy_migrations(apps, schema_editor):
    """
    Create a copy of the migrations table in each active branch.
    """
    Branch = apps.get_model('netbox_branching', 'Branch')

    table = 'django_migrations'
    schema_prefix = get_plugin_config('netbox_branching', 'schema_prefix')

    with connection.cursor() as cursor:
        main_table = f'public.{table}'

        for branch in Branch.objects.filter(status=BranchStatusChoices.READY):
            print(f'\n    Copying migrations for branch {branch.name} ({branch.schema_id})...', end='')
            schema_name = f'{schema_prefix}{branch.schema_id}'
            schema_table = f'{schema_name}.{table}'

            # Copy the migrations table to the branch schema
            cursor.execute(f"CREATE TABLE {schema_table} ( LIKE {main_table} INCLUDING INDEXES )")

            # Copy table data
            cursor.execute(f"INSERT INTO {schema_table} SELECT * FROM {main_table}")

            # Designate id as an identity column
            cursor.execute(f"ALTER TABLE {schema_table} ALTER COLUMN id ADD GENERATED BY DEFAULT AS IDENTITY")

            # Set the next value for the ID sequence
            cursor.execute("SELECT MAX(id) from django_migrations")
            starting_id = cursor.fetchone()[0] + 1
            cursor.execute(f"ALTER SEQUENCE {schema_name}.django_migrations_id_seq RESTART WITH {starting_id}")

    print('\n ', end='')  # Padding for final "OK"


class Migration(migrations.Migration):

    dependencies = [
        ('netbox_branching', '0003_rename_indexes'),
    ]

    operations = [
        migrations.RunPython(
            code=copy_migrations,
            reverse_code=migrations.RunPython.noop
        ),
    ]
