"""add period label to client service task instances

Revision ID: 20260805_task_instance_period_label
Revises: 20260805_periodic_engagements
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy import inspect

revision = "20260805_task_instance_period_label"
down_revision = "20260805_periodic_engagements"
branch_labels = None
depends_on = None

TABLE = "client_service_task_instances"
INDEX = "ix_csti_period_label"


def _inspector():
    return inspect(op.get_bind())


def _columns() -> set[str]:
    return {row["name"] for row in _inspector().get_columns(TABLE)}


def _indexes() -> set[str]:
    return {row.get("name") for row in _inspector().get_indexes(TABLE) if row.get("name")}


def upgrade():
    if "period_label" not in _columns():
        op.add_column(
            TABLE,
            sa.Column("period_label", sa.String(length=30), nullable=False, server_default=""),
        )

    op.execute(
        """
        UPDATE client_service_task_instances AS task
        SET period_label = COALESCE(subscription.period_label, '')
        FROM client_service_subscriptions AS subscription
        WHERE task.subscription_id = subscription.id
          AND COALESCE(task.period_label, '') = ''
        """
    )

    if INDEX not in _indexes():
        op.create_index(INDEX, TABLE, ["period_label"], unique=False)


def downgrade():
    if INDEX in _indexes():
        op.drop_index(INDEX, table_name=TABLE)
    if "period_label" in _columns():
        op.drop_column(TABLE, "period_label")
