"""add vertex ai credentials to secrets

Revision ID: f18f1870b3d7
Revises: b635f502f44f
Create Date: 2026-01-16 13:12:35.418855

"""

import sqlalchemy as sa

from alembic import op
from db_models.custom_types import EncryptedJSON

# revision identifiers, used by Alembic.
revision = "f18f1870b3d7"
down_revision = "b635f502f44f"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("secret_storage", sa.Column("project_id", sa.String(), nullable=True))
    op.add_column("secret_storage", sa.Column("region", sa.String(), nullable=True))
    op.add_column(
        "secret_storage",
        sa.Column("vertex_credentials", EncryptedJSON(), nullable=True),
    )
    op.add_column(
        "secret_storage",
        sa.Column("aws_bedrock_credentials", EncryptedJSON(), nullable=True),
    )
    op.alter_column("secret_storage", "value", existing_type=sa.TEXT(), nullable=True)
    # ### end Alembic commands ###


def downgrade() -> None:
    # Remove vertex ai credentials from secret storage on downgrade since it is unsupported and we would
    # lose the credentials needed to run it if we upgraded again
    op.execute(
        "DELETE FROM secret_storage WHERE secret_type = 'model_provider' AND name = 'vertex_ai'",
    )

    # Remove bedrock provider where value (aka the api key field) is null since we will no longer be able to support it
    op.execute(
        "DELETE FROM secret_storage WHERE secret_type = 'model_provider' AND name = 'bedrock' AND value IS NULL",
    )

    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column("secret_storage", "value", existing_type=sa.TEXT(), nullable=False)
    op.drop_column("secret_storage", "aws_bedrock_credentials")
    op.drop_column("secret_storage", "vertex_credentials")
    op.drop_column("secret_storage", "region")
    op.drop_column("secret_storage", "project_id")
    # ### end Alembic commands ###
