from django.db import migrations


def add_postgres_indexes(apps, schema_editor):
    if schema_editor.connection.vendor != "postgresql":
        return
    statements = [
        "CREATE EXTENSION IF NOT EXISTS pg_trgm",
        "CREATE INDEX IF NOT EXISTS inventory_device_cpu_trgm ON inventory_device USING gin (cpu gin_trgm_ops)",
        "CREATE INDEX IF NOT EXISTS inventory_device_model_trgm ON inventory_device USING gin (model gin_trgm_ops)",
        "CREATE INDEX IF NOT EXISTS inventory_device_ticket_trgm ON inventory_device USING gin (ticket gin_trgm_ops)",
        "CREATE INDEX IF NOT EXISTS inventory_device_donor_trgm ON inventory_device USING gin (donor gin_trgm_ops)",
        "CREATE INDEX IF NOT EXISTS inventory_device_custom_values_gin ON inventory_device USING gin (custom_values jsonb_path_ops)",
    ]
    with schema_editor.connection.cursor() as cursor:
        for statement in statements:
            cursor.execute(statement)


def remove_postgres_indexes(apps, schema_editor):
    if schema_editor.connection.vendor != "postgresql":
        return
    with schema_editor.connection.cursor() as cursor:
        for name in ["inventory_device_cpu_trgm", "inventory_device_model_trgm", "inventory_device_ticket_trgm", "inventory_device_donor_trgm", "inventory_device_custom_values_gin"]:
            cursor.execute(f"DROP INDEX IF EXISTS {name}")


class Migration(migrations.Migration):
    dependencies = [("inventory", "0001_initial")]
    operations = [migrations.RunPython(add_postgres_indexes, remove_postgres_indexes)]