"""sqlalchemy association tables.

Revision ID: 1761860792
Revises: 1761851899
Create Date: 2025-10-30 22:46:32.173185

"""

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op

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


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###

    # Fill NULL index values before making columns NOT NULL
    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) - 1 as row_num
            FROM client_audience_association_table
            WHERE "index" IS NULL
        ) s
        WHERE client_audience_association_table.audience_id = s.audience_id
          AND client_audience_association_table.client_id = s.client_id;
    """)

    op.execute("""
        UPDATE membership_association_table
        SET "index" = s.row_num
        FROM (
            SELECT user_id, group_id, ROW_NUMBER() OVER (PARTITION BY group_id ORDER BY user_id) - 1 as row_num
            FROM membership_association_table
            WHERE "index" IS NULL
        ) s
        WHERE membership_association_table.user_id = s.user_id
          AND membership_association_table.group_id = s.group_id;
    """)

    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) - 1 as row_num
            FROM token_audience_association_table
            WHERE "index" IS NULL
        ) 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(
        "client_audience_association_table", schema=None
    ) as batch_op:
        batch_op.alter_column("index", existing_type=sa.INTEGER(), nullable=False)

    with op.batch_alter_table("membership_association_table", schema=None) as batch_op:
        batch_op.alter_column("index", existing_type=sa.INTEGER(), nullable=False)

    with op.batch_alter_table(
        "token_audience_association_table", schema=None
    ) as batch_op:
        batch_op.alter_column("index", existing_type=sa.INTEGER(), nullable=False)

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table(
        "token_audience_association_table", schema=None
    ) as batch_op:
        batch_op.alter_column("index", existing_type=sa.INTEGER(), nullable=True)

    with op.batch_alter_table("membership_association_table", schema=None) as batch_op:
        batch_op.alter_column("index", existing_type=sa.INTEGER(), nullable=True)

    with op.batch_alter_table(
        "client_audience_association_table", schema=None
    ) as batch_op:
        batch_op.alter_column("index", existing_type=sa.INTEGER(), nullable=True)

    # ### end Alembic commands ###
