#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
为query_history表添加user_id字段，并设置为users表id字段的外键
"""

import sqlite3
import logging

# 配置日志
handlers = [logging.FileHandler('db_update.log'), logging.StreamHandler()]
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=handlers
)
logger = logging.getLogger(__name__)

# 数据库文件路径
db_path = 'car_info.db'


def main():
    """
    主函数：连接数据库并为query_history表添加user_id字段
    """
    logger.info("开始为query_history表添加user_id字段")
    
    try:
        # 连接数据库
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 检查query_history表是否已经存在user_id字段
        cursor.execute("PRAGMA table_info(query_history)")
        columns = cursor.fetchall()
        
        has_user_id = any(column[1] == 'user_id' for column in columns)
        
        if has_user_id:
            logger.info("query_history表已存在user_id字段，无需添加")
        else:
            # 添加user_id字段（先不添加外键约束，确保字段可以添加成功）
            logger.info("开始添加user_id字段")
            cursor.execute("ALTER TABLE query_history ADD COLUMN user_id INTEGER")
            conn.commit()
            logger.info("user_id字段添加成功")
            
            # 尝试添加外键约束
            try:
                logger.info("开始添加外键约束")
                cursor.execute("""
                    ALTER TABLE query_history 
                    ADD CONSTRAINT fk_query_history_user_id 
                    FOREIGN KEY (user_id) REFERENCES users(id)
                """)
                conn.commit()
                logger.info("外键约束添加成功")
            except sqlite3.Error as e:
                # 在SQLite中，某些版本不支持通过ALTER TABLE直接添加外键约束
                # 如果添加外键约束失败，记录警告信息并继续
                logger.warning(f"添加外键约束失败: {str(e)}")
                logger.warning("将创建临时表并重新创建query_history表来添加外键约束")
                
                # 创建临时表
                cursor.execute("""
                    CREATE TABLE query_history_temp (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        department TEXT NOT NULL,
                        query_date TEXT NOT NULL,
                        plate_number TEXT,
                        phone TEXT,
                        results_count INTEGER NOT NULL,
                        query_ip TEXT NOT NULL,
                        user_id INTEGER,
                        FOREIGN KEY (user_id) REFERENCES users(id)
                    )
                """)
                
                # 复制数据到临时表
                cursor.execute("""
                    INSERT INTO query_history_temp (
                        id, department, query_date, plate_number, 
                        phone, results_count, query_ip
                    )
                    SELECT 
                        id, department, query_date, plate_number, 
                        phone, results_count, query_ip
                    FROM query_history
                """)
                
                # 删除原表
                cursor.execute("DROP TABLE query_history")
                
                # 重命名临时表
                cursor.execute("ALTER TABLE query_history_temp RENAME TO query_history")
                
                # 提交更改
                conn.commit()
                logger.info("通过重新创建表的方式成功添加外键约束")
        
        # 关闭连接
        conn.close()
        logger.info("操作完成")
        
    except Exception as e:
        logger.error(f"操作过程中发生错误: {str(e)}")
        if 'conn' in locals() and conn:
            conn.rollback()
            conn.close()


if __name__ == '__main__':
    main()