#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据库迁移脚本
为tasks表添加新字段：paused_at, retry_count, last_error, articles_only
"""
import sqlite3
import os
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def migrate_database():
    """执行数据库迁移"""
    db_path = "data/baijiahao.db"
    
    if not os.path.exists(db_path):
        logger.info("数据库文件不存在，无需迁移")
        return
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 检查表是否存在
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='tasks'")
        if not cursor.fetchone():
            logger.info("tasks表不存在，无需迁移")
            conn.close()
            return
        
        # 获取当前表结构
        cursor.execute("PRAGMA table_info(tasks)")
        columns = {row[1]: row for row in cursor.fetchall()}
        
        logger.info("开始数据库迁移...")
        
        # 添加 paused_at 字段
        if 'paused_at' not in columns:
            logger.info("添加 paused_at 字段...")
            cursor.execute("ALTER TABLE tasks ADD COLUMN paused_at TEXT")
            logger.info("✓ paused_at 字段添加成功")
        else:
            logger.info("✓ paused_at 字段已存在")
        
        # 添加 retry_count 字段
        if 'retry_count' not in columns:
            logger.info("添加 retry_count 字段...")
            cursor.execute("ALTER TABLE tasks ADD COLUMN retry_count INTEGER DEFAULT 0")
            logger.info("✓ retry_count 字段添加成功")
        else:
            logger.info("✓ retry_count 字段已存在")
        
        # 添加 last_error 字段
        if 'last_error' not in columns:
            logger.info("添加 last_error 字段...")
            cursor.execute("ALTER TABLE tasks ADD COLUMN last_error TEXT")
            logger.info("✓ last_error 字段添加成功")
        else:
            logger.info("✓ last_error 字段已存在")
        
        # 添加 articles_only 字段
        if 'articles_only' not in columns:
            logger.info("添加 articles_only 字段...")
            cursor.execute("ALTER TABLE tasks ADD COLUMN articles_only INTEGER DEFAULT 1")
            logger.info("✓ articles_only 字段添加成功")
        else:
            logger.info("✓ articles_only 字段已存在")
        
        conn.commit()
        conn.close()
        
        logger.info("=" * 50)
        logger.info("✅ 数据库迁移完成！")
        logger.info("=" * 50)
        
    except Exception as e:
        logger.error(f"❌ 数据库迁移失败: {e}")
        import traceback
        traceback.print_exc()
        raise

if __name__ == "__main__":
    migrate_database()