"""Add provider multi credential support

Revision ID: e8446f481c1e
Revises: 8bcc02c9bd07
Create Date: 2025-08-09 15:53:54.341341

"""
from alembic import op, context
from libs.uuid_utils import uuidv7
import models as models
import sqlalchemy as sa
from sqlalchemy.sql import table, column


def _is_pg(conn):
    return conn.dialect.name == "postgresql"

# revision identifiers, used by Alembic.
revision = 'e8446f481c1e'
down_revision = 'fa8b0fa6f407'
branch_labels = None
depends_on = None


def upgrade():
    # Create provider_credentials table
    conn = op.get_bind()
    
    if _is_pg(conn):
        op.create_table('provider_credentials',
        sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuidv7()'), nullable=False),
        sa.Column('tenant_id', models.types.StringUUID(), nullable=False),
        sa.Column('provider_name', sa.String(length=255), nullable=False),
        sa.Column('credential_name', sa.String(length=255), nullable=False),
        sa.Column('encrypted_config', sa.Text(), nullable=False),
        sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False),
        sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP'), nullable=False),
        sa.PrimaryKeyConstraint('id', name='provider_credential_pkey')
        )
    else:
        op.create_table('provider_credentials',
        sa.Column('id', models.types.StringUUID(), nullable=False),
        sa.Column('tenant_id', models.types.StringUUID(), nullable=False),
        sa.Column('provider_name', sa.String(length=255), nullable=False),
        sa.Column('credential_name', sa.String(length=255), nullable=False),
        sa.Column('encrypted_config', models.types.LongText(), nullable=False),
        sa.Column('created_at', sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False),
        sa.Column('updated_at', sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False),
        sa.PrimaryKeyConstraint('id', name='provider_credential_pkey')
        )

    # Create index for provider_credentials
    with op.batch_alter_table('provider_credentials', schema=None) as batch_op:
        batch_op.create_index('provider_credential_tenant_provider_idx', ['tenant_id', 'provider_name'], unique=False)

    # Add credential_id to providers table
    with op.batch_alter_table('providers', schema=None) as batch_op:
        batch_op.add_column(sa.Column('credential_id', models.types.StringUUID(), nullable=True))

    # Add credential_id to load_balancing_model_configs table
    with op.batch_alter_table('load_balancing_model_configs', schema=None) as batch_op:
        batch_op.add_column(sa.Column('credential_id', models.types.StringUUID(), nullable=True))

    if not context.is_offline_mode():
        migrate_existing_providers_data()
    else:
        op.execute(
            '-- [IMPORTANT] Data migration skipped!!!\n'
            "-- You should manually run data migration function `migrate_existing_providers_data`\n"
            f"-- inside file {__file__}\n"
            "-- Please review the migration script carefully!"
        )

    # Remove encrypted_config column from providers table after migration
    with op.batch_alter_table('providers', schema=None) as batch_op:
        batch_op.drop_column('encrypted_config')


def migrate_existing_providers_data():
    """migrate providers table data to provider_credentials"""
    conn = op.get_bind()
    # Define table structure for data manipulation
    if _is_pg(conn):
        providers_table = table('providers',
            column('id', models.types.StringUUID()),
            column('tenant_id', models.types.StringUUID()),
            column('provider_name', sa.String()),
            column('encrypted_config', sa.Text()),
            column('created_at', sa.DateTime()),
            column('updated_at', sa.DateTime()),
            column('credential_id', models.types.StringUUID()),
        )
    else:
        providers_table = table('providers',
            column('id', models.types.StringUUID()),
            column('tenant_id', models.types.StringUUID()),
            column('provider_name', sa.String()),
            column('encrypted_config', models.types.LongText()),
            column('created_at', sa.DateTime()),
            column('updated_at', sa.DateTime()),
            column('credential_id', models.types.StringUUID()),
        )

    if _is_pg(conn):
        provider_credential_table = table('provider_credentials',
            column('id', models.types.StringUUID()),
            column('tenant_id', models.types.StringUUID()),
            column('provider_name', sa.String()),
            column('credential_name', sa.String()),
            column('encrypted_config', sa.Text()),
            column('created_at', sa.DateTime()),
            column('updated_at', sa.DateTime())
        )
    else:
        provider_credential_table = table('provider_credentials',
            column('id', models.types.StringUUID()),
            column('tenant_id', models.types.StringUUID()),
            column('provider_name', sa.String()),
            column('credential_name', sa.String()),
            column('encrypted_config', models.types.LongText()),
            column('created_at', sa.DateTime()),
            column('updated_at', sa.DateTime())
        )

    # Get database connection
    conn = op.get_bind()

    # Query all existing providers data
    existing_providers = conn.execute(
        sa.select(providers_table.c.id, providers_table.c.tenant_id,
                 providers_table.c.provider_name, providers_table.c.encrypted_config,
                 providers_table.c.created_at, providers_table.c.updated_at)
        .where(providers_table.c.encrypted_config.isnot(None))
    ).fetchall()

    # Iterate through each provider and insert into provider_credentials
    for provider in existing_providers:
        credential_id = str(uuidv7())
        if not provider.encrypted_config or provider.encrypted_config.strip() == '':
            continue

        # Insert into provider_credentials table
        conn.execute(
            provider_credential_table.insert().values(
                id=credential_id,
                tenant_id=provider.tenant_id,
                provider_name=provider.provider_name,
                credential_name='API_KEY1',  # Use a default name
                encrypted_config=provider.encrypted_config,
                created_at=provider.created_at,
                updated_at=provider.updated_at
            )
        )

        # Update original providers table, set credential_id
        conn.execute(
            providers_table.update()
            .where(providers_table.c.id == provider.id)
            .values(
                credential_id=credential_id,
            )
        )

def downgrade():
    # Re-add encrypted_config column to providers table
    conn = op.get_bind()
    
    if _is_pg(conn):
        with op.batch_alter_table('providers', schema=None) as batch_op:
            batch_op.add_column(sa.Column('encrypted_config', sa.Text(), nullable=True))
    else:
        with op.batch_alter_table('providers', schema=None) as batch_op:
            batch_op.add_column(sa.Column('encrypted_config', models.types.LongText(), nullable=True))

    # Migrate data back from provider_credentials to providers

    if not context.is_offline_mode():
        migrate_data_back_to_providers()
    else:
        op.execute(
            '-- [IMPORTANT] Data migration skipped!!!\n'
            "-- You should manually run data migration function `migrate_data_back_to_providers`\n"
            f"-- inside file {__file__}\n"
            "-- Please review the migration script carefully!"
        )

    # Remove credential_id columns
    with op.batch_alter_table('load_balancing_model_configs', schema=None) as batch_op:
        batch_op.drop_column('credential_id')

    with op.batch_alter_table('providers', schema=None) as batch_op:
        batch_op.drop_column('credential_id')

    # Drop provider_credentials table
    op.drop_table('provider_credentials')


def migrate_data_back_to_providers():
    """Migrate data back from provider_credentials to providers table for downgrade"""

    # Define table structure for data manipulation
    providers_table = table('providers',
        column('id', models.types.StringUUID()),
        column('tenant_id', models.types.StringUUID()),
        column('provider_name', sa.String()),
        column('encrypted_config', sa.Text()),
        column('credential_id', models.types.StringUUID()),
    )

    provider_credential_table = table('provider_credentials',
        column('id', models.types.StringUUID()),
        column('tenant_id', models.types.StringUUID()),
        column('provider_name', sa.String()),
        column('credential_name', sa.String()),
        column('encrypted_config', sa.Text()),
    )

    # Get database connection
    conn = op.get_bind()

    # Query providers that have credential_id
    providers_with_credentials = conn.execute(
        sa.select(providers_table.c.id, providers_table.c.credential_id)
        .where(providers_table.c.credential_id.isnot(None))
    ).fetchall()

    # For each provider, get the credential data and update providers table
    for provider in providers_with_credentials:
        credential = conn.execute(
            sa.select(provider_credential_table.c.encrypted_config)
            .where(provider_credential_table.c.id == provider.credential_id)
        ).fetchone()

        if credential:
            # Update providers table with encrypted_config from credential
            conn.execute(
                providers_table.update()
                .where(providers_table.c.id == provider.id)
                .values(encrypted_config=credential.encrypted_config)
            )
