"""Store avatar legacy ID in the Contact and Room table

Revision ID: abba1ae0edb3
Revises: 8b993243a536
Create Date: 2024-07-29 15:44:41.557388

"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

from slidge.db.models import Contact, Room

# revision identifiers, used by Alembic.
revision: str = "abba1ae0edb3"
down_revision: Union[str, None] = "8b993243a536"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    conn = op.get_bind()
    room_avatars = conn.execute(
        sa.text(
            "select room.id, avatar.legacy_id from room join avatar on room.avatar_id = avatar.id"
        )
    ).all()
    contact_avatars = conn.execute(
        sa.text(
            "select contact.id, avatar.legacy_id from contact join avatar on contact.avatar_id = avatar.id"
        )
    ).all()
    with op.batch_alter_table("contact", schema=None) as batch_op:
        batch_op.add_column(sa.Column("avatar_legacy_id", sa.String(), nullable=True))

    with op.batch_alter_table("room", schema=None) as batch_op:
        batch_op.add_column(sa.Column("avatar_legacy_id", sa.String(), nullable=True))
        if op.get_bind().engine.name != "postgresql":
            batch_op.create_unique_constraint(
                "uq_room_user_account_id_jid", ["user_account_id", "jid"]
            )
            batch_op.create_unique_constraint(
                "uq_room_user_account_id_legacy_id", ["user_account_id", "legacy_id"]
            )

    for room_pk, avatar_legacy_id in room_avatars:
        conn.execute(
            sa.update(Room)
            .where(Room.id == room_pk)
            .values(avatar_legacy_id=avatar_legacy_id)
        )
    for contact_pk, avatar_legacy_id in contact_avatars:
        conn.execute(
            sa.update(Contact)
            .where(Contact.id == contact_pk)
            .values(avatar_legacy_id=avatar_legacy_id)
        )
    # conn.commit()

    with op.batch_alter_table("avatar", schema=None) as batch_op:
        batch_op.drop_column("legacy_id")


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("room", schema=None) as batch_op:
        batch_op.drop_constraint("uq_room_user_account_id_legacy_id", type_="unique")
        batch_op.drop_constraint("uq_room_user_account_id_jid", type_="unique")
        batch_op.drop_column("avatar_legacy_id")

    with op.batch_alter_table("contact", schema=None) as batch_op:
        batch_op.drop_column("avatar_legacy_id")

    with op.batch_alter_table("avatar", schema=None) as batch_op:
        batch_op.add_column(sa.Column("legacy_id", sa.VARCHAR(), nullable=True))

    # ### end Alembic commands ###
