import asyncio
import aiosqlite
import os

async def clean_protocol_data():
    """清理数据库中协议字段异常的数据"""
    
    db_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'db')
    db_path = os.path.join(db_dir, 'proxies.sqlite')
    
    VALID_PROTOCOLS = ['http', 'https', 'socks4', 'socks5']
    
    async with aiosqlite.connect(db_path) as db:
        # 查询异常的协议数据
        async with db.execute('SELECT ip, port, protocol FROM proxies WHERE protocol NOT IN (?, ?, ?, ?)', VALID_PROTOCOLS) as cursor:
            invalid_proxies = await cursor.fetchall()
        
        if invalid_proxies:
            print(f"发现 {len(invalid_proxies)} 条异常协议数据：")
            for ip, port, protocol in invalid_proxies:
                print(f"  - {ip}:{port} (protocol={protocol})")
        
        # 更新所有不是有效协议类型的记录为 'http'
        cursor = await db.execute('''
            UPDATE proxies
            SET protocol = 'http'
            WHERE protocol NOT IN (?, ?, ?, ?)
        ''', VALID_PROTOCOLS)
        
        updated_count = cursor.rowcount
        await db.commit()
        
        print(f"\n已将 {updated_count} 条记录的协议更新为 'http'")
        
        # 统计修复后的协议分布
        print("\n修复后的协议分布：")
        for protocol in VALID_PROTOCOLS:
            async with db.execute('SELECT COUNT(*) FROM proxies WHERE protocol = ?', (protocol,)) as cursor:
                count = (await cursor.fetchone())[0]
                print(f"  - {protocol}: {count} 条")

if __name__ == "__main__":
    print("=" * 60)
    print("开始清理数据库中的异常协议数据...")
    print("=" * 60)
    asyncio.run(clean_protocol_data())
    print("=" * 60)
    print("清理完成！")
    print("=" * 60)