from alembic import op
import sqlalchemy as sa

revision = "20260416_release_b6_cleanup"
down_revision = "20260415_release_b1_b2"
branch_labels = None
depends_on = None


def upgrade():
    conn = op.get_bind()

    rows = conn.execute(
        sa.text(
            """
            SELECT c.id, c.tenant_id, c.branch_id, c.partner_id
            FROM clients c
            LEFT JOIN client_associations a
              ON a.client_id = c.id
            WHERE a.id IS NULL
            """
        )
    ).fetchall()

    for row in rows:
        association_type = "firm" if row.tenant_id else "self_service_unassigned"
        created_source = "system_admin" if row.tenant_id else "self_service"

        conn.execute(
            sa.text(
                """
                INSERT INTO client_associations
                (
                    client_id,
                    association_type,
                    firm_tenant_id,
                    consultant_id,
                    partner_user_id,
                    created_source
                )
                VALUES
                (
                    :client_id,
                    :association_type,
                    :firm_tenant_id,
                    :consultant_id,
                    :partner_user_id,
                    :created_source
                )
                """
            ),
            {
                "client_id": row.id,
                "association_type": association_type,
                "firm_tenant_id": row.tenant_id,
                "consultant_id": None,
                "partner_user_id": row.partner_id,
                "created_source": created_source,
            },
        )


def downgrade():
    conn = op.get_bind()

    conn.execute(
        sa.text(
            """
            DELETE FROM client_associations
            WHERE created_source IN ('system_admin', 'self_service')
            """
        )
    )