"""
数据库迁移脚本: 为sites表添加is_recommended字段
执行方法: python migrations/add_is_recommended.py
"""
import os
import sys

# 添加项目根目录到sys.path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app import create_app
from models import db

def upgrade():
    """添加is_recommended字段"""
    app = create_app(os.getenv('FLASK_ENV', 'development'))

    with app.app_context():
        try:
            # 检查字段是否已存在
            result = db.session.execute(db.text(
                "SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS "
                "WHERE TABLE_NAME = 'sites' AND COLUMN_NAME = 'is_recommended'"
            ))

            if result.fetchone():
                print("字段 is_recommended 已存在，跳过迁移")
                return

            # 添加字段
            print("开始添加 is_recommended 字段...")
            db.session.execute(db.text(
                "ALTER TABLE sites ADD COLUMN is_recommended TINYINT(1) NOT NULL DEFAULT 0 COMMENT '是否推荐'"
            ))
            db.session.commit()
            print("[OK] 成功添加 is_recommended 字段")

        except Exception as e:
            db.session.rollback()
            print(f"[ERROR] 迁移失败: {str(e)}")
            sys.exit(1)

def downgrade():
    """删除is_recommended字段"""
    app = create_app(os.getenv('FLASK_ENV', 'development'))

    with app.app_context():
        try:
            # 删除字段
            print("开始删除 is_recommended 字段...")
            db.session.execute(db.text(
                "ALTER TABLE sites DROP COLUMN is_recommended"
            ))
            db.session.commit()
            print("✓ 成功删除 is_recommended 字段")

        except Exception as e:
            db.session.rollback()
            print(f"✗ 回滚失败: {str(e)}")
            sys.exit(1)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='数据库迁移: 添加is_recommended字段')
    parser.add_argument('--downgrade', action='store_true', help='回滚迁移')
    args = parser.parse_args()

    if args.downgrade:
        downgrade()
    else:
        upgrade()