#!/bin/bash

# 设置错误时立即退出
set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 安装系统依赖
install_dependencies() {
    log_info "安装系统依赖..."
    
    # 安装Python开发相关包
    sudo apt-get update
    sudo apt-get install -y \
        python3-dev \
        python3-pip \
        python3-venv \
        build-essential \
        libssl-dev \
        libffi-dev \
        libxml2-dev \
        libxslt1-dev \
        zlib1g-dev \
        libbz2-dev \
        libreadline-dev \
        libsqlite3-dev \
        libncursesw5-dev \
        xz-utils \
        tk-dev \
        libxmlsec1-dev \
        libffi-dev \
        liblzma-dev
    
    log_info "系统依赖安装完成"
}

# 安装pyenv
install_pyenv() {
    log_info "安装pyenv..."
    
    # 下载pyenv安装脚本
    curl https://pyenv.run | bash
    
    # 配置pyenv环境变量
    cat >> ~/.zshrc << EOF
# pyenv配置
export PYENV_ROOT="\$HOME/.pyenv"
export PATH="\$PYENV_ROOT/bin:\$PATH"
eval "\$(pyenv init --path)"
eval "\$(pyenv init -)"
EOF
    
    # 重新加载配置
    export PYENV_ROOT="$HOME/.pyenv"
    export PATH="$PYENV_ROOT/bin:$PATH"
    eval "$(pyenv init --path)"
    eval "$(pyenv init -)"
    
    log_info "pyenv安装完成"
}

# 安装Python
install_python() {
    log_info "安装Python..."
    
    # 安装最新的Python 3.11
    pyenv install 3.11.7
    
    # 设置全局Python版本
    pyenv global 3.11.7
    
    # 验证安装
    python --version
    
    log_info "Python安装完成"
}

# 配置pip
configure_pip() {
    log_info "配置pip..."
    
    # 升级pip
    python -m pip install --upgrade pip
    
    # 设置pip镜像
    pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
    
    # 安装pip工具
    pip install wheel setuptools
    
    log_info "pip配置完成"
}

# 安装Python包
install_python_packages() {
    log_info "安装Python包..."
    
    # 安装开发工具
    pip install ipython
    pip install jupyter
    pip install jupyterlab
    pip install black
    pip install flake8
    pip install mypy
    pip install pylint
    pip install pytest
    pip install pytest-cov
    
    # 安装数据处理工具
    pip install numpy
    pip install pandas
    pip install scipy
    pip install matplotlib
    pip install seaborn
    
    # 安装Web开发工具
    pip install django
    pip install flask
    pip install fastapi
    pip install uvicorn
    pip install sqlalchemy
    pip install alembic
    
    # 安装机器学习工具
    pip install scikit-learn
    pip install tensorflow
    pip install torch
    pip install transformers
    
    # 安装其他工具
    pip install requests
    pip install python-dotenv
    pip install rich
    pip install tqdm
    
    log_info "Python包安装完成"
}

# 配置虚拟环境
configure_venv() {
    log_info "配置虚拟环境..."
    
    # 创建虚拟环境目录
    mkdir -p ~/.virtualenvs
    
    # 安装virtualenvwrapper
    pip install virtualenvwrapper
    
    # 配置virtualenvwrapper
    cat >> ~/.zshrc << EOF
# virtualenvwrapper配置
export WORKON_HOME=~/.virtualenvs
export PROJECT_HOME=~/projects
source ~/.local/bin/virtualenvwrapper.sh
EOF
    
    log_info "虚拟环境配置完成"
}

# 主函数
main() {
    log_info "开始安装Python开发环境..."
    
    # 安装系统依赖
    install_dependencies
    
    # 安装pyenv
    install_pyenv
    
    # 安装Python
    install_python
    
    # 配置pip
    configure_pip
    
    # 安装Python包
    install_python_packages
    
    # 配置虚拟环境
    configure_venv
    
    log_info "Python开发环境安装完成！"
    log_info "请运行 'source ~/.zshrc' 或重新打开终端以使环境变量生效。"
}

# 执行主函数
main