"""Ordered relationships.

Revision ID: 1742033764
Revises: 1740765703
Create Date: 2025-03-15 11:16:04.256771

"""

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision: str = "1742033764"
down_revision: str | None = "1740765703"
branch_labels: str | Sequence[str] | None = ()
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "client_audience_association_table",
        sa.Column("index", sa.Integer(), nullable=True),
    )
    op.execute("""
        UPDATE client_audience_association_table
        SET 'index' = s.row_num
        FROM (
            SELECT audience_id, client_id, ROW_NUMBER() OVER (ORDER BY audience_id, client_id) as row_num
            FROM client_audience_association_table
        ) s
        WHERE client_audience_association_table.audience_id = s.audience_id AND client_audience_association_table.client_id = s.client_id;
    """)
    with op.batch_alter_table("client_audience_association_table") as batch_op:
        batch_op.alter_column("index", existing_type=sa.Integer(), autoincrement=True)

    op.add_column(
        "membership_association_table",
        sa.Column("index", sa.Integer(), autoincrement=True, nullable=True),
    )
    op.execute("""
        UPDATE membership_association_table
        SET 'index' = s.row_num
        FROM (
            SELECT user_id, group_id, ROW_NUMBER() OVER (ORDER BY user_id, group_id) as row_num
            FROM membership_association_table
        ) s
        WHERE membership_association_table.user_id = s.user_id AND membership_association_table.group_id = s.group_id;
    """)
    with op.batch_alter_table("membership_association_table") as batch_op:
        batch_op.alter_column("index", existing_type=sa.Integer(), autoincrement=True)

    op.add_column(
        "token_audience_association_table",
        sa.Column("index", sa.Integer(), autoincrement=True, nullable=True),
    )
    op.execute("""
        UPDATE token_audience_association_table
        SET 'index' = s.row_num
        FROM (
            SELECT token_id, client_id, ROW_NUMBER() OVER (ORDER BY token_id, client_id) as row_num
            FROM token_audience_association_table
        ) s
        WHERE token_audience_association_table.token_id = s.token_id AND token_audience_association_table.client_id = s.client_id;
    """)
    with op.batch_alter_table("token_audience_association_table") as batch_op:
        batch_op.alter_column("index", existing_type=sa.Integer(), autoincrement=True)

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("token_audience_association_table", "index")
    op.drop_column("membership_association_table", "index")
    op.drop_column("client_audience_association_table", "index")
    # ### end Alembic commands ###
