"""add prepared session fields

Revision ID: 053
Revises: 052
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID

revision = "053"
down_revision = "052"
branch_labels = None
depends_on = None


def upgrade():
    # Add prepared_by_id and assigned_to_id for prepared sessions
    op.add_column("sessions", sa.Column("prepared_by_id", UUID(as_uuid=True), nullable=True))
    op.add_column("sessions", sa.Column("assigned_to_id", UUID(as_uuid=True), nullable=True))

    # Make started_at nullable (prepared sessions have started_at = NULL)
    op.alter_column("sessions", "started_at", existing_type=sa.DateTime(timezone=True), nullable=True)

    # Add foreign key constraints
    op.create_foreign_key("fk_sessions_prepared_by_id", "sessions", "users", ["prepared_by_id"], ["id"])
    op.create_foreign_key("fk_sessions_assigned_to_id", "sessions", "users", ["assigned_to_id"], ["id"])

    # Add indexes for efficient querying
    op.create_index("ix_sessions_prepared_by_id", "sessions", ["prepared_by_id"])
    op.create_index("ix_sessions_assigned_to_id", "sessions", ["assigned_to_id"])


def downgrade():
    op.drop_index("ix_sessions_assigned_to_id", table_name="sessions")
    op.drop_index("ix_sessions_prepared_by_id", table_name="sessions")
    op.drop_constraint("fk_sessions_assigned_to_id", "sessions", type_="foreignkey")
    op.drop_constraint("fk_sessions_prepared_by_id", "sessions", type_="foreignkey")
    op.alter_column("sessions", "started_at", existing_type=sa.DateTime(timezone=True), nullable=False)
    op.drop_column("sessions", "assigned_to_id")
    op.drop_column("sessions", "prepared_by_id")