"""Add TaskStatus.priority

Revision ID: 693cc511d28d
Revises: 328fd44c6347
Create Date: 2024-01-11 15:41:29.314852

"""

from alembic import op
import sqlalchemy as sa
import sqlalchemy_utils
import uuid
from sqlalchemy.orm.session import Session
from sqlalchemy.ext.declarative import declarative_base
from zou.migrations.utils.base import BaseMixin


# revision identifiers, used by Alembic.
revision = "693cc511d28d"
down_revision = "328fd44c6347"
branch_labels = None
depends_on = None

base = declarative_base()


class ProjectTaskStatusLink(base, BaseMixin):
    __tablename__ = "project_task_status_link"
    project_id = sa.Column(
        sqlalchemy_utils.types.UUIDType(binary=False),
        primary_key=True,
    )
    task_status_id = sa.Column(
        sqlalchemy_utils.types.UUIDType(binary=False),
        primary_key=True,
    )
    priority = sa.Column(sa.Integer, default=None)

    __table_args__ = (
        sa.UniqueConstraint(
            "project_id", "task_status_id", name="project_taskstatus_uc"
        ),
    )


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table(
        "project_task_status_link", schema=None
    ) as batch_op:
        batch_op.add_column(sa.Column("priority", sa.Integer(), nullable=True))
        batch_op.add_column(
            sa.Column("created_at", sa.DateTime(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("updated_at", sa.DateTime(), nullable=True)
        )

        batch_op.add_column(
            sa.Column(
                "id",
                sqlalchemy_utils.types.uuid.UUIDType(binary=False),
                default=uuid.uuid4,
                nullable=True,
            )
        )

    with op.batch_alter_table(
        "project_task_status_link", schema=None
    ) as batch_op:
        batch_op.execute(
            f"UPDATE project_task_status_link SET id = '{uuid.uuid4()}'"
        )

        session = Session(bind=op.get_bind())
        for link in session.query(ProjectTaskStatusLink).all():
            link.id = uuid.uuid4()
            session.add(link)
            session.commit()
        op.alter_column("project_task_status_link", "id", nullable=False)

        batch_op.create_unique_constraint(
            "project_taskstatus_uc", ["project_id", "task_status_id"]
        )

    with op.batch_alter_table("task_status", schema=None) as batch_op:
        batch_op.add_column(sa.Column("priority", sa.Integer(), nullable=True))

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("task_status", schema=None) as batch_op:
        batch_op.drop_column("priority")

    with op.batch_alter_table(
        "project_task_status_link", schema=None
    ) as batch_op:
        batch_op.drop_constraint("project_taskstatus_uc", type_="unique")
        batch_op.drop_column("priority")
        batch_op.drop_column("updated_at")
        batch_op.drop_column("created_at")
        batch_op.drop_column("id")

    # ### end Alembic commands ###
