"""Add restaurant_id to indirectly-scoped tables and backfill.

Revision ID: 0022_add_restaurant_id_to_indirect_tables
Revises: 0021_add_fts_gin_indexes
Create Date: 2026-03-12
"""

from collections.abc import Sequence

import sqlalchemy as sa

from alembic import op

revision: str = "0022_add_restaurant_id_to_indirect_tables"
down_revision: str | None = "0021_add_fts_gin_indexes"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    # 1) Add nullable restaurant_id columns
    op.add_column("chair", sa.Column("restaurant_id", sa.String(), nullable=True))
    op.create_index("ix_chair_restaurant_id", "chair", ["restaurant_id"], unique=False)
    op.create_foreign_key(
        "fk_chair_restaurant_id_restaurant",
        "chair",
        "restaurant",
        ["restaurant_id"],
        ["id"],
    )

    op.add_column("order_item", sa.Column("restaurant_id", sa.String(), nullable=True))
    op.create_index("ix_order_item_restaurant_id", "order_item", ["restaurant_id"], unique=False)
    op.create_foreign_key(
        "fk_order_item_restaurant_id_restaurant",
        "order_item",
        "restaurant",
        ["restaurant_id"],
        ["id"],
    )

    op.add_column("message", sa.Column("restaurant_id", sa.String(), nullable=True))
    op.create_index("ix_message_restaurant_id", "message", ["restaurant_id"], unique=False)
    op.create_foreign_key(
        "fk_message_restaurant_id_restaurant",
        "message",
        "restaurant",
        ["restaurant_id"],
        ["id"],
    )

    op.add_column("service_block_zones", sa.Column("restaurant_id", sa.String(), nullable=True))
    op.create_index(
        "ix_service_block_zones_restaurant_id",
        "service_block_zones",
        ["restaurant_id"],
        unique=False,
    )
    op.create_foreign_key(
        "fk_service_block_zones_restaurant_id_restaurant",
        "service_block_zones",
        "restaurant",
        ["restaurant_id"],
        ["id"],
    )

    op.add_column("combined_chair_config", sa.Column("restaurant_id", sa.String(), nullable=True))
    op.create_index(
        "ix_combined_chair_config_restaurant_id",
        "combined_chair_config",
        ["restaurant_id"],
        unique=False,
    )
    op.create_foreign_key(
        "fk_combined_chair_config_restaurant_id_restaurant",
        "combined_chair_config",
        "restaurant",
        ["restaurant_id"],
        ["id"],
    )

    # 2) Backfill from parent tables
    op.execute(
        """
        UPDATE chair c
        SET restaurant_id = ft.restaurant_id
        FROM floor_table ft
        WHERE c.table_id = ft.id AND c.restaurant_id IS NULL
        """
    )

    op.execute(
        """
        UPDATE order_item oi
        SET restaurant_id = o.restaurant_id
        FROM "order" o
        WHERE oi.order_id = o.id AND oi.restaurant_id IS NULL
        """
    )

    op.execute(
        """
        UPDATE message m
        SET restaurant_id = conv.restaurant_id
        FROM conversation conv
        WHERE m.conversation_id = conv.id AND m.restaurant_id IS NULL
        """
    )

    op.execute(
        """
        UPDATE service_block_zones sbz
        SET restaurant_id = sb.restaurant_id
        FROM service_block sb
        WHERE sbz.service_block_id = sb.id AND sbz.restaurant_id IS NULL
        """
    )

    op.execute(
        """
        UPDATE combined_chair_config ccc
        SET restaurant_id = tc.restaurant_id
        FROM table_combination tc
        WHERE ccc.combination_id = tc.id AND ccc.restaurant_id IS NULL
        """
    )

    # 3) Enforce NOT NULL after backfill
    op.alter_column("chair", "restaurant_id", nullable=False)
    op.alter_column("order_item", "restaurant_id", nullable=False)
    op.alter_column("message", "restaurant_id", nullable=False)
    op.alter_column("service_block_zones", "restaurant_id", nullable=False)
    op.alter_column("combined_chair_config", "restaurant_id", nullable=False)


def downgrade() -> None:
    # Reverse NOT NULL, drop FKs and indexes, then drop columns
    for table, ix, fk in [
        (
            "combined_chair_config",
            "ix_combined_chair_config_restaurant_id",
            "fk_combined_chair_config_restaurant_id_restaurant",
        ),
        (
            "service_block_zones",
            "ix_service_block_zones_restaurant_id",
            "fk_service_block_zones_restaurant_id_restaurant",
        ),
        ("message", "ix_message_restaurant_id", "fk_message_restaurant_id_restaurant"),
        ("order_item", "ix_order_item_restaurant_id", "fk_order_item_restaurant_id_restaurant"),
        ("chair", "ix_chair_restaurant_id", "fk_chair_restaurant_id_restaurant"),
    ]:
        from contextlib import suppress

        with suppress(Exception):
            op.alter_column(table, "restaurant_id", nullable=True)
        op.drop_constraint(fk, table_name=table, type_="foreignkey")
        op.drop_index(ix, table_name=table)
        op.drop_column(table, "restaurant_id")