"""fix sqlite agent update

Revision ID: 8f4846812221
Revises: 07acd03b1173
Create Date: 2026-02-07 11:50:21.774021

"""
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 import inspect


# revision identifiers, used by Alembic.
revision: str = '8f4846812221'
down_revision: Union[str, Sequence[str], None] = '07acd03b1173'
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.Integer(), 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='记忆库表，存储记忆库基本信息'
        )

    with op.batch_alter_table('memory_base', schema=None) as batch_op:
        if not index_exists('memory_base', 'idx_space_id_memory_base'):
            batch_op.create_index('idx_space_id_memory_base', ['space_id'], unique=False)
        if not index_exists('memory_base', 'idx_space_mdb_memory_base'):
            batch_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_memory_base'):
            batch_op.create_index('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_memory_base'):
            batch_op.create_index('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_memory_base'):
            batch_op.create_index('ix_memory_base_mdb_id_memory_base', ['mdb_id'], unique=False)

    if table_exists('knowledge_base'):
        with op.batch_alter_table('knowledge_base', schema=None) as batch_op:
            if not column_exists('knowledge_base', 'index_manager_type'):
                batch_op.add_column(sa.Column('index_manager_type',
                                                 sa.String(length=200), nullable=True,
                                                 comment='Whether Chroma or Milvus is used for indexing'))

    if table_exists('knowledge_base_document'):
        with op.batch_alter_table('knowledge_base_document', schema=None) as batch_op:
            if not column_exists('knowledge_base_document', 'index_manager_type'):
                batch_op.add_column(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'):
                batch_op.add_column(sa.Column('index_id', sa.String(length=200),
                                                 nullable=True, comment='ES中的索引ID'))
            if not column_exists('knowledge_base_document', 'index_name'):
                batch_op.add_column(sa.Column('index_name', sa.String(length=200),
                                                 nullable=True, comment='ES索引名称'))
            if column_exists('knowledge_base_document', 'es_index_name'):
                batch_op.drop_column('es_index_name')
            if column_exists('knowledge_base_document', 'es_index_id'):
                batch_op.drop_column('es_index_id')

    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    if table_exists('knowledge_base_document'):
        with op.batch_alter_table('knowledge_base_document', schema=None) as batch_op:
            if not column_exists('knowledge_base_document', 'es_index_id'):
                batch_op.add_column(sa.Column('es_index_id', sa.VARCHAR(length=200), nullable=True))
            if not column_exists('knowledge_base_document', 'es_index_name'):
                batch_op.add_column(sa.Column('es_index_name', sa.VARCHAR(length=200), nullable=True))
            if column_exists('knowledge_base_document', 'index_name'):
                batch_op.drop_column('index_name')
            if column_exists('knowledge_base_document', 'index_id'):
                batch_op.drop_column('index_id')
            if column_exists('knowledge_base_document', 'index_manager_type'):
                batch_op.drop_column('index_manager_type')

    if table_exists('knowledge_base'):
        with op.batch_alter_table('knowledge_base', schema=None) as batch_op:
            if column_exists('knowledge_base', 'index_manager_type'):
                batch_op.drop_column('index_manager_type')

    if table_exists('memory_base'):
        with op.batch_alter_table('memory_base', schema=None) as batch_op:
            if index_exists('memory_base', 'ix_memory_base_mdb_id_memory_base'):
                batch_op.drop_index('ix_memory_base_mdb_id_memory_base')
            if index_exists('memory_base', 'ix_memory_base_llm_model_config_id_memory_base'):
                batch_op.drop_index('ix_memory_base_llm_model_config_id_memory_base')
            if index_exists('memory_base', 'ix_memory_base_embedding_model_config_id_memory_base'):
                batch_op.drop_index('ix_memory_base_embedding_model_config_id_memory_base')
            if index_exists('memory_base', 'idx_space_mdb_memory_base'):
                batch_op.drop_index('idx_space_mdb_memory_base')
            if index_exists('memory_base', 'idx_space_id_memory_base'):
                batch_op.drop_index('idx_space_id_memory_base')

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