# 设置错误操作首选项
$ErrorActionPreference = "Stop"

# 检查管理员权限
$isAdmin = ([Security.Principal.WindowsPrincipal] [Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
if (-not $isAdmin) {
    Write-Warning "建议使用管理员权限运行此脚本"
    Start-Sleep -Seconds 3
}

# 检查 pyenv-win 是否安装
if (!(Get-Command pyenv -ErrorAction SilentlyContinue)) {
    Write-Host "pyenv not found. Installing..."
    try {
        # 下载并安装 pyenv-win
        Invoke-WebRequest -UseBasicParsing -Uri "https://raw.githubusercontent.com/pyenv-win/pyenv-win/master/pyenv-win/install-pyenv-win.ps1" -OutFile "./install-pyenv-win.ps1"
        & ./install-pyenv-win.ps1
        
        # 添加环境变量
        $env:PYENV = "$env:USERPROFILE\.pyenv\pyenv-win"
        $env:Path = "$env:PYENV\bin;$env:PYENV\shims;$env:Path"
        
        # 刷新环境变量
        $env:Path = [System.Environment]::GetEnvironmentVariable("Path","Machine") + ";" + [System.Environment]::GetEnvironmentVariable("Path","User")
    }
    catch {
        Write-Error "Failed to install pyenv: $_"
        exit 1
    }
}

try {
    # 安装指定版本的 Python
    Write-Host "Installing Python 3.11.8..."
    pyenv install 3.11.8
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to install Python 3.11.8"
    }

    # 设置本地 Python 版本
    Write-Host "Setting local Python version..."
    pyenv local 3.11.8
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to set local Python version"
    }

    # 验证 Python 版本
    $pythonVersion = python -V
    if (-not $pythonVersion.Contains("3.11.8")) {
        throw "Wrong Python version: $pythonVersion"
    }
    Write-Host "Using Python version: $pythonVersion"

    # 创建虚拟环境
    Write-Host "Creating virtual environment..."
    python -m venv .venv

    # 激活虚拟环境
    Write-Host "Activating virtual environment..."
    .\.venv\Scripts\Activate.ps1

    # 升级 pip 和构建工具
    Write-Host "Upgrading pip and build tools..."
    python -m pip install --upgrade pip setuptools wheel

    # 分步安装依赖以确保正确的顺序和版本
    Write-Host "Installing database dependencies..."
    pip install mysql-connector-python==8.0.33

    Write-Host "Installing PyTorch and related packages..."
    pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cpu

    Write-Host "Installing basic dependencies..."
    pip install numpy==1.26.4 pandas==2.2.1

    Write-Host "Installing machine learning packages..."
    pip install scikit-learn==1.5.2

    # 安装开发依赖
    Write-Host "Installing development dependencies..."
    pip install -e ".[dev]"
    if ($LASTEXITCODE -ne 0) {
        Write-Warning "Failed to install development dependencies. Installing core package..."
        pip install -e .
    }

    # 验证安装
    Write-Host "Verifying installations..."
    python -c "import torch; print(f'PyTorch version: {torch.__version__}')"
    python -c "import numpy; print(f'NumPy version: {numpy.__version__}')"
    python -c "import pandas; print(f'Pandas version: {pandas.__version__}')"
    python -c "import sklearn; print(f'Scikit-learn version: {sklearn.__version__}')"

    Write-Host "Environment setup complete!" -ForegroundColor Green
}
catch {
    Write-Error "An error occurred: $_"
    exit 1
}
finally {
    # 清理临时文件
    if (Test-Path "./install-pyenv-win.ps1") {
        Remove-Item "./install-pyenv-win.ps1"
    }
}

# 显示使用说明
Write-Host @"

环境设置完成！使用说明：
1. 虚拟环境已激活，命令提示符前应该显示 (.venv)
2. 要退出虚拟环境，运行: deactivate
3. 要重新激活虚拟环境，运行: .\.venv\Scripts\Activate.ps1
4. 项目依赖已安装，可以开始开发了

如果遇到问题，请检查：
- Python 版本: python -V
- PyTorch 安装: python -c "import torch; print(torch.__version__)"
- 虚拟环境状态: 确保看到 (.venv) 前缀

"@ -ForegroundColor Cyan