# Generated by Django 5.2 on 2025-08-07 18:52

from django.db import migrations

CREATE_FUNCTION_SQL = """
CREATE OR REPLACE FUNCTION is_valid_record_type(record_type_id INTEGER, record_is_type BOOLEAN)
RETURNS BOOLEAN AS $$
BEGIN
    -- Record with no type is valid
    IF record_type_id IS NULL THEN
        RETURN TRUE;
    END IF;

    -- If current record is a type, it can only reference schema-less types
    IF record_is_type THEN
        RETURN EXISTS (
            SELECT 1 FROM lamindb_record r
            WHERE r.id = record_type_id AND r.is_type AND r.schema_id IS NULL
        );
    END IF;

    -- Regular records can reference any type
    RETURN EXISTS (
        SELECT 1 FROM lamindb_record r
        WHERE r.id = record_type_id AND r.is_type
    );
END;
$$ LANGUAGE plpgsql;
"""

ADD_CONSTRAINT_SQL = """
ALTER TABLE lamindb_record
ADD CONSTRAINT record_type_is_valid_fk
CHECK (is_valid_record_type(type_id, is_type));
"""

DROP_CONSTRAINT_SQL = (
    "ALTER TABLE lamindb_record DROP CONSTRAINT IF EXISTS record_type_is_valid_fk;"
)
DROP_FUNCTION_SQL = "DROP FUNCTION IF EXISTS is_valid_record_type(INTEGER, BOOLEAN);"


def apply_postgres_constraint(apps, schema_editor):
    if schema_editor.connection.vendor == "postgresql":
        schema_editor.execute(CREATE_FUNCTION_SQL)
        schema_editor.execute(ADD_CONSTRAINT_SQL)


def revert_postgres_constraint(apps, schema_editor):
    if schema_editor.connection.vendor == "postgresql":
        schema_editor.execute(DROP_CONSTRAINT_SQL)
        schema_editor.execute(DROP_FUNCTION_SQL)


class Migration(migrations.Migration):
    dependencies = [
        ("lamindb", "0119_squashed"),
    ]

    operations = [
        migrations.RunPython(
            apply_postgres_constraint, reverse_code=revert_postgres_constraint
        ),
    ]
