import os
import sys
import argparse
from alembic.config import Config
from alembic import command

def create_migration(message=None):
    """
    创建数据库迁移
    
    Args:
        message: 迁移消息，描述此次迁移的内容
    """
    # 获取Alembic配置
    alembic_cfg = Config("alembic.ini")
    
    # 创建迁移
    if message:
        command.revision(alembic_cfg, message=message, autogenerate=True)
    else:
        command.revision(alembic_cfg, message="自动生成的迁移", autogenerate=True)
        
def upgrade_database():
    """更新数据库到最新版本"""
    alembic_cfg = Config("alembic.ini")
    command.upgrade(alembic_cfg, "head")
    
def downgrade_database(revision):
    """回滚数据库到指定版本"""
    alembic_cfg = Config("alembic.ini")
    command.downgrade(alembic_cfg, revision)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="数据库迁移工具")
    subparsers = parser.add_subparsers(dest="command", help="命令")
    
    # 创建迁移命令
    create_parser = subparsers.add_parser("create", help="创建数据库迁移")
    create_parser.add_argument("-m", "--message", help="迁移描述消息")
    
    # 更新数据库命令
    upgrade_parser = subparsers.add_parser("upgrade", help="更新数据库到最新版本")
    
    # 回滚数据库命令
    downgrade_parser = subparsers.add_parser("downgrade", help="回滚数据库到指定版本")
    downgrade_parser.add_argument("revision", help="目标版本，例如：-1表示回滚一个版本")
    
    args = parser.parse_args()
    
    if args.command == "create":
        create_migration(args.message)
        print("迁移文件已创建，请查看 alembic/versions 目录")
    elif args.command == "upgrade":
        upgrade_database()
        print("数据库已更新到最新版本")
    elif args.command == "downgrade":
        downgrade_database(args.revision)
        print(f"数据库已回滚到版本: {args.revision}")
    else:
        parser.print_help()