"""Add eplace version with checksum

Revision ID: b75d279ba985
Revises: fd966d4ad99c
Create Date: 2022-04-25 13:52:07.514013

"""

import sqlalchemy as sa
from alembic import op

import prefect
from prefect.blocks.core import Block

# revision identifiers, used by Alembic.
revision = "b75d279ba985"
down_revision = "fd966d4ad99c"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "block_type",
        sa.Column(
            "id",
            prefect.server.utilities.database.UUID(),
            server_default=sa.text(
                "(\n    (\n        lower(hex(randomblob(4))) \n        || '-' \n       "
                " || lower(hex(randomblob(2))) \n        || '-4' \n        ||"
                " substr(lower(hex(randomblob(2))),2) \n        || '-' \n        ||"
                " substr('89ab',abs(random()) % 4 + 1, 1) \n        ||"
                " substr(lower(hex(randomblob(2))),2) \n        || '-' \n        ||"
                " lower(hex(randomblob(6)))\n    )\n    )"
            ),
            nullable=False,
        ),
        sa.Column(
            "created",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("(strftime('%Y-%m-%d %H:%M:%f000', 'now'))"),
            nullable=False,
        ),
        sa.Column(
            "updated",
            prefect.server.utilities.database.Timestamp(timezone=True),
            server_default=sa.text("(strftime('%Y-%m-%d %H:%M:%f000', 'now'))"),
            nullable=False,
        ),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("logo_url", sa.String(), nullable=True),
        sa.Column("documentation_url", sa.String(), nullable=True),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_block_type")),
    )
    with op.batch_alter_table("block_type", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_block_type__updated"), ["updated"], unique=False
        )
        batch_op.create_index(batch_op.f("uq_block_type__name"), ["name"], unique=True)

    with op.batch_alter_table("block_document", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column(
                "block_type_id", prefect.server.utilities.database.UUID(), nullable=True
            )
        )
        batch_op.drop_index("uq_block__schema_id_name")
        batch_op.create_index(
            "uq_block__type_id_name", ["block_type_id", "name"], unique=True
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_block_document__block_type_id__block_type"),
            "block_type",
            ["block_type_id"],
            ["id"],
            ondelete="cascade",
        )
        batch_op.drop_constraint("fk_block__block_schema_id__block_schema")

    with op.batch_alter_table(
        "block_schema",
        schema=None,
    ) as batch_op:
        batch_op.add_column(
            sa.Column(
                "block_type_id", prefect.server.utilities.database.UUID(), nullable=True
            )
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_block_schema__block_type_id__block_type"),
            "block_type",
            ["block_type_id"],
            ["id"],
            ondelete="cascade",
        )
        batch_op.drop_column("version")
        batch_op.add_column(sa.Column("checksum", sa.String(), nullable=True))
        batch_op.drop_index("uq_block_schema__name_version")
        batch_op.create_index(
            batch_op.f("ix_block_schema__checksum"), ["checksum"], unique=False
        )
        batch_op.create_index("uq_block_schema__checksum", ["checksum"], unique=True)

    # Add checksums and block types for existing block schemas
    connection = op.get_bind()
    meta_data = sa.MetaData()
    meta_data.reflect(connection)
    BLOCK_SCHEMA = meta_data.tables["block_schema"]
    BLOCK_TYPE = meta_data.tables["block_type"]
    BLOCK_DOCUMENT = meta_data.tables["block_document"]
    results = connection.execute(
        sa.select(BLOCK_SCHEMA.c.id, BLOCK_SCHEMA.c.name, BLOCK_SCHEMA.c.fields)
    )
    for id, name, fields in results:
        schema_checksum = Block._calculate_schema_checksum(fields)
        # Add checksum
        connection.execute(
            sa.update(BLOCK_SCHEMA)
            .where(BLOCK_SCHEMA.c.id == id)
            .values(checksum=schema_checksum)
        )
        # Create corresponding block type
        block_type_result = connection.execute(
            sa.select(BLOCK_TYPE.c.id).where(BLOCK_TYPE.c.name == name)
        ).first()
        if block_type_result is None:
            # Create block type if it doesn't already exist
            connection.execute(sa.insert(BLOCK_TYPE).values(name=name))
        block_type_result = connection.execute(
            sa.select(BLOCK_TYPE.c.id).where(BLOCK_TYPE.c.name == name)
        ).first()
        new_block_type_id = block_type_result[0]
        connection.execute(
            sa.update(BLOCK_SCHEMA)
            .where(BLOCK_SCHEMA.c.id == id)
            .values(block_type_id=new_block_type_id)
        )
        # Associate new block type will all block documents for this block schema
        block_document_results = connection.execute(
            sa.select(BLOCK_DOCUMENT.c.id).where(BLOCK_DOCUMENT.c.block_schema_id == id)
        ).all()
        for (block_document_id,) in block_document_results:
            connection.execute(
                sa.update(BLOCK_DOCUMENT)
                .where(BLOCK_DOCUMENT.c.id == block_document_id)
                .values(block_type_id=new_block_type_id)
            )

    with op.batch_alter_table(
        "block_schema",
        schema=None,
    ) as batch_op:
        batch_op.drop_column("name")
        batch_op.alter_column("checksum", existing_type=sa.VARCHAR(), nullable=False)
        batch_op.alter_column(
            "block_type_id", existing_type=sa.VARCHAR(), nullable=False
        )

    with op.batch_alter_table(
        "block_document",
        schema=None,
    ) as batch_op:
        batch_op.alter_column(
            "block_type_id", existing_type=sa.VARCHAR(), nullable=False
        )

        batch_op.create_foreign_key(
            batch_op.f("fk_block__block_schema_id__block_schema"),
            "block_schema",
            ["block_schema_id"],
            ["id"],
            ondelete="cascade",
        )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("block_schema", schema=None) as batch_op:
        batch_op.drop_index("uq_block_schema__checksum")
        batch_op.drop_index(batch_op.f("ix_block_schema__checksum"))
        batch_op.create_index(
            "uq_block_schema__name_version", ["name", "version"], unique=False
        )
        batch_op.drop_column("checksum")
        batch_op.add_column(sa.Column("version", sa.VARCHAR(), nullable=True))
        batch_op.add_column(sa.Column("name", sa.VARCHAR(), nullable=True))
        batch_op.drop_constraint(
            batch_op.f("fk_block_schema__block_type_id__block_type"), type_="foreignkey"
        )
        batch_op.drop_column("block_type_id")

    with op.batch_alter_table("block_document", schema=None) as batch_op:
        batch_op.drop_constraint(
            batch_op.f("fk_block_document__block_type_id__block_type"),
            type_="foreignkey",
        )
        batch_op.drop_index("uq_block__type_id_name")
        batch_op.create_index(
            "uq_block__schema_id_name", ["block_schema_id", "name"], unique=False
        )
        batch_op.drop_column("block_type_id")

    with op.batch_alter_table("block_type", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_block_type__updated"))
        batch_op.drop_index(batch_op.f("uq_block_type__name"))

    op.drop_table("block_type")

    # ### end Alembic commands ###
