"""Initialize empty Tenant codes

Revision ID: 867deef0888b
Revises: cab899dbb213
Create Date: 2025-04-03 09:01:20.446536

"""
from alembic import op
import sqlalchemy as sa
import uuid
from sqlalchemy.sql import table, column, select, or_
from sqlalchemy.orm import Session


# revision identifiers, used by Alembic.
revision = '867deef0888b'
down_revision = 'cab899dbb213'
branch_labels = None
depends_on = None


def upgrade():
    # Create a reference to the tenant table
    tenant_table = table('tenant',
                         column('id', sa.Integer),
                         column('code', sa.String(50)),
                         schema='public'  # Assuming the table is in the 'public' schema
                         )

    # Get a connection
    connection = op.get_bind()
    session = Session(bind=connection)

    try:
        # Find all tenants with empty or null code
        # Note the updated select syntax for SQLAlchemy 2.0
        query = select(tenant_table.c.id).where(
            or_(
                tenant_table.c.code == None,
                tenant_table.c.code == ''
            )
        )

        results = connection.execute(query)

        # Update each tenant with a UUID-based code
        for row in results:
            tenant_id = row[0]
            code = f"TENANT-{str(uuid.uuid4())}"

            # Update the tenant record
            update_stmt = tenant_table.update().where(
                tenant_table.c.id == tenant_id
            ).values(
                code=code
            )

            connection.execute(update_stmt)

        # Commit changes
        session.commit()

        # Log how many records were updated
        count_query = select(sa.func.count()).select_from(tenant_table).where(
            tenant_table.c.code.like('TENANT-%')
        )

        updated_count = connection.execute(count_query).scalar()

        print(f"Updated {updated_count} tenant records with UUID-based codes")

    except Exception as e:
        session.rollback()
        print(f"Error updating tenant codes: {str(e)}")
        raise e
    finally:
        session.close()


def downgrade():
    # No downgrade needed for this data migration
    pass