#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
直接使用 alembic API 执行数据库迁移的脚本

用法:
    python scripts/direct_migration.py [--revision 版本号]

示例:
    # 升级到最新版本
    python scripts/direct_migration.py
    
    # 升级到指定版本
    python scripts/direct_migration.py --revision 001
    
    # 降级到指定版本
    python scripts/direct_migration.py --revision 001 --downgrade
"""

import argparse
import os
import sys
from pathlib import Path
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('migration_direct.log', encoding='utf-8')
    ]
)
logger = logging.getLogger('alembic_direct')

# 将项目根目录添加到 Python 路径
root_dir = Path(__file__).parent.parent
sys.path.insert(0, str(root_dir))

def run_migration(revision=None, downgrade=False, verbose=False):
    """
    执行数据库迁移
    
    Args:
        revision (str): 版本号，为 None 表示升级到最新版本
        downgrade (bool): 是否降级
        verbose (bool): 是否显示详细日志
    """
    try:
        # 导入 alembic 相关模块
        from alembic.config import Config
        from alembic import command
        
        # 设置日志级别
        if verbose:
            logging.getLogger('alembic').setLevel(logging.DEBUG)
        
        # 获取配置文件路径
        config_path = root_dir / "migrations" / "alembic.ini.tmp"
        if not config_path.exists():
            # 如果临时配置文件不存在，使用原始配置文件
            config_path = root_dir / "migrations" / "alembic.ini"
        
        logger.info(f"使用配置文件: {config_path}")
        
        # 创建 alembic 配置对象
        alembic_cfg = Config(str(config_path))
        
        # 设置脚本位置
        script_location = root_dir / "migrations"
        alembic_cfg.set_main_option("script_location", str(script_location))
        
        # 从 config/database.py 导入数据库 URL
        from config.database import DATABASE_URL
        logger.info(f"数据库 URL: {DATABASE_URL}")
        
        # 设置数据库 URL
        alembic_cfg.set_main_option("sqlalchemy.url", DATABASE_URL)
        
        # 执行迁移
        if downgrade:
            logger.info(f"正在降级到版本: {revision or 'base'}")
            command.downgrade(alembic_cfg, revision or "base")
        else:
            logger.info(f"正在升级到版本: {revision or 'head'}")
            command.upgrade(alembic_cfg, revision or "head")
        
        logger.info("迁移完成")
        return True
    except ImportError as e:
        logger.error(f"导入错误: {e}")
        return False
    except Exception as e:
        logger.error(f"执行迁移时发生错误: {e}", exc_info=True)
        return False


def main():
    parser = argparse.ArgumentParser(description="执行数据库迁移")
    parser.add_argument("--revision", help="版本号，为空表示升级到最新版本")
    parser.add_argument("--downgrade", action="store_true", help="是否降级")
    parser.add_argument("--verbose", "-v", action="store_true", help="显示详细日志")
    
    args = parser.parse_args()
    
    success = run_migration(args.revision, args.downgrade, args.verbose)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()