#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库迁移管理器
负责处理Flask-Migrate的初始化、迁移生成和执行
"""

import os
import sys
import subprocess
from pathlib import Path

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

def run_command(command, description):
    """执行命令并处理错误"""
    print(f"🔄 {description}...")
    try:
        result = subprocess.run(
            command, 
            shell=True, 
            check=True, 
            capture_output=True, 
            text=True,
            cwd=project_root
        )
        print(f"✅ {description} 成功")
        if result.stdout:
            print(f"输出: {result.stdout}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ {description} 失败")
        print(f"错误: {e.stderr}")
        return False

def check_migrations_directory():
    """检查migrations目录是否存在"""
    migrations_dir = project_root / "database" / "migrations"
    return migrations_dir.exists()

def check_migration_files():
    """检查是否有迁移文件"""
    versions_dir = project_root / "database" / "migrations" / "versions"
    if not versions_dir.exists():
        return 0
    
    py_files = list(versions_dir.glob("*.py"))
    return len(py_files)

def initialize_migrations():
    """初始化Flask-Migrate"""
    if not check_migrations_directory():
        print("📋 初始化Flask-Migrate...")
        os.environ['FLASK_APP'] = 'app.py'
        return run_command('flask db init', '初始化Flask-Migrate')
    else:
        print("✅ Flask-Migrate已初始化")
        return True

def stamp_current_state():
    """标记当前数据库状态（用于已有数据库）"""
    print("📋 标记当前数据库状态为最新版本...")
    os.environ['FLASK_APP'] = 'app.py'

    # 先检查是否有迁移文件
    migration_count = check_migration_files()

    if migration_count == 0:
        print("⚠️  没有迁移文件，需要先生成初始迁移")
        return False

    # 标记数据库为最新版本（head），这样就不会重复执行已有的迁移
    return run_command('flask db stamp head', '标记数据库状态为最新版本')

def generate_migration(message="Auto migration"):
    """生成迁移文件"""
    print(f"📋 生成迁移文件: {message}")
    os.environ['FLASK_APP'] = 'app.py'
    return run_command(f'flask db migrate -m "{message}"', '生成迁移文件')



def upgrade_database():
    """执行数据库迁移"""
    print("📋 执行数据库迁移...")
    os.environ['FLASK_APP'] = 'app.py'
    return run_command('flask db upgrade', '执行数据库迁移')

def get_current_revision():
    """获取当前迁移版本"""
    print("📋 检查当前迁移状态...")
    os.environ['FLASK_APP'] = 'app.py'
    return run_command('flask db current', '获取当前迁移状态')

def check_alembic_version_table():
    """检查是否存在alembic_version表"""
    try:
        from app import create_app
        from database.models import db
        from sqlalchemy import text

        app = create_app()
        with app.app_context():
            # 尝试查询alembic_version表
            result = db.session.execute(text("SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = DATABASE() AND table_name = 'alembic_version'"))
            count = result.scalar()
            return count > 0
    except Exception as e:
        print(f"⚠️  检查alembic_version表失败: {e}")
        return False



def main():
    """主函数 - 完整的迁移流程"""
    print("🚀 数据库迁移管理器启动")
    print(f"📍 项目根目录: {project_root}")

    # 步骤1: 初始化Flask-Migrate
    if not initialize_migrations():
        print("❌ Flask-Migrate初始化失败")
        return False

    # 步骤2: 检查迁移文件
    migration_count = check_migration_files()
    print(f"📊 发现 {migration_count} 个迁移文件")

    # 步骤3: 检查是否为首次部署（线上已有数据库但无alembic_version表）
    has_alembic_table = check_alembic_version_table()

    if not has_alembic_table and migration_count > 0:
        print("🔍 检测到首次部署到已有数据库，执行stamp操作...")
        # 标记数据库为最新版本，避免重复创建已存在的表
        if not stamp_current_state():
            print("❌ 数据库状态标记失败")
            return False
        print("✅ 数据库状态标记完成")
    else:
        # 步骤4: 获取当前状态
        get_current_revision()

        # 步骤4.5: 执行迁移
        if not upgrade_database():
            print("❌ 数据库迁移失败")
            return False

    print("🎉 数据库迁移完成")
    return True

if __name__ == '__main__':
    success = main()
    sys.exit(0 if success else 1)
