"""drop auth_type from non source connectors

Revision ID: 51f41d4cb9db
Revises: 22fcfb994985
Create Date: 2025-09-15 12:49:15.161797

"""

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = "51f41d4cb9db"
down_revision = "22fcfb994985"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###

    # Drop auth_type from non-source connectors
    op.drop_column("auth_provider", "auth_type")
    op.drop_column("destination", "auth_type")
    op.drop_column("embedding_model", "auth_type")

    # Add new columns as nullable first
    op.add_column(
        "integration_credential", sa.Column("authentication_method", sa.String(), nullable=True)
    )
    op.add_column("integration_credential", sa.Column("oauth_type", sa.String(), nullable=True))

    # Migrate data from auth_type to new columns
    # Map old AuthType enum values to new AuthenticationMethod and OAuthType
    op.execute(
        """
        UPDATE integration_credential
        SET authentication_method = CASE
            WHEN auth_type = 'oauth2' THEN 'oauth_browser'
            WHEN auth_type = 'oauth2_with_refresh' THEN 'oauth_browser'
            WHEN auth_type = 'oauth2_with_refresh_rotating' THEN 'oauth_browser'
            WHEN auth_type = 'trello_auth' THEN 'oauth_browser'
            WHEN auth_type = 'api_key' THEN 'direct'
            WHEN auth_type = 'config_class' THEN 'direct'
            WHEN auth_type = 'native_functionality' THEN 'direct'
            WHEN auth_type = 'none' THEN 'direct'
            ELSE 'direct'  -- Default fallback
        END,
        oauth_type = CASE
            WHEN auth_type = 'oauth2' THEN 'access_only'
            WHEN auth_type = 'oauth2_with_refresh' THEN 'with_refresh'
            WHEN auth_type = 'oauth2_with_refresh_rotating' THEN 'with_rotating_refresh'
            WHEN auth_type = 'trello_auth' THEN 'access_only'  -- Trello uses simple OAuth
            ELSE NULL  -- Non-OAuth methods
        END
    """
    )

    # Now make authentication_method NOT NULL after data is migrated
    op.alter_column("integration_credential", "authentication_method", nullable=False)

    # Drop the old auth_type column
    op.drop_column("integration_credential", "auth_type")
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###

    # Add auth_type column back as nullable first
    op.add_column(
        "integration_credential",
        sa.Column("auth_type", sa.VARCHAR(), autoincrement=False, nullable=True),
    )

    # Migrate data back from new columns to auth_type
    op.execute(
        """
        UPDATE integration_credential
        SET auth_type = CASE
            WHEN authentication_method = 'oauth_browser' AND oauth_type = 'access_only' THEN 'oauth2'
            WHEN authentication_method = 'oauth_browser' AND oauth_type = 'with_refresh' THEN 'oauth2_with_refresh'
            WHEN authentication_method = 'oauth_browser' AND oauth_type = 'with_rotating_refresh' THEN 'oauth2_with_refresh_rotating'
            WHEN authentication_method = 'oauth_browser' AND oauth_type IS NULL THEN 'oauth2'  -- Default OAuth
            WHEN authentication_method = 'direct' AND auth_config_class LIKE '%ApiKey%' THEN 'api_key'
            WHEN authentication_method = 'direct' AND auth_config_class IS NOT NULL THEN 'config_class'
            WHEN authentication_method = 'direct' THEN 'api_key'  -- Default direct auth
            WHEN authentication_method = 'oauth_token' THEN 'oauth2'
            WHEN authentication_method = 'oauth_byoc' THEN 'oauth2'
            WHEN authentication_method = 'auth_provider' THEN 'native_functionality'
            ELSE 'none'  -- Fallback
        END
    """
    )

    # Make auth_type NOT NULL after migration
    op.alter_column("integration_credential", "auth_type", nullable=False)

    # Drop the new columns
    op.drop_column("integration_credential", "oauth_type")
    op.drop_column("integration_credential", "authentication_method")

    # Re-add auth_type to other tables
    op.add_column(
        "embedding_model",
        sa.Column(
            "auth_type",
            postgresql.ENUM(
                "oauth2",
                "oauth2_with_refresh",
                "oauth2_with_refresh_rotating",
                "trello_auth",
                "api_key",
                "native_functionality",
                "config_class",
                "none",
                name="authtype",
            ),
            autoincrement=False,
            nullable=True,
        ),
    )
    op.add_column(
        "destination",
        sa.Column(
            "auth_type",
            postgresql.ENUM(
                "oauth2",
                "oauth2_with_refresh",
                "oauth2_with_refresh_rotating",
                "trello_auth",
                "api_key",
                "native_functionality",
                "config_class",
                "none",
                name="authtype",
            ),
            autoincrement=False,
            nullable=True,
        ),
    )
    op.add_column(
        "auth_provider",
        sa.Column(
            "auth_type",
            postgresql.ENUM(
                "oauth2",
                "oauth2_with_refresh",
                "oauth2_with_refresh_rotating",
                "trello_auth",
                "api_key",
                "native_functionality",
                "config_class",
                "none",
                name="authtype",
            ),
            autoincrement=False,
            nullable=False,
        ),
    )
    # ### end Alembic commands ###
