from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import table, column, select, update
from sqlalchemy.orm.session import Session

# revision identifiers, used by Alembic.
revision = '20251107_180000_link_listings_to_nav_ids'
down_revision = '20251107_163500_add_regular_price_and_oe_list_price'
branch_labels = None
depends_on = None

def upgrade():
    # Add new nullable columns first
    op.add_column('listings', sa.Column('top_id', sa.Integer(), nullable=True))
    op.add_column('listings', sa.Column('sub_id', sa.Integer(), nullable=True))

    bind = op.get_bind()
    session = Session(bind=bind)

    nav_tops = sa.table(
        'nav_tops',
        sa.column('id', sa.Integer),
        sa.column('slug', sa.String),
    )
    nav_subs = sa.table(
        'nav_subs',
        sa.column('id', sa.Integer),
        sa.column('slug', sa.String),
        sa.column('top_id', sa.Integer),
    )
    listings = sa.table(
        'listings',
        sa.column('id', sa.Integer),
        sa.column('top_slug', sa.String),
        sa.column('sub_slug', sa.String),
        sa.column('top_id', sa.Integer),
        sa.column('sub_id', sa.Integer),
    )

    # Map top_slug -> top_id
    top_slug_to_id = {
        slug: id_ for id_, slug in session.execute(select(nav_tops.c.id, nav_tops.c.slug))
    }

    sub_slug_to_id = {
        (top_id, slug): id_ for id_, slug, top_id in session.execute(
            select(nav_subs.c.id, nav_subs.c.slug, nav_subs.c.top_id)
        )
    }

    for row in session.execute(select(listings.c.id, listings.c.top_slug, listings.c.sub_slug)):
        listing_id, top_slug, sub_slug = row
        top_id = top_slug_to_id.get(top_slug)
        sub_id = sub_slug_to_id.get((top_id, sub_slug)) if sub_slug else None
        session.execute(
            listings.update()
            .where(listings.c.id == listing_id)
            .values(top_id=top_id, sub_id=sub_id)
        )
    session.commit()

    # Add foreign keys and constraints
    op.create_foreign_key(None, 'listings', 'nav_tops', ['top_id'], ['id'])
    op.create_foreign_key(None, 'listings', 'nav_subs', ['sub_id'], ['id'])
    op.alter_column('listings', 'top_id', nullable=False)

    # Optional: remove old slug fields
    # op.drop_column('listings', 'top_slug')
    # op.drop_column('listings', 'sub_slug')

def downgrade():
    raise NotImplementedError("No downgrade")
