# Generated by Django 3.2.25 on 2024-07-31 16:38

from django.conf import settings
from django.db import migrations, models
from django.db.migrations.operations.special import RunSQL
from django.db.migrations.operations.base import Operation
from django.db import connection
from core.redis import start_job_async_or_sync
from core.models import AsyncMigrationStatus
import logging
logger = logging.getLogger(__name__)


IS_SQLITE = connection.vendor == 'sqlite'
migration_name = '0017_auto_20240731_1638'

def create_index_sql(table_name, index_name, column_name):
    return f"""
    CREATE INDEX CONCURRENTLY IF NOT EXISTS "{index_name}" ON "{table_name}" ("{column_name}");
    """

def create_fk_sql(table_name, constraint_name, column_name, referenced_table, referenced_column):
    return f"""
    ALTER TABLE "{table_name}" DROP CONSTRAINT IF EXISTS "{constraint_name}";
    ALTER TABLE "{table_name}" ADD CONSTRAINT "{constraint_name}" FOREIGN KEY ("{column_name}") REFERENCES "{referenced_table}" ("{referenced_column}") DEFERRABLE INITIALLY DEFERRED;
    """

def drop_index_sql(table_name, index_name, column_name):
    return f"""
    DROP INDEX CONCURRENTLY IF EXISTS "{index_name}";
    """

tables = [
    {
        "model_name": "azureblobexportstoragelink",
        "table_name": "io_storages_azureblobexportstoragelink",
        "index_name": "io_storages_azureblobexportstoragelink_annotation_id_6cc15c83",
        "fk_constraint": "io_storages_azureblo_annotation_id_6cc15c83_fk_task_comp",
        "column_name": "annotation_id"
    },
    {
        "model_name": "gcsexportstoragelink",
        "table_name": "io_storages_gcsexportstoragelink",
        "index_name": "io_storages_gcsexportstoragelink_annotation_id_2df715a6",
        "fk_constraint": "io_storages_gcsexpor_annotation_id_2df715a6_fk_task_comp",
        "column_name": "annotation_id"
    },
    {
        "model_name": "localfilesexportstoragelink",
        "table_name": "io_storages_localfilesexportstoragelink",
        "index_name": "io_storages_localfilesexportstoragelink_annotation_id_fc4f9825",
        "fk_constraint": "io_storages_localfil_annotation_id_fc4f9825_fk_task_comp",
        "column_name": "annotation_id"
    },
    {
        "model_name": "redisexportstoragelink",
        "table_name": "io_storages_redisexportstoragelink",
        "index_name": "io_storages_redisexportstoragelink_annotation_id_8547e508",
        "fk_constraint": "io_storages_redisexp_annotation_id_8547e508_fk_task_comp",
        "column_name": "annotation_id"
    },
    {
        "model_name": "s3exportstoragelink",
        "table_name": "io_storages_s3exportstoragelink",
        "index_name": "io_storages_s3exportstoragelink_annotation_id_729994fe",
        "fk_constraint": "io_storages_s3export_annotation_id_729994fe_fk_task_comp",
        "column_name": "annotation_id"
    }
]


def forward_migration(migration_name, db_alias):
    migration = AsyncMigrationStatus.objects.using(db_alias).create(
        name=migration_name,
        status=AsyncMigrationStatus.STATUS_STARTED,
    )
    logger.debug(
        f'Start async migration {migration_name}'
    )

    # Get db cursor
    from django.db import connections
    cursor = connections[db_alias].cursor()
    for table in tables:
        index_sql = create_index_sql(table['table_name'], table['index_name'], table['column_name'])
        fk_sql = create_fk_sql(table['table_name'], table['fk_constraint'], table['column_name'], "task_completion",
                               "id")

        # Run index_sql
        cursor.execute(index_sql)
        cursor.execute(fk_sql)

    migration.status = AsyncMigrationStatus.STATUS_FINISHED
    migration.save(using=db_alias)
    logger.debug(
        f'Async migration {migration_name} complete'
    )

def reverse_migration(migration_name, db_alias):
    migration = AsyncMigrationStatus.objects.using(db_alias).create(
        name=migration_name,
        status=AsyncMigrationStatus.STATUS_STARTED,
    )
    logger.debug(
        f'Start async migration {migration_name}'
    )

    # Get db cursor
    from django.db import connections
    cursor = connections[db_alias].cursor()
    for table in tables:
        reverse_sql = drop_index_sql(table['table_name'], table['index_name'], table['column_name'])
        # Run reverse_sql
        cursor.execute(reverse_sql)

    migration.status = AsyncMigrationStatus.STATUS_FINISHED
    migration.save(using=db_alias)
    logger.debug(
        f'Async migration {migration_name} complete'
    )


def forwards(apps, schema_editor):
    # Dispatch migrations to rqworkers
    db_alias = schema_editor.connection.alias
    start_job_async_or_sync(forward_migration, migration_name=migration_name, db_alias=db_alias)


def backwards(apps, schema_editor):
    db_alias = schema_editor.connection.alias
    start_job_async_or_sync(reverse_migration, migration_name=migration_name, db_alias=db_alias)


def get_operations():
    logger.info(f'IS_SQLITE: {IS_SQLITE}')
    if not IS_SQLITE:
        return [
            migrations.RunPython(forwards, backwards),
        ]

    operations = []
    # Use standard migration for SQLITE
    for table in tables:
        operations.append(
            migrations.AlterField(
                model_name=table['model_name'],
                name='annotation',
                field=models.ForeignKey(on_delete=models.deletion.CASCADE,
                                        related_name=table['table_name'],
                                        to='tasks.annotation'),
            ),
        )
    return operations


class Migration(migrations.Migration):
    atomic = False

    dependencies = [
        ('tasks', '0047_merge_20240318_2210'),
        ('io_storages', '0016_add_aws_sse_kms_key'),
    ]

    operations = get_operations()