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

import os
import sqlite3
import sys

# 数据库路径
db_path = '/root/autodl-tmp/TravelContentCreator/distribution.db'

def main():
    print(f"开始修复数据库: {db_path}")
    
    # 连接数据库
    try:
        conn = sqlite3.connect(db_path)
        conn.execute("PRAGMA foreign_keys = OFF")  # 禁用外键约束
        cursor = conn.cursor()
        print(f"已连接到数据库")
        
        # 备份用户数据
        print("正在备份用户数据...")
        users_data = []
        try:
            cursor.execute("SELECT id, email, username, created_at FROM users")
            users_data = cursor.fetchall()
            print(f"成功备份 {len(users_data)} 条用户记录")
        except sqlite3.Error as e:
            print(f"备份用户数据时出错: {e}")
            print("继续执行修复操作...")
        
        # 检查数据库表结构
        print("检查数据库表结构...")
        try:
            cursor.execute("PRAGMA table_info(users)")
            columns = {col[1] for col in cursor.fetchall()}
            
            # 打印当前表结构
            print(f"当前用户表字段: {', '.join(columns)}")
            
            # 创建新的用户表
            print("创建新的用户表结构...")
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS users_new (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                email TEXT NOT NULL,
                username TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                xhs_user_id TEXT UNIQUE
            )
            """)
            
            # 复制用户数据到新表
            if users_data:
                print("将数据迁移到新表...")
                for user in users_data:
                    user_id, email, username, created_at = user
                    # 使用email作为临时xhs_user_id，稍后可以更新
                    cursor.execute("""
                    INSERT INTO users_new (id, email, username, created_at, xhs_user_id) 
                    VALUES (?, ?, ?, ?, NULL)
                    """, (user_id, email, username, created_at))
                print(f"已将 {len(users_data)} 条记录迁移到新表")
            
            # 查询是否存在users表，如果存在就删除
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='users'")
            if cursor.fetchone():
                cursor.execute("DROP TABLE users")
                print("已删除旧的users表")
            
            # 重命名新表
            cursor.execute("ALTER TABLE users_new RENAME TO users")
            print("已将新表重命名为users")
            
            # 为邮箱创建非唯一索引
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_users_email ON users(email)")
            print("已为email字段创建索引")
            
            # 为xhs_user_id创建唯一索引
            cursor.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_users_xhs_id ON users(xhs_user_id) WHERE xhs_user_id IS NOT NULL")
            print("已为xhs_user_id字段创建唯一索引")
            
            # 提交事务
            conn.commit()
            print("数据库结构修复完成！")
            
        except sqlite3.Error as e:
            print(f"修复表结构时出错: {e}")
            conn.rollback()
            return False
            
    except sqlite3.Error as e:
        print(f"数据库连接错误: {e}")
        return False
    finally:
        print("关闭数据库连接")
        conn.close()
        
    return True

if __name__ == "__main__":
    if main():
        print("数据库修复成功!")
    else:
        print("数据库修复失败!")
        sys.exit(1)