"""quick experience-environment: Add system embedding model and llm model, Adjust user

Revision ID: 07acd03b1173
Revises: 69cb68688942
Create Date: 2026-02-03 16:03:43.432989

"""
from typing import Sequence, Union

from alembic import op
from openjiuwen_studio.core.database.migration_utils import (
    table_exists, column_exists, index_exists, unique_constraint_exists
)

import sqlalchemy as sa
from sqlalchemy.dialects import sqlite

# revision identifiers, used by Alembic.
revision: str = '07acd03b1173'
down_revision: Union[str, Sequence[str], None] = '69cb68688942'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    if not table_exists('system_embedding_model'):
        op.create_table('system_embedding_model',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('model_name', sa.String(length=100), nullable=False),
        sa.Column('protocol', sa.String(length=50), nullable=False),
        sa.Column('model_id', sa.String(length=100), nullable=False),
        sa.Column('api_key', sa.Text(), nullable=True),
        sa.Column('api_base', sa.String(length=500), nullable=False),
        sa.Column('max_batch_size', sa.Integer(), nullable=False),
        sa.Column('is_active', sa.Boolean(), server_default='1', nullable=True),
        sa.Column('created_at', sa.DateTime(timezone=True),
                   server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True),
        sa.Column('updated_at', sa.DateTime(timezone=True),
                   server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('model_name')
        )
    with op.batch_alter_table('system_embedding_model', schema=None) as batch_op:
        if not index_exists('system_embedding_model', 'ix_system_embedding_model_id_system_embedding_model'):
            batch_op.create_index('ix_system_embedding_model_id_system_embedding_model', ['id'], unique=False)
        if not index_exists('system_embedding_model',
                          'ix_system_embedding_model_is_active_system_embedding_model'):
            batch_op.create_index(
                'ix_system_embedding_model_is_active_system_embedding_model',
                ['is_active'], unique=False)
        if not index_exists('system_embedding_model',
                          'ix_system_embedding_model_protocol_system_embedding_model'):
            batch_op.create_index(
                'ix_system_embedding_model_protocol_system_embedding_model',
                ['protocol'], unique=False)

    if not table_exists('system_llm_model'):
        op.create_table('system_llm_model',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=100), nullable=False),
        sa.Column('provider', sa.String(length=50), nullable=False),
        sa.Column('model_type', sa.String(length=100), nullable=False),
        sa.Column('description', sa.Text(), nullable=True),
        sa.Column('api_key', sa.Text(), nullable=True),
        sa.Column('base_url', sa.String(length=500), nullable=True),
        sa.Column('is_active', sa.Boolean(), server_default='1', nullable=True),
        sa.Column('parameters', sa.JSON(), nullable=True),
        sa.Column('timeout', sa.Integer(), nullable=True),
        sa.Column('created_at', sa.DateTime(timezone=True),
                   server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True),
        sa.Column('updated_at', sa.DateTime(timezone=True),
                   server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('name')
        )
    with op.batch_alter_table('system_llm_model', schema=None) as batch_op:
        if not index_exists('system_llm_model', 'ix_system_llm_model_id_system_llm_model'):
            batch_op.create_index('ix_system_llm_model_id_system_llm_model', ['id'], unique=False)
        if not index_exists('system_llm_model',
                          'ix_system_llm_model_is_active_system_llm_model'):
            batch_op.create_index(
                'ix_system_llm_model_is_active_system_llm_model',
                ['is_active'], unique=False)
        if not index_exists('system_llm_model', 'ix_system_llm_model_provider_system_llm_model'):
            batch_op.create_index('ix_system_llm_model_provider_system_llm_model', ['provider'], unique=False)

    with op.batch_alter_table('embedding_model_configs', schema=None) as batch_op:
        if not column_exists('embedding_model_configs', 'is_system_model'):
            batch_op.add_column(sa.Column('is_system_model', sa.Boolean(), server_default='0', nullable=True))
        if not column_exists('embedding_model_configs', 'system_model_id'):
            batch_op.add_column(sa.Column('system_model_id', sa.Integer(), nullable=True))
        if not index_exists('embedding_model_configs',
                          'ix_embedding_model_configs_system_model_id_embedding_model_configs'):
            batch_op.create_index(
                'ix_embedding_model_configs_system_model_id_embedding_model_configs',
                ['system_model_id'], unique=False)

    with op.batch_alter_table('model_configs', schema=None) as batch_op:
        if not column_exists('model_configs', 'is_system_model'):
            batch_op.add_column(sa.Column('is_system_model', sa.Boolean(), server_default='0', nullable=True))
        if not column_exists('model_configs', 'system_model_id'):
            batch_op.add_column(sa.Column('system_model_id', sa.Integer(), nullable=True))
        if not index_exists('model_configs', 'ix_model_configs_system_model_id_model_configs'):
            batch_op.create_index('ix_model_configs_system_model_id_model_configs', ['system_model_id'], unique=False)

    with op.batch_alter_table('user', schema=None) as batch_op:
        if not column_exists('user', 'refresh_token'):
            batch_op.add_column(sa.Column('refresh_token', sa.String(length=512), server_default='', nullable=True))
        if unique_constraint_exists('user', 'idx_session_key'):
            batch_op.drop_constraint('idx_session_key', type_='unique')

    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###

    with op.batch_alter_table('user', schema=None) as batch_op:
        if not unique_constraint_exists('user', 'idx_session_key'):
            batch_op.create_unique_constraint('idx_session_key', ['session_key'])
        if column_exists('user', 'refresh_token'):
            batch_op.drop_column('refresh_token')

    with op.batch_alter_table('model_configs', schema=None) as batch_op:
        if index_exists('model_configs', 'ix_model_configs_system_model_id_model_configs'):
            batch_op.drop_index('ix_model_configs_system_model_id_model_configs')
        if column_exists('model_configs', 'system_model_id'):
            batch_op.drop_column('system_model_id')
        if column_exists('model_configs', 'is_system_model'):
            batch_op.drop_column('is_system_model')

    with op.batch_alter_table('embedding_model_configs', schema=None) as batch_op:
        if index_exists('embedding_model_configs',
                          'ix_embedding_model_configs_system_model_id_embedding_model_configs'):
            batch_op.drop_index(
                'ix_embedding_model_configs_system_model_id_embedding_model_configs')
        if column_exists('embedding_model_configs', 'system_model_id'):
            batch_op.drop_column('system_model_id')
        if column_exists('embedding_model_configs', 'is_system_model'):
            batch_op.drop_column('is_system_model')

    with op.batch_alter_table('system_llm_model', schema=None) as batch_op:
        if index_exists('system_llm_model', 'ix_system_llm_model_provider_system_llm_model'):
            batch_op.drop_index('ix_system_llm_model_provider_system_llm_model')
        if index_exists('system_llm_model', 'ix_system_llm_model_is_active_system_llm_model'):
            batch_op.drop_index('ix_system_llm_model_is_active_system_llm_model')
        if index_exists('system_llm_model', 'ix_system_llm_model_id_system_llm_model'):
            batch_op.drop_index('ix_system_llm_model_id_system_llm_model')

    if table_exists('system_llm_model'):
        op.drop_table('system_llm_model')
    with op.batch_alter_table('system_embedding_model', schema=None) as batch_op:
        if index_exists('system_embedding_model', 'ix_system_embedding_model_protocol_system_embedding_model'):
            batch_op.drop_index('ix_system_embedding_model_protocol_system_embedding_model')
        if index_exists('system_embedding_model', 'ix_system_embedding_model_is_active_system_embedding_model'):
            batch_op.drop_index('ix_system_embedding_model_is_active_system_embedding_model')
        if index_exists('system_embedding_model', 'ix_system_embedding_model_id_system_embedding_model'):
            batch_op.drop_index('ix_system_embedding_model_id_system_embedding_model')

    if table_exists('system_embedding_model'):
        op.drop_table('system_embedding_model')
    # ### end Alembic commands ###