"""fix mysql agent update

Revision ID: 072ac1293a02
Revises: 54be202e347e
Create Date: 2026-02-07 11:16:02.133215

"""
from typing import Sequence, Union

from alembic import op
from openjiuwen_studio.core.database.migration_utils import table_exists, column_exists, index_exists
import sqlalchemy as sa
from sqlalchemy.dialects import mysql
from sqlalchemy import inspect

# revision identifiers, used by Alembic.
revision: str = '072ac1293a02'
down_revision: Union[str, Sequence[str], None] = '54be202e347e'
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('memory_base'):
        op.create_table('memory_base',
        sa.Column('id', sa.BigInteger(), autoincrement=True, nullable=False),
        sa.Column('space_id', sa.String(length=100), nullable=False,
                     comment='空间ID，用于多租户隔离'),
        sa.Column('mdb_id', sa.String(length=100), nullable=False, comment='记忆库ID，唯一标识'),
        sa.Column('name', sa.String(length=100), nullable=False, comment='记忆库名称'),
        sa.Column('description', sa.String(length=100), nullable=False, comment='记忆库名描述'),
        sa.Column('embedding_model_config_id', sa.Integer(), nullable=False,
                     comment='Embedding 模型配置ID'),
        sa.Column('llm_model_config_id', sa.Integer(), nullable=False,
                     comment='LLM 模型配置ID'),
        sa.Column('create_time', sa.BigInteger(), nullable=False, comment='创建时间'),
        sa.Column('update_time', sa.BigInteger(), nullable=False, comment='更新时间'),
        sa.PrimaryKeyConstraint('id'),
        sa.UniqueConstraint('mdb_id', name='uix_mdb_id'),
        comment='记忆库表，存储记忆库基本信息'
        )

    if not index_exists('memory_base', 'idx_space_id'):
        op.create_index('idx_space_id', 'memory_base', ['space_id'], unique=False)
    if not index_exists('memory_base', 'idx_space_mdb'):
        op.create_index('idx_space_mdb', 'memory_base', ['space_id', 'mdb_id'], unique=False)
    if not index_exists('memory_base', 'ix_memory_base_embedding_model_config_id'):
        op.create_index(op.f('ix_memory_base_embedding_model_config_id'),
                     'memory_base', ['embedding_model_config_id'], unique=False)
    if not index_exists('memory_base', 'ix_memory_base_llm_model_config_id'):
        op.create_index(op.f('ix_memory_base_llm_model_config_id'),
                     'memory_base', ['llm_model_config_id'], unique=False)
    if not index_exists('memory_base', 'ix_memory_base_mdb_id'):
        op.create_index(op.f('ix_memory_base_mdb_id'), 'memory_base', ['mdb_id'], unique=False)

    if not column_exists('knowledge_base', 'index_manager_type'):
        op.add_column('knowledge_base',
                   sa.Column('index_manager_type', sa.String(length=200),
                             nullable=True,
                             comment='Whether Chroma or Milvus is used for indexing'))
    if not column_exists('knowledge_base_document', 'index_manager_type'):
        op.add_column('knowledge_base_document',
                   sa.Column('index_manager_type', sa.String(length=200),
                             nullable=True,
                             comment='Whether Chroma or Milvus is used for indexing'))
    if not column_exists('knowledge_base_document', 'index_id'):
        op.add_column('knowledge_base_document',
                   sa.Column('index_id', sa.String(length=200), nullable=True,
                             comment='ES中的索引ID'))
    if not column_exists('knowledge_base_document', 'index_name'):
        op.add_column('knowledge_base_document',
                   sa.Column('index_name', sa.String(length=200), nullable=True,
                             comment='ES索引名称'))
    
    if column_exists('knowledge_base_document', 'es_index_name'):
        op.drop_column('knowledge_base_document', 'es_index_name')
    if column_exists('knowledge_base_document', 'es_index_id'):
        op.drop_column('knowledge_base_document', 'es_index_id')
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    if not column_exists('knowledge_base_document', 'es_index_id'):
        op.add_column('knowledge_base_document',
                   sa.Column('es_index_id',
                             mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=200),
                             nullable=True, comment='ES中的索引ID'))
    if not column_exists('knowledge_base_document', 'es_index_name'):
        op.add_column('knowledge_base_document',
                   sa.Column('es_index_name',
                             mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=200),
                             nullable=True, comment='ES索引名称'))
    
    if column_exists('knowledge_base_document', 'index_name'):
        op.drop_column('knowledge_base_document', 'index_name')
    if column_exists('knowledge_base_document', 'index_id'):
        op.drop_column('knowledge_base_document', 'index_id')
    if column_exists('knowledge_base_document', 'index_manager_type'):
        op.drop_column('knowledge_base_document', 'index_manager_type')
    if column_exists('knowledge_base', 'index_manager_type'):
        op.drop_column('knowledge_base', 'index_manager_type')

    if index_exists('memory_base', 'ix_memory_base_mdb_id'):
        op.drop_index(op.f('ix_memory_base_mdb_id'), table_name='memory_base')
    if index_exists('memory_base', 'ix_memory_base_llm_model_config_id'):
        op.drop_index(op.f('ix_memory_base_llm_model_config_id'), table_name='memory_base')
    if index_exists('memory_base', 'ix_memory_base_embedding_model_config_id'):
        op.drop_index(op.f('ix_memory_base_embedding_model_config_id'), table_name='memory_base')
    if index_exists('memory_base', 'idx_space_mdb'):
        op.drop_index('idx_space_mdb', table_name='memory_base')
    if index_exists('memory_base', 'idx_space_id'):
        op.drop_index('idx_space_id', table_name='memory_base')
    
    if table_exists('memory_base'):
        op.drop_table('memory_base')
    # ### end Alembic commands ###
