"""Fix primary key and unique index on PluginTaskInputParameter and PluginTaskOutputParameter

Revision ID: 5b80d059cbb4
Revises: b5cc97666e3d
Create Date: 2024-06-27 13:22:22.111192

"""

from typing import Annotated

import sqlalchemy as sa
from alembic import op
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    MappedAsDataclass,
    mapped_column,
    sessionmaker,
)

# revision identifiers, used by Alembic.
revision = "5b80d059cbb4"
down_revision = "b5cc97666e3d"
branch_labels = None
depends_on = None


# Migration data models
intpk = Annotated[
    int,
    mapped_column(sa.BigInteger().with_variant(sa.Integer, "sqlite"), primary_key=True),
]
text_ = Annotated[str, mapped_column(sa.Text())]
bool_ = Annotated[bool, mapped_column(sa.Boolean())]


class UpgradeBase(DeclarativeBase, MappedAsDataclass):
    pass


class PluginTaskInputParameterUpgrade(UpgradeBase):
    __tablename__ = "plugin_task_input_parameters"

    plugin_file_resource_snapshot_id: Mapped[intpk]
    plugin_task_name: Mapped[text_] = mapped_column(primary_key=True)
    name: Mapped[text_] = mapped_column(nullable=False, primary_key=True)
    required: Mapped[bool_] = mapped_column(nullable=False)


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    bind = op.get_bind()
    Session = sessionmaker(bind=bind)

    with op.batch_alter_table("plugin_task_input_parameters", schema=None) as batch_op:
        batch_op.add_column(sa.Column("required", sa.Boolean(), nullable=True))

    # Set the new required field on any preexisting input_parameters to True.
    with Session() as session:
        for param in session.scalars(sa.select(PluginTaskInputParameterUpgrade)):
            param.required = True

        session.commit()

    # Set the required field to nullable=False now that it is populated with values.
    with op.batch_alter_table("plugin_task_input_parameters", schema=None) as batch_op:
        batch_op.alter_column("required", nullable=False)

    with op.batch_alter_table("plugin_task_input_parameters", schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f("pk_plugin_task_input_parameters"))
        batch_op.drop_index(
            batch_op.f(
                "ix_plugin_task_input_parameters_plugin_file_resource_snapshot_id"
            ),
        )
        batch_op.create_primary_key(
            batch_op.f("pk_plugin_task_input_parameters"),
            ["plugin_file_resource_snapshot_id", "plugin_task_name", "name"],
        )
        batch_op.create_index(
            batch_op.f(
                "ix_plugin_task_input_parameters_plugin_file_resource_snapshot_id"
            ),
            ["plugin_file_resource_snapshot_id", "plugin_task_name", "name"],
            unique=True,
        )

    with op.batch_alter_table("plugin_task_output_parameters", schema=None) as batch_op:
        batch_op.drop_constraint(batch_op.f("pk_plugin_task_output_parameters"))
        batch_op.drop_index(
            batch_op.f(
                "ix_plugin_task_output_parameters_plugin_file_resource_snapshot_id"
            )
        )
        batch_op.create_primary_key(
            batch_op.f("pk_plugin_task_output_parameters"),
            ["plugin_file_resource_snapshot_id", "plugin_task_name", "name"],
        )
        batch_op.create_index(
            batch_op.f(
                "ix_plugin_task_output_parameters_plugin_file_resource_snapshot_id"
            ),
            ["plugin_file_resource_snapshot_id", "plugin_task_name", "name"],
            unique=True,
        )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # NOTE: We are not changing the primary key and unique index back to their original
    # state because it would require dropping data from existing tables. Keeping the
    # redefinition should not cause any issues in prior versions, and redoing the
    # upgrade just drops and replaces the same index/primary key.
    with op.batch_alter_table("plugin_task_input_parameters", schema=None) as batch_op:
        batch_op.drop_column("required")

    # ### end Alembic commands ###
