"""add_warehouse_and_bin_location_to_inventory

Revision ID: e1bfb453fbe9
Revises: j8e9f0a1b2c3
Create Date: 2025-12-25 12:21:24.006548

"""
from collections.abc import Sequence

import sqlalchemy as sa
from sqlalchemy import text

from alembic import op

# revision identifiers, used by Alembic.
revision: str = "e1bfb453fbe9"
down_revision: str | None = "j8e9f0a1b2c3"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def get_column_names(conn, table_name: str) -> set:
    """Get column names for a table (PostgreSQL)."""
    result = conn.execute(text(
        "SELECT column_name FROM information_schema.columns "
        "WHERE table_name = :table AND table_schema = 'public'"
    ), {"table": table_name})
    return {row[0] for row in result.fetchall()}


def get_index_names(conn, table_name: str) -> set:
    """Get index names for a table (PostgreSQL)."""
    result = conn.execute(text(
        "SELECT indexname FROM pg_indexes "
        "WHERE tablename = :table AND schemaname = 'public'"
    ), {"table": table_name})
    return {row[0] for row in result.fetchall()}


def upgrade() -> None:
    conn = op.get_bind()

    # Check if columns already exist (idempotent)
    columns = get_column_names(conn, "inventory")

    if "warehouse" not in columns:
        op.add_column("inventory", sa.Column("warehouse", sa.String(), nullable=False, server_default="strassen"))

    if "bin_location" not in columns:
        op.add_column("inventory", sa.Column("bin_location", sa.String(), nullable=False, server_default=""))

    # Migrate existing data: copy location to bin_location, set default warehouse
    conn.execute(text("""
        UPDATE inventory
        SET bin_location = COALESCE(location, 'UNKNOWN'),
            warehouse = 'strassen'
        WHERE bin_location IS NULL OR bin_location = ''
    """))

    # Create indexes if they don't exist
    existing_indexes = get_index_names(conn, "inventory")

    if "idx_inventory_warehouse_bin" not in existing_indexes:
        op.create_index("idx_inventory_warehouse_bin", "inventory", ["warehouse", "bin_location"], unique=False)
    if "ix_inventory_bin_location" not in existing_indexes:
        op.create_index(op.f("ix_inventory_bin_location"), "inventory", ["bin_location"], unique=False)
    if "ix_inventory_warehouse" not in existing_indexes:
        op.create_index(op.f("ix_inventory_warehouse"), "inventory", ["warehouse"], unique=False)


def downgrade() -> None:
    conn = op.get_bind()

    # Check which indexes exist before dropping
    existing_indexes = get_index_names(conn, "inventory")

    if "ix_inventory_warehouse" in existing_indexes:
        op.drop_index(op.f("ix_inventory_warehouse"), table_name="inventory")
    if "ix_inventory_bin_location" in existing_indexes:
        op.drop_index(op.f("ix_inventory_bin_location"), table_name="inventory")
    if "idx_inventory_warehouse_bin" in existing_indexes:
        op.drop_index("idx_inventory_warehouse_bin", table_name="inventory")

    # Check if columns exist before dropping
    columns = get_column_names(conn, "inventory")

    if "bin_location" in columns:
        op.drop_column("inventory", "bin_location")
    if "warehouse" in columns:
        op.drop_column("inventory", "warehouse")
