"""move system_dependencies to egress and remove

Revision ID: 76c02f99eec1
Revises: e798f37f0c26
Create Date: 2023-05-10 17:59:14.536666

"""

import json

import sqlalchemy as sa
from alembic import op
from sqlalchemy import text
from sqlalchemy.dialects import postgresql
from sqlalchemy.exc import IntegrityError
from sqlalchemy.sql.elements import TextClause

# revision identifiers, used by Alembic.
revision = "76c02f99eec1"
down_revision = "e798f37f0c26"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    bind = op.get_bind()
    existing_system_dependencies = bind.execute(
        text(
            "SELECT id, egress, system_dependencies FROM ctl_systems where system_dependencies is not null;"
        )
    )
    for row in existing_system_dependencies:
        row_id = row["id"]
        existing_egress = row["egress"]
        system_dependencies = row["system_dependencies"]
        updated_egress = []
        updated_egress_keys = []
        if existing_egress:
            for egress in existing_egress:
                updated_egress.append(egress)
                updated_egress_keys.append(egress["fides_key"])
        for system in system_dependencies:
            if system not in updated_egress_keys:
                egress = {
                    "fides_key": system,
                    "type": "system",
                    "data_categories": None,
                }
                updated_egress.append(egress)

        update_query: TextClause = text(
            "UPDATE ctl_systems "
            "SET egress = :updated_egress ::jsonb "
            "WHERE id = :row_id"
        )
        update_params = {}
        update_params["row_id"] = row_id
        update_params["updated_egress"] = f"{json.dumps(updated_egress)}"

        try:
            bind.execute(
                update_query,
                update_params,
            )
        except IntegrityError as exc:
            raise Exception(
                f"Fides attempted to copy ctl_systems.system_dependencies into ctl_systems.egress but got error: {exc}. "
            )
    op.drop_column("ctl_systems", "system_dependencies")
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "ctl_systems",
        sa.Column(
            "system_dependencies",
            postgresql.ARRAY(sa.VARCHAR()),
            autoincrement=False,
            nullable=True,
        ),
    )
    # ### end Alembic commands ###
