from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import table, column, select

# revision identifiers, used by Alembic.
revision = 'a9e886c56482'
down_revision = None
branch_labels = None
depends_on = None

def upgrade():
    # Create room table
    op.create_table('room',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=128), nullable=False),
        sa.Column('title', sa.String(length=128), nullable=True),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('name')
    )

    # Add room_id column to message table
    op.add_column('message', sa.Column('room_id', sa.Integer(), nullable=True))

    # Temporary table objects
    message_table = table('message',
        column('id', sa.Integer),
        column('username', sa.String),
        column('content', sa.String),
        column('room', sa.String),
        column('room_id', sa.Integer),
    )
    room_table = table('room',
        column('id', sa.Integer),
        column('name', sa.String)
    )

    # Execution context
    conn = op.get_bind()

    # Insert distinct rooms into room table and create mapping
    distinct_rooms = conn.execute(select(message_table.c.room).distinct())
    room_name_to_id = {}
    for room_name, in distinct_rooms:
        conn.execute(room_table.insert().values(name=room_name))
        room_id = conn.execute(select(room_table.c.id).where(room_table.c.name == room_name)).scalar()
        room_name_to_id[room_name] = room_id

    # Update message table with room_id
    for room_name, room_id in room_name_to_id.items():
        conn.execute(message_table.update().where(message_table.c.room == room_name).values(room_id=room_id))

    # Create new_message table
    new_message_table = op.create_table('new_message',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('username', sa.String(length=128), nullable=False),
        sa.Column('content', sa.String(length=1024), nullable=False),
        sa.Column('room_id', sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(['room_id'], ['room.id']),
        sa.PrimaryKeyConstraint('id')
    )

    # Copy data from old message table to new_message table
    old_messages = conn.execute(sa.select(message_table)).fetchall()
    for old_message in old_messages:
        conn.execute(new_message_table.insert().values(
            id=old_message.id,
            username=old_message.username,
            content=old_message.content,
            room_id=old_message.room_id
        ))

    # Drop old message table and rename new_message to message
    op.drop_table('message')
    op.rename_table('new_message', 'message')

def downgrade():
    # Recreate old_message table with 'room' column
    old_message_table = op.create_table('old_message',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('username', sa.String(length=128), nullable=False),
        sa.Column('content', sa.String(length=1024), nullable=False),
        sa.Column('room', sa.String(length=128), nullable=False),
        sa.PrimaryKeyConstraint('id')
    )

    # Copy data back from message to old_message
    message_table = table('message',
        column('id', sa.Integer),
        column('username', sa.String),
        column('content', sa.String),
        column('room_id', sa.Integer)
    )
    room_table = table('room',
        column('id', sa.Integer),
        column('name', sa.String)
    )

    conn = op.get_bind()
    messages = conn.execute(select(message_table)).fetchall()
    for message in messages:
        room_name = conn.execute(select(room_table.c.name).where(room_table.c.id == message.room_id)).scalar()
        conn.execute(old_message_table.insert().values(
            id=message.id,
            username=message.username,
            content=message.content,
            room=room_name
        ))

    # Drop current message table and rename old_message to message
    op.drop_table('message')
    op.rename_table('old_message', 'message')
    op.drop_table('room')