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

"""
自动安装脚本
自动安装所需依赖并验证环境
"""

import subprocess
import sys
import time
import os

def print_header(title):
    """打印标题"""
    print("\n" + "=" * 60)
    print(f"  {title}")
    print("=" * 60)

def run_command(command, description):
    """运行命令并显示结果"""
    print(f"\n🔄 {description}...")
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True)
        if result.returncode == 0:
            print(f"✅ {description} 成功完成")
            return True
        else:
            print(f"❌ {description} 失败")
            print(f"错误信息: {result.stderr}")
            return False
    except Exception as e:
        print(f"❌ {description} 执行异常: {e}")
        return False

def check_python_version():
    """检查Python版本"""
    print_header("检查Python版本")
    
    version = sys.version_info
    print(f"当前Python版本: {version.major}.{version.minor}.{version.micro}")
    
    if version.major < 3 or (version.major == 3 and version.minor < 8):
        print("❌ Python版本过低，需要Python 3.8或更高版本")
        return False
    else:
        print("✅ Python版本满足要求")
        return True

def install_requirements():
    """安装依赖包"""
    print_header("安装依赖包")
    
    if not os.path.exists("requirements.txt"):
        print("❌ requirements.txt 文件不存在")
        return False
    
    # 升级pip
    if not run_command(f"{sys.executable} -m pip install --upgrade pip", "升级pip"):
        print("⚠️ pip升级失败，继续安装依赖")
    
    # 安装基础依赖
    success1 = run_command(f"{sys.executable} -m pip install -r requirements.txt", "安装基础依赖包")
    
    # 安装Flask（Web功能所需）
    success2 = install_flask()
    
    return success1 and success2

def install_flask():
    """安装Flask及其依赖"""
    print("\n🔧 安装Flask（Web功能支持）...")
    
    try:
        # 检查是否已安装
        import flask
        print(f"✅ Flask已安装，版本: {flask.__version__}")
        return True
    except ImportError:
        pass
    
    # 尝试使用清华镜像源
    result = subprocess.run([
        sys.executable, "-m", "pip", "install", 
        "flask==2.3.3", 
        "-i", "https://pypi.tuna.tsinghua.edu.cn/simple/",
        "--trusted-host", "pypi.tuna.tsinghua.edu.cn"
    ], capture_output=True, text=True, timeout=300)
    
    if result.returncode == 0:
        print("✅ Flask安装成功！")
        return True
    else:
        # 尝试官方源
        result = subprocess.run([
            sys.executable, "-m", "pip", "install", "flask==2.3.3"
        ], capture_output=True, text=True, timeout=300)
        
        if result.returncode == 0:
            print("✅ Flask安装成功！")
            return True
        else:
            print(f"❌ Flask安装失败: {result.stderr}")
            return False

def verify_installation():
    """验证安装"""
    print_header("验证安装")
    
    modules_to_check = [
        ("cv2", "OpenCV"),
        ("numpy", "NumPy"),
        ("ultralytics", "Ultralytics"),
        ("torch", "PyTorch"),
        ("flask", "Flask"),
    ]
    
    all_success = True
    
    for module, name in modules_to_check:
        try:
            __import__(module)
            print(f"✅ {name} 安装成功")
        except ImportError:
            print(f"❌ {name} 安装失败")
            all_success = False
    
    return all_success

def download_yolo_model():
    """预下载YOLO模型"""
    print_header("下载YOLO模型")
    
    try:
        from ultralytics import YOLO
        print("🔄 正在下载YOLOv8n模型，请稍候...")
        model = YOLO('yolov8n.pt')
        print("✅ YOLOv8n模型下载成功")
        
        # 测试模块结构
        try:
            from src import PersonDetector, DistanceCalculator
            print("✅ 重构后的模块结构测试成功")
        except ImportError as e:
            print(f"⚠️ 模块结构测试失败: {e}")
        
        return True
    except Exception as e:
        print(f"❌ YOLO模型下载失败: {e}")
        return False

def run_test():
    """运行测试"""
    print_header("运行系统测试")
    
    if os.path.exists("test_modules.py"):
        return run_command(f"{sys.executable} test_modules.py", "运行系统测试")
    else:
        print("⚠️ 测试脚本不存在，跳过测试")
        return True

def create_desktop_shortcut():
    """创建桌面快捷方式（Windows）"""
    try:
        import platform
        if platform.system() == "Windows":
            print_header("创建桌面快捷方式")
            
            desktop_path = os.path.join(os.path.expanduser("~"), "Desktop")
            if os.path.exists(desktop_path):
                shortcut_content = f"""
@echo off
cd /d "{os.getcwd()}"
python run.py
pause
"""
                shortcut_path = os.path.join(desktop_path, "人体距离检测.bat")
                with open(shortcut_path, "w", encoding="gbk") as f:
                    f.write(shortcut_content)
                print(f"✅ 桌面快捷方式已创建: {shortcut_path}")
                return True
            else:
                print("⚠️ 未找到桌面路径")
                return False
    except Exception as e:
        print(f"⚠️ 创建快捷方式失败: {e}")
        return False

def main():
    """主安装函数"""
    print("🚀 人体距离检测系统 - 自动安装程序")
    print("此程序将自动安装所需依赖并配置环境")
    
    steps = [
        ("检查Python版本", check_python_version),
        ("安装依赖包", install_requirements),
        ("验证安装", verify_installation),
        ("下载YOLO模型", download_yolo_model),
        ("运行系统测试", run_test),
    ]
    
    # 可选步骤
    optional_steps = [
        ("创建桌面快捷方式", create_desktop_shortcut),
    ]
    
    print(f"\n📋 安装计划:")
    for i, (name, _) in enumerate(steps, 1):
        print(f"  {i}. {name}")
    
    print(f"\n可选步骤:")
    for i, (name, _) in enumerate(optional_steps, 1):
        print(f"  {i}. {name}")
    
    input("\n按Enter键开始安装...")
    
    start_time = time.time()
    success_count = 0
    total_steps = len(steps)
    
    # 执行主要安装步骤
    for i, (name, func) in enumerate(steps, 1):
        print(f"\n📦 步骤 {i}/{total_steps}: {name}")
        
        if func():
            success_count += 1
            print(f"✅ 步骤 {i} 完成")
        else:
            print(f"❌ 步骤 {i} 失败")
            
            # 询问是否继续
            choice = input("是否继续安装？(y/n): ").lower()
            if choice != 'y':
                print("安装已取消")
                return
    
    # 执行可选步骤
    for name, func in optional_steps:
        choice = input(f"\n是否执行: {name}？(y/n): ").lower()
        if choice == 'y':
            func()
    
    # 显示安装结果
    elapsed_time = time.time() - start_time
    print_header("安装完成")
    
    print(f"✅ 安装步骤: {success_count}/{total_steps} 完成")
    print(f"⏱️ 总耗时: {elapsed_time:.1f}秒")
    
    if success_count == total_steps:
        print("🎉 所有步骤都已成功完成！")
        print("\n📖 使用指南:")
        print("  1. 运行 'python run.py' 启动系统")
        print("  2. 选择运行模式（Web或传统模式）")
        print("  3. 如需配置摄像头位置，运行 setup_camera_location.py")
        print("  4. 如遇问题，运行 test_modules.py 进行诊断")
    else:
        print("⚠️ 部分步骤未成功，请检查错误信息")
        print("💡 如需帮助，请查看README.md文档")
    
    print("\n安装程序结束！")

if __name__ == "__main__":
    main()