# Generated by Django 5.2 on 2025-10-20 09:47

from django.db import migrations, models

CREATE_IS_VALID_FEATURE_TYPE_FUNCTION = """
CREATE OR REPLACE FUNCTION is_valid_feature_type(feature_type_id INTEGER)
RETURNS BOOLEAN AS $$
BEGIN
    -- Feature with no type is valid
    IF feature_type_id IS NULL THEN
        RETURN TRUE;
    END IF;

    -- Type must have is_type = TRUE
    RETURN EXISTS (
        SELECT 1 FROM lamindb_feature f
        WHERE f.id = feature_type_id AND f.is_type
    );
END;
$$ LANGUAGE plpgsql;
"""

CREATE_IS_VALID_FEATURE_TYPE_CONSTRAINT = """
ALTER TABLE lamindb_feature
ADD CONSTRAINT feature_type_is_valid_fk
CHECK (is_valid_feature_type(type_id));
"""


CREATE_IS_VALID_SCHEMA_TYPE_FUNCTION = """
CREATE OR REPLACE FUNCTION is_valid_schema_type(schema_type_id INTEGER)
RETURNS BOOLEAN AS $$
BEGIN
    IF schema_type_id IS NULL THEN
        RETURN TRUE;
    END IF;

    RETURN EXISTS (
        SELECT 1 FROM lamindb_schema s
        WHERE s.id = schema_type_id AND s.is_type
    );
END;
$$ LANGUAGE plpgsql;
"""

CREATE_IS_VALID_SCHEMA_TYPE_CONSTRAINT = """
ALTER TABLE lamindb_schema
ADD CONSTRAINT schema_type_is_valid_fk
CHECK (is_valid_schema_type(type_id));
"""


CREATE_IS_VALID_PROJECT_TYPE_FUNCTION = """
CREATE OR REPLACE FUNCTION is_valid_project_type(project_type_id INTEGER)
RETURNS BOOLEAN AS $$
BEGIN
    IF project_type_id IS NULL THEN
        RETURN TRUE;
    END IF;

    RETURN EXISTS (
        SELECT 1 FROM lamindb_project p
        WHERE p.id = project_type_id AND p.is_type
    );
END;
$$ LANGUAGE plpgsql;
"""

CREATE_IS_VALID_PROJECT_TYPE_CONSTRAINT = """
ALTER TABLE lamindb_project
ADD CONSTRAINT project_type_is_valid_fk
CHECK (is_valid_project_type(type_id));
"""


CREATE_IS_VALID_REFERENCE_TYPE_FUNCTION = """
CREATE OR REPLACE FUNCTION is_valid_reference_type(reference_type_id INTEGER)
RETURNS BOOLEAN AS $$
BEGIN
    IF reference_type_id IS NULL THEN
        RETURN TRUE;
    END IF;

    RETURN EXISTS (
        SELECT 1 FROM lamindb_reference r
        WHERE r.id = reference_type_id AND r.is_type
    );
END;
$$ LANGUAGE plpgsql;
"""

CREATE_IS_VALID_REFERENCE_TYPE_CONSTRAINT = """
ALTER TABLE lamindb_reference
ADD CONSTRAINT reference_type_is_valid_fk
CHECK (is_valid_reference_type(type_id));
"""


CREATE_IS_VALID_ULABEL_TYPE_FUNCTION = """
CREATE OR REPLACE FUNCTION is_valid_ulabel_type(ulabel_type_id INTEGER)
RETURNS BOOLEAN AS $$
BEGIN
    IF ulabel_type_id IS NULL THEN
        RETURN TRUE;
    END IF;

    RETURN EXISTS (
        SELECT 1 FROM lamindb_ulabel u
        WHERE u.id = ulabel_type_id AND u.is_type
    );
END;
$$ LANGUAGE plpgsql;
"""

CREATE_IS_VALID_ULABEL_TYPE_CONSTRAINT = """
ALTER TABLE lamindb_ulabel
ADD CONSTRAINT ulabel_type_is_valid_fk
CHECK (is_valid_ulabel_type(type_id));
"""


def apply_constraints(apps, schema_editor):
    if schema_editor.connection.vendor == "postgresql":
        schema_editor.execute(CREATE_IS_VALID_FEATURE_TYPE_FUNCTION)
        schema_editor.execute(CREATE_IS_VALID_FEATURE_TYPE_CONSTRAINT)
        schema_editor.execute(CREATE_IS_VALID_SCHEMA_TYPE_FUNCTION)
        schema_editor.execute(CREATE_IS_VALID_SCHEMA_TYPE_CONSTRAINT)
        schema_editor.execute(CREATE_IS_VALID_PROJECT_TYPE_FUNCTION)
        schema_editor.execute(CREATE_IS_VALID_PROJECT_TYPE_CONSTRAINT)
        schema_editor.execute(CREATE_IS_VALID_REFERENCE_TYPE_FUNCTION)
        schema_editor.execute(CREATE_IS_VALID_REFERENCE_TYPE_CONSTRAINT)
        schema_editor.execute(CREATE_IS_VALID_ULABEL_TYPE_FUNCTION)
        schema_editor.execute(CREATE_IS_VALID_ULABEL_TYPE_CONSTRAINT)


class Migration(migrations.Migration):
    dependencies = [
        ("lamindb", "0135_alter_collection_hash_and_more"),
    ]

    operations = [
        migrations.RemoveConstraint(
            model_name="record",
            name="unique_name_type_space",
        ),
        migrations.AddConstraint(
            model_name="feature",
            constraint=models.UniqueConstraint(
                condition=models.Q(("branch_id", -1), _negated=True),
                fields=("name", "type", "space"),
                name="unique_feature_name_type_space",
            ),
        ),
        migrations.AddConstraint(
            model_name="project",
            constraint=models.UniqueConstraint(
                condition=models.Q(("branch_id", -1), _negated=True),
                fields=("name", "type", "space"),
                name="unique_project_name_type_space",
            ),
        ),
        migrations.AddConstraint(
            model_name="record",
            constraint=models.UniqueConstraint(
                condition=models.Q(("branch_id", -1), _negated=True),
                fields=("name", "type", "space"),
                name="unique_record_name_type_space",
            ),
        ),
        migrations.AddConstraint(
            model_name="reference",
            constraint=models.UniqueConstraint(
                condition=models.Q(("branch_id", -1), _negated=True),
                fields=("name", "type", "space"),
                name="unique_reference_name_type_space",
            ),
        ),
        migrations.AddConstraint(
            model_name="schema",
            constraint=models.UniqueConstraint(
                condition=models.Q(("branch_id", -1), _negated=True),
                fields=("name", "type", "space"),
                name="unique_schema_name_type_space",
            ),
        ),
        migrations.AddConstraint(
            model_name="ulabel",
            constraint=models.UniqueConstraint(
                condition=models.Q(("branch_id", -1), _negated=True),
                fields=("name", "type", "space"),
                name="unique_ulabel_name_type_space",
            ),
        ),
        migrations.RunPython(apply_constraints),
    ]
