"""
数据库迁移脚本 - 为Sites表添加news_keywords字段
版本：v2.3.0
日期：2025-12-31
"""
import pymysql
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

def migrate():
    """执行数据库迁移"""
    # 数据库配置
    db_config = {
        'host': os.environ.get('DB_HOST', 'localhost'),
        'port': int(os.environ.get('DB_PORT', 3306)),
        'user': os.environ.get('DB_USER', 'root'),
        'password': os.environ.get('DB_PASSWORD', ''),
        'database': os.environ.get('DB_NAME', 'ai_nav'),
        'charset': 'utf8mb4'
    }

    try:
        # 连接数据库
        connection = pymysql.connect(**db_config)
        cursor = connection.cursor()

        print("=" * 60)
        print("开始执行数据库迁移 v2.3.0")
        print("=" * 60)

        # 检查字段是否已存在
        cursor.execute("""
            SELECT COLUMN_NAME
            FROM INFORMATION_SCHEMA.COLUMNS
            WHERE TABLE_SCHEMA = %s
            AND TABLE_NAME = 'sites'
            AND COLUMN_NAME = 'news_keywords'
        """, (db_config['database'],))

        existing_columns = [row[0] for row in cursor.fetchall()]

        # 添加 news_keywords 字段
        if 'news_keywords' not in existing_columns:
            print("\n1. 添加 news_keywords 字段...")
            cursor.execute("""
                ALTER TABLE sites
                ADD COLUMN news_keywords VARCHAR(200)
                COMMENT '新闻获取关键词（用于精准匹配相关新闻）'
                AFTER features
            """)
            print(">>> news_keywords 字段添加成功")
        else:
            print("\n1. news_keywords 字段已存在，跳过")

        # 可选：为现有网站设置默认关键词（使用网站名称）
        if 'news_keywords' not in existing_columns:
            print("\n2. 为现有网站设置默认关键词（使用网站名称）...")
            cursor.execute("""
                UPDATE sites
                SET news_keywords = name
                WHERE news_keywords IS NULL OR news_keywords = ''
            """)
            affected_rows = cursor.rowcount
            print(f">>> 已更新 {affected_rows} 个网站的默认关键词")

        # 提交事务
        connection.commit()

        print("\n" + "=" * 60)
        print(">>> 数据库迁移完成！")
        print("=" * 60)

        # 显示表结构
        print("\n当前 sites 表结构：")
        cursor.execute("DESCRIBE sites")
        for row in cursor.fetchall():
            print(f"  - {row[0]}: {row[1]} {row[2]}")

        # 显示示例数据
        print("\n示例数据（前5条）：")
        cursor.execute("""
            SELECT id, name, news_keywords
            FROM sites
            LIMIT 5
        """)
        for row in cursor.fetchall():
            print(f"  ID: {row[0]}, 名称: {row[1]}, 关键词: {row[2] or '(空)'}")

    except Exception as e:
        print(f"\n>>> 迁移失败：{str(e)}")
        if 'connection' in locals():
            connection.rollback()
        raise

    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'connection' in locals():
            connection.close()
            print("\n数据库连接已关闭")


if __name__ == '__main__':
    migrate()