#!/usr/bin/env python3
"""
Cursor自动化服务 - 修复脚本
添加虚拟环境支持并修复MySQL连接问题
"""
import os
import sys
import subprocess
import shutil
import time

# 彩色输出函数
class Colors:
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    BLUE = '\033[94m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'

def print_success(msg):
    print(f"{Colors.GREEN}✓ {msg}{Colors.ENDC}")

def print_warning(msg):
    print(f"{Colors.YELLOW}⚠ {msg}{Colors.ENDC}")

def print_error(msg):
    print(f"{Colors.RED}✗ {msg}{Colors.ENDC}")

def print_title(title):
    print("\n" + "=" * 60)
    print(f"{Colors.BOLD}{Colors.BLUE}{title.center(60)}{Colors.ENDC}")
    print("=" * 60 + "\n")

def check_virtual_env():
    """检查是否在虚拟环境中运行"""
    return hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix)

def setup_virtual_env():
    """设置Python虚拟环境"""
    print_title("虚拟环境设置")
    
    if check_virtual_env():
        print_success("当前已在虚拟环境中运行")
        return True
    
    print_warning("当前未在虚拟环境中运行")
    create_venv = input("是否创建虚拟环境? (推荐) (y/n, 默认: y): ").strip().lower()
    if create_venv == 'n':
        print_warning("跳过虚拟环境创建")
        return True
    
    venv_path = input("请输入虚拟环境路径 (默认: ./venv): ").strip() or "./venv"
    
    try:
        # 创建虚拟环境
        print(f"创建虚拟环境: {venv_path}")
        subprocess.check_call([sys.executable, "-m", "venv", venv_path])
        
        # 获取激活脚本路径
        if sys.platform.startswith('win'):
            activate_script = os.path.join(venv_path, "Scripts", "activate.bat")
            python_exe = os.path.join(venv_path, "Scripts", "python.exe")
        else:
            activate_script = os.path.join(venv_path, "bin", "activate")
            python_exe = os.path.join(venv_path, "bin", "python")
        
        # 创建安装脚本
        print("创建安装脚本...")
        setup_script = "setup_venv.sh"
        with open(setup_script, "w") as f:
            if sys.platform.startswith('win'):
                f.write(f"@echo off\n")
                f.write(f"call {activate_script}\n")
                f.write(f"pip install -r requirements.txt\n")
                f.write(f"pip install cryptography\n")
                f.write(f"python setup_environment.py\n")
                f.write(f"pause\n")
            else:
                f.write(f"#!/bin/bash\n")
                f.write(f"source {activate_script}\n")
                f.write(f"pip install -r requirements.txt\n")
                f.write(f"pip install cryptography\n")
                f.write(f"python setup_environment.py\n")
        
        if not sys.platform.startswith('win'):
            os.chmod(setup_script, 0o755)
        
        print_success(f"虚拟环境创建成功: {venv_path}")
        print_success(f"安装脚本已创建: {setup_script}")
        
        print("\n请运行以下命令完成安装:")
        if sys.platform.startswith('win'):
            print(f"  {setup_script}")
        else:
            print(f"  ./{setup_script}")
        
        return False
    except Exception as e:
        print_error(f"创建虚拟环境失败: {str(e)}")
        return True

def fix_mysql_format_issue():
    """修复MySQL格式化问题"""
    print_title("修复MySQL连接问题")
    
    if not os.path.exists("setup_environment.py"):
        print_error("未找到setup_environment.py文件")
        return False
    
    backup_file = f"setup_environment.py.bak.{int(time.time())}"
    try:
        # 创建备份
        shutil.copy2("setup_environment.py", backup_file)
        print_success(f"已创建备份: {backup_file}")
        
        # 读取文件内容
        with open("setup_environment.py", "r", encoding="utf-8") as f:
            content = f.read()
        
        # 修复格式化问题
        fixed_content = content
        
        # 1. 添加cryptography到required_packages
        if "cryptography" not in content:
            fixed_content = fixed_content.replace(
                'required_packages = [\n    "loguru",\n    "pymysql",\n    "aiomysql",\n    "redis",\n    "pyyaml"',
                'required_packages = [\n    "loguru",\n    "pymysql",\n    "aiomysql",\n    "redis",\n    "pyyaml",\n    "cryptography"'
            )
            print_success("已添加cryptography到依赖包列表")
        
        # 2. 修复CREATE USER语句的格式化问题
        if "cursor.execute(f\"CREATE USER '{db_user}'@'%'" in content:
            fixed_content = fixed_content.replace(
                "cursor.execute(f\"CREATE USER '{db_user}'@'%' IDENTIFIED BY %s\", (db_password,))",
                "cursor.execute(\"CREATE USER %s@'%%' IDENTIFIED BY %s\", (db_user, db_password))"
            )
            print_success("已修复CREATE USER语句")
        
        # 3. 修复ALTER USER语句的格式化问题
        if "cursor.execute(f\"ALTER USER '{db_user}'@'%'" in content:
            fixed_content = fixed_content.replace(
                "cursor.execute(f\"ALTER USER '{db_user}'@'%' IDENTIFIED BY %s\", (db_password,))",
                "cursor.execute(\"ALTER USER %s@'%%' IDENTIFIED BY %s\", (db_user, db_password))"
            )
            print_success("已修复ALTER USER语句")
        
        # 写入修复后的内容
        with open("setup_environment.py", "w", encoding="utf-8") as f:
            f.write(fixed_content)
        
        print_success("MySQL连接问题修复完成")
        return True
    except Exception as e:
        print_error(f"修复过程出错: {str(e)}")
        try:
            if os.path.exists(backup_file):
                shutil.copy2(backup_file, "setup_environment.py")
                print_warning("已恢复原始文件")
        except:
            pass
        return False

def main():
    print_title("Cursor自动化服务 - 修复工具")
    
    # 设置虚拟环境
    if not setup_virtual_env():
        # 如果创建了虚拟环境并生成了安装脚本，直接退出
        sys.exit(0)
    
    # 修复MySQL格式化问题
    if not fix_mysql_format_issue():
        print_error("修复失败，请手动检查setup_environment.py文件")
        sys.exit(1)
    
    print_title("修复完成")
    print("现在可以安全地运行设置向导:")
    print("  python setup_environment.py")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n操作已取消")
        sys.exit(1)