import asyncio
import sys

from loguru import logger

from core.config import Config
from core.database import DatabaseManager
from core.logger import setup_logger


async def add_extracted_field():
    # 初始化
    config = Config.from_yaml()
    logger = setup_logger(config)
    db_manager = DatabaseManager(config)
    
    try:
        await db_manager.initialize()
        
        # 检查字段是否存在
        async with db_manager.get_connection() as conn:
            # SQLite中查询表结构的方法
            cursor = await conn.execute("PRAGMA table_info(email_accounts)")
            columns = await cursor.fetchall()
            
            # 检查是否已存在extracted字段
            extracted_exists = any(column[1] == 'extracted' for column in columns)
            
            if extracted_exists:
                logger.info("extracted字段已存在，无需添加")
                return
            
            # 添加extracted字段
            logger.info("正在添加extracted字段...")
            await conn.execute("""
                ALTER TABLE email_accounts 
                ADD COLUMN extracted BOOLEAN DEFAULT 0
            """)
            await conn.commit()
            
            logger.success("成功添加extracted字段")
            
            # 初始化字段值
            logger.info("正在初始化extracted字段值...")
            await conn.execute("""
                UPDATE email_accounts
                SET extracted = 0
            """)
            await conn.commit()
            
            logger.success("初始化完成，所有账号的extracted字段已设置为0")
            
    except Exception as e:
        logger.error(f"添加字段时出错: {str(e)}")
    finally:
        await db_manager.cleanup()


if __name__ == "__main__":
    asyncio.run(add_extracted_field())