"""202512051405

Revision ID: 20a742ef1d93
Revises: fc9664b8e4a1
Create Date: 2025-12-05 14:24:33.545865

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy import inspect
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision: str = '20a742ef1d93'
down_revision: Union[str, None] = 'fc9664b8e4a1'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def table_exists(table_name: str) -> bool:
    """检查表是否存在"""
    bind = op.get_bind()
    inspector = inspect(bind)
    return table_name in inspector.get_table_names()


def column_exists(table_name: str, column_name: str) -> bool:
    """检查列是否存在"""
    bind = op.get_bind()
    inspector = inspect(bind)
    if not table_exists(table_name):
        return False
    columns = [col['name'] for col in inspector.get_columns(table_name)]
    return column_name in columns


def index_exists(table_name: str, index_name: str) -> bool:
    """检查索引是否存在"""
    bind = op.get_bind()
    inspector = inspect(bind)
    if not table_exists(table_name):
        return False
    indexes = [idx['name'] for idx in inspector.get_indexes(table_name)]
    return index_name in indexes


def constraint_exists(table_name: str, constraint_name: str) -> bool:
    """检查约束是否存在（外键、唯一约束等）"""
    bind = op.get_bind()
    inspector = inspect(bind)
    if not table_exists(table_name):
        return False
    
    # 检查外键约束
    foreign_keys = [fk['name'] for fk in inspector.get_foreign_keys(table_name) if fk['name']]
    if constraint_name in foreign_keys:
        return True
    
    # 检查唯一约束
    unique_constraints = [uc['name'] for uc in inspector.get_unique_constraints(table_name) if uc['name']]
    if constraint_name in unique_constraints:
        return True
    
    # 检查检查约束
    check_constraints = [cc['name'] for cc in inspector.get_check_constraints(table_name) if cc['name']]
    if constraint_name in check_constraints:
        return True
    
    return False


def trigger_exists(trigger_name: str) -> bool:
    """检查触发器是否存在（PostgreSQL）"""
    bind = op.get_bind()
    result = bind.execute(sa.text(
        "SELECT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = :trigger_name)"
    ), {"trigger_name": trigger_name})
    return result.scalar()


def sequence_exists(sequence_name: str) -> bool:
    """检查序列是否存在（PostgreSQL）"""
    bind = op.get_bind()
    result = bind.execute(sa.text(
        "SELECT EXISTS (SELECT 1 FROM pg_class WHERE relkind = 'S' AND relname = :sequence_name)"
    ), {"sequence_name": sequence_name})
    return result.scalar()


def enum_exists(enum_name: str) -> bool:
    """检查枚举类型是否存在（PostgreSQL）"""
    bind = op.get_bind()
    result = bind.execute(sa.text(
        "SELECT EXISTS (SELECT 1 FROM pg_type WHERE typname = :enum_name)"
    ), {"enum_name": enum_name})
    return result.scalar()


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('data_config', schema=None) as batch_op:
        batch_op.alter_column('llm',
               existing_type=sa.VARCHAR(length=255),
               comment='LLM模型配置ID',
               existing_comment='LS',
               existing_nullable=True)

    with op.batch_alter_table('end_users', schema=None) as batch_op:
        batch_op.alter_column('id',
               existing_type=sa.UUID(),
               server_default=None,
               existing_nullable=False)

    with op.batch_alter_table('knowledges', schema=None) as batch_op:
        batch_op.alter_column('permission_id',
               existing_type=sa.VARCHAR(),
               comment='permission ID:Private|Share|Memory',
               existing_comment='permission ID:Private|Share',
               existing_nullable=True)

    with op.batch_alter_table('memory_increments', schema=None) as batch_op:
        batch_op.alter_column('id',
               existing_type=sa.UUID(),
               server_default=None,
               existing_nullable=False)
        batch_op.alter_column('created_at',
               existing_type=postgresql.TIMESTAMP(precision=0),
               server_default=None,
               existing_nullable=True)
        batch_op.alter_column('updated_at',
               existing_type=postgresql.TIMESTAMP(precision=0),
               server_default=None,
               existing_nullable=True)

    # 为 model_configs 添加 tenant_id 列和相关约束
    if table_exists('model_configs'):
        # 添加列（分步骤处理 NOT NULL 约束）
        if not column_exists('model_configs', 'tenant_id'):
            # 步骤1: 先添加可空列
            with op.batch_alter_table('model_configs', schema=None) as batch_op:
                batch_op.add_column(sa.Column('tenant_id', sa.UUID(), nullable=True, comment='租户ID'))
            
            # 步骤2: 为现有数据填充默认值
            # 注意：这里需要根据实际业务逻辑填充，这里假设获取第一个租户ID
            bind = op.get_bind()
            result = bind.execute(sa.text("SELECT id FROM tenants LIMIT 1"))
            first_tenant = result.scalar()
            
            if first_tenant:
                # 更新现有记录
                bind.execute(
                    sa.text("UPDATE model_configs SET tenant_id = :tenant_id WHERE tenant_id IS NULL"),
                    {"tenant_id": first_tenant}
                )
            
            # 步骤3: 将列改为 NOT NULL
            with op.batch_alter_table('model_configs', schema=None) as batch_op:
                batch_op.alter_column('tenant_id', nullable=False)
        
        # 创建索引
        if not index_exists('model_configs', 'ix_model_configs_tenant_id'):
            op.create_index('ix_model_configs_tenant_id', 'model_configs', ['tenant_id'], unique=False)
        
        # 创建外键
        if not constraint_exists('model_configs', 'model_configs_tenant_id_fkey'):
            op.create_foreign_key('model_configs_tenant_id_fkey', 'model_configs', 'tenants', ['tenant_id'], ['id'])

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    # 回滚 model_configs 的修改
    if table_exists('model_configs'):
        # 删除外键
        if constraint_exists('model_configs', 'model_configs_tenant_id_fkey'):
            op.drop_constraint('model_configs_tenant_id_fkey', 'model_configs', type_='foreignkey')
        
        # 删除索引
        if index_exists('model_configs', 'ix_model_configs_tenant_id'):
            op.drop_index('ix_model_configs_tenant_id', table_name='model_configs')
        
        # 删除列
        if column_exists('model_configs', 'tenant_id'):
            with op.batch_alter_table('model_configs', schema=None) as batch_op:
                batch_op.drop_column('tenant_id')

    with op.batch_alter_table('memory_increments', schema=None) as batch_op:
        batch_op.alter_column('updated_at',
               existing_type=postgresql.TIMESTAMP(precision=0),
               server_default=sa.text('CURRENT_TIMESTAMP'),
               existing_nullable=True)
        batch_op.alter_column('created_at',
               existing_type=postgresql.TIMESTAMP(precision=0),
               server_default=sa.text('CURRENT_TIMESTAMP'),
               existing_nullable=True)
        batch_op.alter_column('id',
               existing_type=sa.UUID(),
               server_default=sa.text('gen_random_uuid()'),
               existing_nullable=False)

    with op.batch_alter_table('knowledges', schema=None) as batch_op:
        batch_op.alter_column('permission_id',
               existing_type=sa.VARCHAR(),
               comment='permission ID:Private|Share',
               existing_comment='permission ID:Private|Share|Memory',
               existing_nullable=True)

    with op.batch_alter_table('end_users', schema=None) as batch_op:
        batch_op.alter_column('id',
               existing_type=sa.UUID(),
               server_default=sa.text('gen_random_uuid()'),
               existing_nullable=False)

    with op.batch_alter_table('data_config', schema=None) as batch_op:
        batch_op.alter_column('llm',
               existing_type=sa.VARCHAR(length=255),
               comment='LS',
               existing_comment='LLM模型配置ID',
               existing_nullable=True)

    # ### end Alembic commands ###