"""add local_path

Revision ID: 004c73a5c09e
Revises: 0f7dde8b3c4b
Create Date: 2024-11-04 13:47:01.461783

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import sqlmodel
import gpustack


# revision identifiers, used by Alembic.
revision: str = '004c73a5c09e'
down_revision: Union[str, None] = '0f7dde8b3c4b'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

# https://alembic.sqlalchemy.org/en/latest/batch.html#dropping-unnamed-or-named-foreign-key-constraints
naming_convention = {
    "fk":
    "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
}

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('model_instances', schema=None) as batch_op:
        batch_op.add_column(sa.Column('local_path', sqlmodel.sql.sqltypes.AutoString(), nullable=True))

    with op.batch_alter_table('models', schema=None) as batch_op:
        batch_op.add_column(sa.Column('backend_version', sqlmodel.sql.sqltypes.AutoString(), nullable=True))
        batch_op.add_column(sa.Column('local_path', sqlmodel.sql.sqltypes.AutoString(), nullable=True))
        # The default is False, but is not valid in SQLite which uses interger for boolean. Use a generic representation.
        batch_op.add_column(sa.Column('image_only', sa.Boolean(), nullable=False, server_default="0"))
        batch_op.add_column(sa.Column('speech_to_text', sa.Boolean(),
                            nullable=False, server_default="0"))
        batch_op.add_column(sa.Column('text_to_speech', sa.Boolean(),
                            nullable=False, server_default="0"))


    with op.batch_alter_table('model_usages', naming_convention=naming_convention) as batch_op:
        batch_op.alter_column('operation',
                              existing_type=sa.VARCHAR(length=16),
                              type_=sa.Enum('COMPLETION', 'CHAT_COMPLETION', 'EMBEDDING', 'RERANK',
                                            'IMAGE_GENERATION', 'AUDIO_SPEECH', 'AUDIO_TRANSCRIPTION', name='operationenum'),
                              existing_nullable=False)
        batch_op.drop_constraint('fk_model_usages_user_id_users', type_='foreignkey')
        batch_op.drop_constraint('fk_model_usages_model_id_models', type_='foreignkey')
        batch_op.create_foreign_key('fk_model_usages_user_id_users', 'users', ['user_id'], ['id'], ondelete='CASCADE')
        batch_op.create_foreign_key('fk_model_usages_model_id_models', 'models', ['model_id'], ['id'], ondelete='CASCADE')
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('model_usages', naming_convention=naming_convention) as batch_op:
        batch_op.drop_constraint('fk_model_usages_user_id_users', type_='foreignkey')
        batch_op.drop_constraint('fk_model_usages_model_id_models', type_='foreignkey')
        batch_op.create_foreign_key('fk_model_usages_user_id_users', 'users', ['user_id'], ['id'])
        batch_op.create_foreign_key('fk_model_usages_model_id_models', 'models', ['model_id'], ['id'])

    with op.batch_alter_table('models', schema=None) as batch_op:
        batch_op.drop_column('local_path')
        batch_op.drop_column('image_only')
        batch_op.drop_column('text_to_speech')
        batch_op.drop_column('speech_to_text')
        batch_op.drop_column('backend_version')

    with op.batch_alter_table('model_instances', schema=None) as batch_op:
        batch_op.drop_column('local_path')

    # ### end Alembic commands ###
