"""
创建新数据库并初始化所有表
数据库名: fa3_new
"""
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
import os
from dotenv import load_dotenv

load_dotenv()

def create_database():
    """创建新数据库"""
    # 连接到 postgres 数据库（默认数据库）
    conn = psycopg2.connect(
        host=os.getenv("DB_HOST", "192.168.3.195"),
        user=os.getenv("DB_USER", "value"),
        password=os.getenv("DB_PASSWORD", "Value609!"),
        dbname="postgres",  # 连接到默认数据库
        port=os.getenv("DB_PORT", "5432")
    )
    conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
    cursor = conn.cursor()
    
    try:
        # 检查数据库是否存在
        cursor.execute("SELECT 1 FROM pg_database WHERE datname='fa3_new'")
        exists = cursor.fetchone()
        
        if exists:
            print("⚠️  数据库 fa3_new 已存在")
            response = input("是否删除并重建？(yes/no): ")
            if response.lower() == 'yes':
                # 断开所有连接
                cursor.execute("""
                    SELECT pg_terminate_backend(pg_stat_activity.pid)
                    FROM pg_stat_activity
                    WHERE pg_stat_activity.datname = 'fa3_new'
                    AND pid <> pg_backend_pid()
                """)
                cursor.execute("DROP DATABASE fa3_new")
                print("✅ 已删除旧数据库")
            else:
                print("❌ 取消操作")
                return False
        
        # 创建新数据库（使用 template0 以支持 UTF-8）
        cursor.execute("""
            CREATE DATABASE fa3_new 
            WITH OWNER = value 
            ENCODING = 'UTF8'
            TEMPLATE = template0
        """)
        print("✅ 成功创建数据库: fa3_new")
        return True
        
    except Exception as e:
        print(f"❌ 创建数据库失败: {e}")
        return False
    finally:
        cursor.close()
        conn.close()

def initialize_database():
    """初始化数据库表结构"""
    conn = psycopg2.connect(
        host=os.getenv("DB_HOST", "192.168.3.195"),
        user=os.getenv("DB_USER", "value"),
        password=os.getenv("DB_PASSWORD", "Value609!"),
        dbname="fa3_new",  # 连接到新数据库
        port=os.getenv("DB_PORT", "5432")
    )
    conn.set_client_encoding('UTF8')
    conn.autocommit = True
    cursor = conn.cursor()
    
    try:
        # 读取初始化脚本
        with open('migrations/init_new_database.sql', 'r', encoding='utf-8') as f:
            sql = f.read()
        
        # 执行SQL
        cursor.execute(sql)
        
        print("✅ 数据库表结构初始化完成")
        
        # 显示创建的表
        cursor.execute("""
            SELECT schemaname, tablename 
            FROM pg_tables 
            WHERE schemaname = 'public'
            ORDER BY tablename
        """)
        tables = cursor.fetchall()
        print(f"\n📊 已创建 {len(tables)} 个表:")
        for schema, table in tables:
            print(f"   - {table}")
        
        return True
        
    except Exception as e:
        print(f"❌ 初始化数据库失败: {e}")
        return False
    finally:
        cursor.close()
        conn.close()

def main():
    print("="*60)
    print("FA3 新数据库创建工具")
    print("="*60)
    
    # 1. 创建数据库
    if not create_database():
        return
    
    # 2. 初始化表结构
    print("\n开始初始化表结构...")
    if not initialize_database():
        return
    
    print("\n" + "="*60)
    print("✅ 新数据库创建完成！")
    print("="*60)
    print("\n下一步:")
    print("1. 更新 .env 文件中的 DB_NAME=fa3_new")
    print("2. 或使用环境变量: export DB_NAME=fa3_new")
    print("3. 重启后端服务即可使用新数据库")

if __name__ == "__main__":
    main()