#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MySQL数据库创建脚本
"""

import pymysql
import sys

def create_database():
    """创建MySQL数据库"""
    print("=" * 50)
    print("MySQL数据库创建")
    print("=" * 50)
    
    try:
        # 连接MySQL服务器（不指定数据库）
        connection = pymysql.connect(
            host='localhost',
            user='root',
            password='123456',
            charset='utf8mb4'
        )
        
        cursor = connection.cursor()
        
        # 创建数据库
        database_name = 'tsp_assistant'
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {database_name} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
        print(f"✓ 数据库 {database_name} 创建成功")
        
        # 创建用户（可选）
        try:
            cursor.execute("CREATE USER IF NOT EXISTS 'tsp_user'@'localhost' IDENTIFIED BY 'tsp_password'")
            cursor.execute(f"GRANT ALL PRIVILEGES ON {database_name}.* TO 'tsp_user'@'localhost'")
            cursor.execute("FLUSH PRIVILEGES")
            print("✓ 用户 tsp_user 创建成功")
        except Exception as e:
            print(f"⚠ 用户创建失败（可能已存在）: {e}")
        
        connection.commit()
        cursor.close()
        connection.close()
        
        print("✓ MySQL数据库设置完成")
        return True
        
    except Exception as e:
        print(f"✗ MySQL数据库创建失败: {e}")
        print("\n请检查：")
        print("1. MySQL服务是否已启动")
        print("2. 用户名和密码是否正确")
        print("3. 是否有创建数据库的权限")
        return False

def test_connection():
    """测试数据库连接"""
    print("\n" + "=" * 50)
    print("测试数据库连接")
    print("=" * 50)
    
    try:
        connection = pymysql.connect(
            host='localhost',
            user='root',
            password='123456',
            database='tsp_assistant',
            charset='utf8mb4'
        )
        
        cursor = connection.cursor()
        cursor.execute("SELECT VERSION()")
        version = cursor.fetchone()
        print(f"✓ MySQL连接成功，版本: {version[0]}")
        
        cursor.close()
        connection.close()
        return True
        
    except Exception as e:
        print(f"✗ MySQL连接失败: {e}")
        return False

def main():
    """主函数"""
    print("TSP助手MySQL数据库设置工具")
    print("=" * 50)
    
    # 创建数据库
    if create_database():
        # 测试连接
        if test_connection():
            print("\n" + "=" * 50)
            print("MySQL数据库设置成功！")
            print("=" * 50)
            print("现在您可以运行以下命令初始化数据库：")
            print("python init_database.py")
        else:
            print("\n" + "=" * 50)
            print("数据库连接测试失败！")
            print("=" * 50)
    else:
        print("\n" + "=" * 50)
        print("MySQL数据库设置失败！")
        print("=" * 50)

if __name__ == "__main__":
    main()