"""Add eligible task roles to templates and engagement task snapshots."""

from alembic import op
import sqlalchemy as sa

revision = "20260808_task_role_eligibility"
down_revision = "20260807_service_reg_master"
branch_labels = None
depends_on = None

TABLES = (
    "service_default_task_templates",
    "firm_service_task_templates",
    "client_service_task_instances",
)

def upgrade():
    bind = op.get_bind()
    inspector = sa.inspect(bind)
    tables = set(inspector.get_table_names())
    for table_name in TABLES:
        if table_name not in tables:
            continue
        columns = {c["name"] for c in inspector.get_columns(table_name)}
        if "eligible_role_names" not in columns:
            with op.batch_alter_table(table_name) as batch:
                batch.add_column(sa.Column("eligible_role_names", sa.String(length=200), nullable=True))
        bind.execute(sa.text(
            f"""
            UPDATE {table_name}
               SET eligible_role_names = default_role_name
             WHERE (eligible_role_names IS NULL OR TRIM(eligible_role_names) = '')
               AND default_role_name IS NOT NULL
               AND TRIM(default_role_name) <> ''
            """
        ))

def downgrade():
    bind = op.get_bind()
    inspector = sa.inspect(bind)
    tables = set(inspector.get_table_names())
    for table_name in reversed(TABLES):
        if table_name not in tables:
            continue
        columns = {c["name"] for c in inspector.get_columns(table_name)}
        if "eligible_role_names" in columns:
            with op.batch_alter_table(table_name) as batch:
                batch.drop_column("eligible_role_names")