"""Change data models to accomodate for geometry-less global data model

Revision ID: 2ce65db29472
Revises: aa42863f0b44
Create Date: 2024-10-28 13:40:12.793386

"""

import sqlalchemy as sa
from alembic import op
from geoalchemy2 import Geometry
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = "2ce65db29472"
down_revision = "aa42863f0b44"
branch_labels = None
depends_on = None


old_geodatatype = sa.Enum(
    "geo_data",
    "generated_geo_data",
    "proposition_geo_data",
    name="geodatatype",
)
new_geodatatype = sa.Enum(
    "geo_data",
    "generated_geo_data",
    "proposition_geo_data",
    "global_data",
    name="geodatatype",
)
tmp_geodatatype = sa.Enum(
    "geo_data",
    "generated_geo_data",
    "proposition_geo_data",
    "global_data",
    name="_geodatatype",
)
featuretype = sa.Enum("feature", "geo_feature", name="featuretype")

old_datatype = sa.Enum("geo_data", "generator", name="datatype")
new_datatype = sa.Enum("geo_data", "generator", "global_data", name="datatype")
tmp_datatype = sa.Enum(
    "geo_data", "generator", "global_data", name="_datatype"
)

base_data = sa.Table(
    "base_data",
    sa.MetaData(),
    sa.Column("type", new_geodatatype, nullable=True),
    sa.Column("name", sa.String(), nullable=False),
    sa.Column("properties_modified_at", sa.DateTime(), nullable=True),
    sa.Column("created_at", sa.DateTime(), nullable=True),
    sa.Column("modified_at", sa.DateTime(), nullable=True),
    sa.Column("id", sa.Integer(), nullable=False),
    sa.Column(
        "extent",
        Geometry(
            spatial_index=False, from_text="ST_GeomFromEWKT", name="geometry"
        ),
        nullable=True,
    ),
)

feature = sa.Table(
    "feature",
    sa.MetaData(),
    sa.Column(
        "geom",
        Geometry(
            spatial_index=False, from_text="ST_GeomFromEWKT", name="geometry"
        ),
        nullable=True,
    ),
    sa.Column("geo_data_id", sa.Integer(), nullable=True),
    sa.Column(
        "explainability",
        postgresql.JSONB(astext_type=sa.Text()),
        nullable=True,
    ),
    sa.Column("id", sa.Integer(), nullable=False),
)

geo_feature = sa.Table(
    "geo_feature",
    sa.MetaData(),
    sa.Column("id", sa.Integer(), nullable=False),
    sa.Column(
        "geom",
        Geometry(
            spatial_index=False, from_text="ST_GeomFromEWKT", name="geometry"
        ),
        nullable=True,
    ),
)

project_data = sa.Table(
    "project_data",
    sa.MetaData(),
    sa.Column("data_type", new_datatype, nullable=True),
    sa.Column("id", sa.Integer(), nullable=False),
    sa.Column("data_id", sa.Integer(), nullable=True),
    sa.Column("last_update", sa.DateTime(), nullable=True),
    sa.Column("project_id", sa.Integer(), nullable=False),
    sa.Column("description", sa.String(), nullable=True),
    sa.Column("name", sa.String(), nullable=True),
    sa.Column("active_model_id", sa.Integer(), nullable=True),
)


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # Replace olf geodatatype with new one
    # https://stackoverflow.com/questions/14845203/altering-an-enum-field-using-alembic  # noqa: E501
    tmp_geodatatype.create(op.get_bind(), checkfirst=False)
    op.execute(
        "ALTER TABLE base_geo_data ALTER COLUMN type TYPE _geodatatype"
        " USING type::text::_geodatatype"
    )
    old_geodatatype.drop(op.get_bind(), checkfirst=False)
    new_geodatatype.create(op.get_bind(), checkfirst=False)
    op.execute(
        "ALTER TABLE base_geo_data ALTER COLUMN type TYPE geodatatype"
        " USING type::text::geodatatype"
    )
    tmp_geodatatype.drop(op.get_bind(), checkfirst=False)

    # Replace olf datatype with new one
    # https://stackoverflow.com/questions/14845203/altering-an-enum-field-using-alembic  # noqa: E501
    tmp_datatype.create(op.get_bind(), checkfirst=False)
    op.execute(
        "ALTER TABLE project_data ALTER COLUMN data_type TYPE _datatype"
        " USING data_type::text::_datatype"
    )
    old_datatype.drop(op.get_bind(), checkfirst=False)
    new_datatype.create(op.get_bind(), checkfirst=False)
    op.execute(
        "ALTER TABLE project_data ALTER COLUMN data_type TYPE datatype"
        " USING data_type::text::datatype"
    )
    tmp_datatype.drop(op.get_bind(), checkfirst=False)
    op.execute(
        project_data.update()
        .where(project_data.c.data_type == None)  # noqa: E711 (using ORM)
        .values(data_type="global_data")
    )

    # Remove trigger on feature (will be replaced by trigger on geo_feature)
    op.execute(sa.DDL("DROP TRIGGER feature_make_valid ON feature;"))

    with op.batch_alter_table("base_geo_data", schema=None) as batch_op:
        batch_op.drop_geospatial_index(
            "idx_base_geo_data_extent",
            postgresql_using="gist",
            column_name="extent",
        )
    op.rename_table("base_geo_data", "base_data")
    with op.batch_alter_table("base_data", schema=None) as batch_op:
        batch_op.alter_column(
            "type",
            type_=postgresql.ENUM(
                "geo_data",
                "generated_geo_data",
                "proposition_geo_data",
                "global_data",
                name="geodatatype",
                create_type=False,
            ),
        )
        batch_op.create_geospatial_index(
            "idx_base_data_extent",
            ["extent"],
            unique=False,
            postgresql_using="gist",
            postgresql_ops={},
        )

    op.create_geospatial_table(
        "geo_feature",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column(
            "geom",
            Geometry(
                spatial_index=False,
                from_text="ST_GeomFromEWKT",
                name="geometry",
            ),
            nullable=True,
        ),
        sa.ForeignKeyConstraint(
            ["id"], ["feature.id"], name=op.f("fk_geo_feature_id_feature")
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_geo_feature")),
    )
    with op.batch_alter_table("geo_feature", schema=None) as batch_op:
        batch_op.create_geospatial_index(
            "idx_geo_feature_geom",
            ["geom"],
            unique=False,
            postgresql_using="gist",
            postgresql_ops={},
        )

    # Populate geo_feature with all features (as they are all geo_feature)
    op.execute(
        geo_feature.insert().from_select(
            names=["id", "geom"], select=sa.select(feature.c["id", "geom"])
        )
    )

    op.create_table(
        "global_data",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("feature_id", sa.Integer(), nullable=True),
        sa.Column("description", sa.String(), nullable=True),
        sa.Column("is_public", sa.Boolean(), nullable=False),
        sa.Column("upload_user_id", sa.Integer(), nullable=True),
        sa.ForeignKeyConstraint(
            ["feature_id"],
            ["feature.id"],
            name=op.f("fk_global_data_feature_id_feature"),
        ),
        sa.ForeignKeyConstraint(
            ["id"], ["base_data.id"], name=op.f("fk_global_data_id_base_data")
        ),
        sa.ForeignKeyConstraint(
            ["upload_user_id"],
            ["user.id"],
            name=op.f("fk_global_data_upload_user_id_user"),
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_global_data")),
    )
    op.create_table(
        "global_data_permission",
        sa.Column("object_id", sa.Integer(), nullable=True),
        sa.Column("user_id", sa.Integer(), nullable=False),
        sa.Column("id", sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(
            ["object_id"],
            ["global_data.id"],
            name=op.f("fk_global_data_permission_object_id_global_data"),
        ),
        sa.ForeignKeyConstraint(
            ["user_id"],
            ["user.id"],
            name=op.f("fk_global_data_permission_user_id_user"),
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_global_data_permission")),
    )
    with op.batch_alter_table("data_attribute", schema=None) as batch_op:
        batch_op.alter_column("geo_data_id", new_column_name="data_id")
        batch_op.drop_index("ix_data_attribute_geo_data_id")
        batch_op.create_index(
            batch_op.f("ix_data_attribute_data_id"), ["data_id"], unique=False
        )
        batch_op.drop_constraint(
            "fk_data_attribute_geo_data_id_base_geo_data", type_="foreignkey"
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_data_attribute_data_id_base_data"),
            "base_data",
            ["data_id"],
            ["id"],
        )

    # Create FeatureType enum type (else it fails)
    featuretype.create(op.get_bind(), checkfirst=False)
    with op.batch_alter_table("feature", schema=None) as batch_op:
        batch_op.add_column(sa.Column("type", featuretype, nullable=True))
        batch_op.alter_column("geo_data_id", new_column_name="data_id")
        batch_op.drop_geospatial_index(
            "idx_feature_geom", postgresql_using="gist", column_name="geom"
        )
        batch_op.drop_index("ix_feature_geo_data_id")
        batch_op.drop_constraint(
            "fk_feature_geo_data_id_base_geo_data", type_="foreignkey"
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_feature_data_id_base_data"),
            "base_data",
            ["data_id"],
            ["id"],
        )
        batch_op.drop_geospatial_column("geom")
    # Only geometry feature existed before
    op.execute("UPDATE feature SET type = 'geo_feature'")

    with op.batch_alter_table("geo_data", schema=None) as batch_op:
        batch_op.drop_constraint(
            "fk_geo_data_id_base_geo_data", type_="foreignkey"
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_geo_data_id_base_data"), "base_data", ["id"], ["id"]
        )

    with op.batch_alter_table("project_data", schema=None) as batch_op:
        batch_op.drop_constraint(
            "fk_project_data_data_id_base_geo_data", type_="foreignkey"
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_project_data_data_id_base_data"),
            "base_data",
            ["data_id"],
            ["id"],
        )

    with op.batch_alter_table("zone_proposition", schema=None) as batch_op:
        batch_op.drop_constraint(
            "fk_zone_proposition_data_id_base_geo_data", type_="foreignkey"
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_zone_proposition_data_id_base_data"),
            "base_data",
            ["data_id"],
            ["id"],
        )

    # ### end Alembic commands ###
    op.execute(
        sa.DDL(
            """
/**
 * Trigger to correct the invalid geometries before an INSERT or an UPDATE.
 */
CREATE OR REPLACE FUNCTION feature_make_valid() RETURNS trigger AS $feature_make_valid$
BEGIN
    IF ST_IsValid(NEW.geom) = 'f' THEN
        NEW.geom := ST_MakeValid(NEW.geom);
    END IF;
    RETURN NEW;
END;
$feature_make_valid$ LANGUAGE plpgsql;

CREATE TRIGGER feature_make_valid BEFORE INSERT OR UPDATE ON geo_feature
    FOR EACH ROW EXECUTE PROCEDURE feature_make_valid();
    """  # noqa: E501
        )
    )
    op.execute(
        sa.DDL(
            """
/**
 * Dissolve adjacent geometries if they have the same attributes.
 *
 * Warning: the function deletes the old features before inserting the new ones.
 */
CREATE OR REPLACE FUNCTION dissolve_adjacent (data_id int) RETURNS void AS $$
BEGIN
    -- Use a temporary table to store the attributes
    CREATE TEMP TABLE temp_dissolve AS
    WITH merge AS (
        SELECT g.geom AS geom, json_object_agg(v.attribute_id, v.value) AS props
        FROM feature f INNER JOIN geo_feature g ON f.id = g.id
        INNER JOIN data_value v ON f.id = v.feature_id
        WHERE f.data_id = data_id
        GROUP BY f.id
    )
    -- Dissolve the adjacent geometries based on their attributes
    SELECT ST_SetSRID(ST_UnaryUnion(d.geom), 4326) AS geom, d.props AS props
    FROM (
        SELECT UNNEST(ST_ClusterIntersecting(m.geom)) AS geom, m.props::jsonb AS props
        FROM merge m
        GROUP BY m.props::jsonb
    ) AS d;

    -- Delete old features
    WITH delf AS (
        DELETE FROM feature f WHERE f.data_id = data_id
        RETURNING f.id
    )
    DELETE FROM geo_feature WHERE id IN (SELECT * FROM delf);

    -- Insert the new dissolved features
    WITH ins AS (
        INSERT INTO feature (data_id)
        SELECT geom, data_id
        FROM temp_dissolve
        RETURNING id, geom
    )
    INSERT INTO geo_feature (id, geom)
    SELECT id, geom FROM ins;
    INSERT INTO data_value (attribute_id, value, feature_id, type)
    SELECT d.key, d.value, ins.id, a.type
    FROM ins INNER JOIN (
        SELECT t.geom, j.key::int, j.value
        FROM temp_dissolve t, jsonb_each(t.props) j
    ) AS d ON ins.geom = d.geom
    INNER JOIN data_attribute a ON d.key = a.id;

    -- Delete the temporary table
    DROP TABLE temp_dissolve;
END;
$$ LANGUAGE plpgsql;
    """  # noqa: E501
        )
    )


def downgrade():
    op.execute(sa.DDL("DROP TRIGGER feature_make_valid ON geo_feature;"))

    # Delete all ProjectGlobalData objects
    # Need to delete all objects referencing project global data first
    op.execute(
        "DELETE FROM data_input_association pd WHERE EXISTS "
        "(SELECT 1 FROM project_data p WHERE p.id = pd.input_data_id AND "
        "p.data_type = 'global_data');"
    )
    op.execute(
        "DELETE FROM project_shared_data pd WHERE EXISTS "
        "(SELECT 1 FROM project_data p WHERE p.id = pd.project_data_id AND "
        "p.data_type = 'global_data');"
    )
    op.execute(
        project_data.delete().where(project_data.c.data_type == "global_data")
    )

    # Replace olf datatype with new one
    # https://stackoverflow.com/questions/14845203/altering-an-enum-field-using-alembic  # noqa: E501
    tmp_datatype.create(op.get_bind(), checkfirst=False)
    op.execute(
        "ALTER TABLE project_data ALTER COLUMN data_type TYPE _datatype"
        " USING data_type::text::_datatype"
    )
    new_datatype.drop(op.get_bind(), checkfirst=False)
    old_datatype.create(op.get_bind(), checkfirst=False)
    op.execute(
        "ALTER TABLE project_data ALTER COLUMN data_type TYPE datatype"
        " USING data_type::text::datatype"
    )
    tmp_datatype.drop(op.get_bind(), checkfirst=False)

    # Change geodatatype enum columns
    # https://stackoverflow.com/questions/14845203/altering-an-enum-field-using-alembic  # noqa: E501
    op.execute(
        base_data.update()
        .where(base_data.c.type == "global_data")
        .values(type=None)
    )
    tmp_geodatatype.create(op.get_bind(), checkfirst=False)
    op.execute(
        "ALTER TABLE base_data ALTER COLUMN type TYPE _geodatatype"
        " USING type::text::_geodatatype"
    )
    new_geodatatype.drop(op.get_bind(), checkfirst=False)
    old_geodatatype.create(op.get_bind(), checkfirst=False)
    op.execute(
        "ALTER TABLE base_data ALTER COLUMN type TYPE geodatatype"
        " USING type::text::geodatatype"
    )
    tmp_geodatatype.drop(op.get_bind(), checkfirst=False)
    # ### commands auto generated by Alembic - please adjust! ###

    with op.batch_alter_table("base_data", schema=None) as batch_op:
        batch_op.drop_geospatial_index(
            "idx_base_data_extent",
            postgresql_using="gist",
            column_name="extent",
        )
    op.rename_table("base_data", "base_geo_data")
    with op.batch_alter_table("base_geo_data", schema=None) as batch_op:
        batch_op.alter_column(
            "type",
            type_=postgresql.ENUM(
                "geo_data",
                "generated_geo_data",
                "proposition_geo_data",
                name="geodatatype",
                create_type=False,
            ),
        )
        batch_op.create_geospatial_index(
            "idx_base_geo_data_extent",
            ["extent"],
            unique=False,
            postgresql_using="gist",
            postgresql_ops={},
        )

    with op.batch_alter_table("zone_proposition", schema=None) as batch_op:
        batch_op.drop_constraint(
            batch_op.f("fk_zone_proposition_data_id_base_data"),
            type_="foreignkey",
        )
        batch_op.create_foreign_key(
            "fk_zone_proposition_data_id_base_geo_data",
            "base_geo_data",
            ["data_id"],
            ["id"],
        )

    with op.batch_alter_table("project_data", schema=None) as batch_op:
        batch_op.drop_constraint(
            batch_op.f("fk_project_data_data_id_base_data"), type_="foreignkey"
        )
        batch_op.create_foreign_key(
            "fk_project_data_data_id_base_geo_data",
            "base_geo_data",
            ["data_id"],
            ["id"],
        )

    with op.batch_alter_table("geo_data", schema=None) as batch_op:
        batch_op.drop_constraint(
            batch_op.f("fk_geo_data_id_base_data"), type_="foreignkey"
        )
        batch_op.create_foreign_key(
            "fk_geo_data_id_base_geo_data", "base_geo_data", ["id"], ["id"]
        )

    with op.batch_alter_table("feature", schema=None) as batch_op:
        batch_op.alter_column("data_id", new_column_name="geo_data_id")
        batch_op.add_geospatial_column(
            sa.Column(
                "geom",
                Geometry(
                    spatial_index=False,
                    from_text="ST_GeomFromEWKT",
                    name="geometry",
                ),
                autoincrement=False,
                nullable=True,
            )
        )
        batch_op.drop_constraint(
            batch_op.f("fk_feature_data_id_base_data"), type_="foreignkey"
        )
        batch_op.create_foreign_key(
            "fk_feature_geo_data_id_base_geo_data",
            "base_geo_data",
            ["geo_data_id"],
            ["id"],
        )
        batch_op.create_index(
            "ix_feature_geo_data_id", ["geo_data_id"], unique=False
        )
        batch_op.create_geospatial_index(
            "idx_feature_geom",
            ["geom"],
            unique=False,
            postgresql_using="gist",
            postgresql_ops={},
        )
        batch_op.drop_column("type")

    # Populate feature geometry with all features (as they are all geo_feature)
    op.execute(
        "UPDATE feature SET geom = geo_feature.geom "
        "FROM geo_feature WHERE feature.id = geo_feature.id"
    )

    with op.batch_alter_table("data_attribute", schema=None) as batch_op:
        batch_op.alter_column("data_id", new_column_name="geo_data_id")
        batch_op.drop_constraint(
            batch_op.f("fk_data_attribute_data_id_base_data"),
            type_="foreignkey",
        )
        batch_op.create_foreign_key(
            "fk_data_attribute_geo_data_id_base_geo_data",
            "base_geo_data",
            ["geo_data_id"],
            ["id"],
        )
        batch_op.drop_index(batch_op.f("ix_data_attribute_data_id"))
        batch_op.create_index(
            "ix_data_attribute_geo_data_id", ["geo_data_id"], unique=False
        )

    op.drop_table("global_data_permission")
    op.drop_table("global_data")
    with op.batch_alter_table("geo_feature", schema=None) as batch_op:
        batch_op.drop_geospatial_index(
            "idx_geo_feature_geom", postgresql_using="gist", column_name="geom"
        )

    op.drop_geospatial_table("geo_feature")

    # Delete all base_geo_data without type (i.e all global data)
    # Need to delete all objects referencing global data first
    # (easier to do there, as we don't need to dereference them from base_data)
    op.execute(
        "DELETE FROM data_value WHERE feature_id IN "
        "(SELECT f.id AS feature_id FROM feature f JOIN base_geo_data g ON "
        "f.geo_data_id = g.id WHERE g.type IS null)"
    )
    op.execute(
        "DELETE FROM data_attribute WHERE geo_data_id IN "
        "(SELECT id AS geo_data_id FROM base_geo_data WHERE type IS null)"
    )
    op.execute(
        "DELETE FROM feature WHERE geo_data_id IN "
        "(SELECT id AS geo_data_id FROM base_geo_data WHERE type IS null)"
    )
    op.execute("DELETE FROM base_geo_data WHERE type IS null")

    featuretype.drop(op.get_bind(), checkfirst=False)
    # ### end Alembic commands ###

    op.execute(
        sa.DDL(
            """
/**
 * Trigger to correct the invalid geometries before an INSERT or an UPDATE.
 */
CREATE OR REPLACE FUNCTION feature_make_valid() RETURNS trigger AS $feature_make_valid$
BEGIN
    IF ST_IsValid(NEW.geom) = 'f' THEN
        NEW.geom := ST_MakeValid(NEW.geom);
    END IF;
    RETURN NEW;
END;
$feature_make_valid$ LANGUAGE plpgsql;

CREATE TRIGGER feature_make_valid BEFORE INSERT OR UPDATE ON feature
    FOR EACH ROW EXECUTE PROCEDURE feature_make_valid();
        """  # noqa: E501
        )
    )
    op.execute(
        sa.DDL(
            """
/**
 * Dissolve adjacent geometries if they have the same attributes.
 *
 * Warning: the function deletes the old features before inserting the new ones.
 */
CREATE OR REPLACE FUNCTION dissolve_adjacent (data_id int) RETURNS void AS $$
BEGIN
    -- Use a temporary table to store the attributes
    CREATE TEMP TABLE temp_dissolve AS
    WITH merge AS (
        SELECT f.geom AS geom, json_object_agg(v.attribute_id, v.value) AS props
        FROM feature f INNER JOIN data_value v ON f.id = v.feature_id
        WHERE f.geo_data_id = data_id
        GROUP BY f.id
    )
    -- Dissolve the adjacent geometries based on their attributes
    SELECT ST_SetSRID(ST_UnaryUnion(d.geom), 4326) AS geom, d.props AS props
    FROM (
        SELECT UNNEST(ST_ClusterIntersecting(m.geom)) AS geom, m.props::jsonb AS props
        FROM merge m
        GROUP BY m.props::jsonb
    ) AS d;

    -- Delete old features
    DELETE FROM feature WHERE geo_data_id = data_id;

    -- Insert the new dissolved features
    WITH ins AS (
        INSERT INTO feature (geom, geo_data_id)
        SELECT geom, data_id
        FROM temp_dissolve
        RETURNING id, geom
    )
    INSERT INTO data_value (attribute_id, value, feature_id, type)
    SELECT d.key, d.value, ins.id, a.type
    FROM ins INNER JOIN (
        SELECT t.geom, j.key::int, j.value
        FROM temp_dissolve t, jsonb_each(t.props) j
    ) AS d ON ins.geom = d.geom
    INNER JOIN data_attribute a ON d.key = a.id;

    -- Delete the temporary table
    DROP TABLE temp_dissolve;
END;
$$ LANGUAGE plpgsql;
        """  # noqa: E501
        )
    )
