#!/bin/bash

# 导入公共函数
source "$(dirname "$0")/../utils/common.sh"

# 安装系统依赖
install_dependencies() {
    log_info "安装系统依赖..."
    
    # 安装Python开发相关包
    check_dependencies \
        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_success "系统依赖安装完成"
}

# 安装pyenv
install_pyenv() {
    log_info "安装pyenv..."
    
    # 切换到用户目录
    cd "${USER_HOME}"
    
    # 以用户身份下载并安装pyenv
    su - $CURRENT_USER << EOF
    # 下载pyenv安装脚本
    curl -L https://pyenv.run | bash
    
    # 设置环境变量
    export PYENV_ROOT="\$HOME/.pyenv"
    export PATH="\$PYENV_ROOT/bin:\$PATH"
    eval "\$(pyenv init --path)"
    eval "\$(pyenv init -)"
    eval "\$(pyenv virtualenv-init -)"
EOF
    
    # 备份现有配置
    manage_config "backup" "${USER_HOME}/.zshrc"
    
    # 配置pyenv
    cat >> "${USER_HOME}/.zshrc" << EOF
# pyenv配置
export PYENV_ROOT="\$HOME/.pyenv"
export PATH="\$PYENV_ROOT/bin:\$PATH"
eval "\$(pyenv init --path)"
eval "\$(pyenv init -)"
eval "\$(pyenv virtualenv-init -)"
EOF
    
    # 确保.pyenv目录权限正确
    chown -R $CURRENT_USER:$CURRENT_USER "${USER_HOME}/.pyenv"
    
    log_success "pyenv安装完成"
}

# 安装Python
install_python() {
    log_info "安装Python..."
    
    # 切换到用户目录
    cd "${USER_HOME}"
    
    # 以用户身份运行命令
    su - $CURRENT_USER << EOF
    # 设置环境变量
    export PYENV_ROOT="\$HOME/.pyenv"
    export PATH="\$PYENV_ROOT/bin:\$PATH"
    eval "\$(pyenv init --path)"
    eval "\$(pyenv init -)"
    
    # 安装Python依赖
    sudo apt-get install -y make build-essential libssl-dev zlib1g-dev \
    libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
    
    # 安装最新的Python 3.11
    pyenv install 3.11.7
    
    # 设置全局Python版本
    pyenv global 3.11.7
    
    # 验证安装
    python --version
EOF
    
    log_success "Python安装完成"
}

# 配置pip
configure_pip() {
    log_info "配置pip..."
    
    # 切换到用户目录
    cd "${USER_HOME}"
    
    # 以用户身份运行命令
    su - $CURRENT_USER << EOF
    # 设置环境变量
    export PYENV_ROOT="\$HOME/.pyenv"
    export PATH="\$PYENV_ROOT/bin:\$PATH"
    eval "\$(pyenv init --path)"
    eval "\$(pyenv init -)"
    
    # 升级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
EOF
    
    log_success "pip配置完成"
}

# 安装基础Python包
install_python_packages() {
    log_info "安装基础Python包..."
    
    # 切换到用户目录
    cd "${USER_HOME}"
    
    # 以用户身份运行命令
    su - $CURRENT_USER << EOF
    # 设置环境变量
    export PYENV_ROOT="\$HOME/.pyenv"
    export PATH="\$PYENV_ROOT/bin:\$PATH"
    eval "\$(pyenv init --path)"
    eval "\$(pyenv init -)"
    
    # 创建并激活虚拟环境
    python -m venv \$HOME/.venv
    source \$HOME/.venv/bin/activate
    
    # 升级pip
    python -m pip install --upgrade pip
    
    # 设置pip镜像
    pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
    
    # 安装开发工具
    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
    
    # 安装Web开发工具
    pip install django
    pip install flask
    pip install fastapi
    pip install uvicorn
    pip install sqlalchemy
    pip install alembic
    
    # 安装其他工具
    pip install requests
    pip install python-dotenv
    pip install rich
    pip install tqdm
    
    # 退出虚拟环境
    deactivate
EOF
    
    # 备份现有配置
    manage_config "backup" "${USER_HOME}/.zshrc"
    
    # 添加虚拟环境激活命令到配置文件
    cat >> "${USER_HOME}/.zshrc" << EOF
# 激活Python虚拟环境
source \$HOME/.venv/bin/activate 2>/dev/null || true
EOF
    
    log_success "基础Python包安装完成"
}

# 安装机器学习包
install_ml_packages() {
    log_info "安装机器学习包..."
    
    # 切换到用户目录
    cd "${USER_HOME}"
    
    # 以用户身份运行命令
    su - $CURRENT_USER << EOF
    # 设置环境变量
    export PYENV_ROOT="\$HOME/.pyenv"
    export PATH="\$PYENV_ROOT/bin:\$PATH"
    eval "\$(pyenv init --path)"
    eval "\$(pyenv init -)"
    
    # 激活虚拟环境
    source \$HOME/.venv/bin/activate
    
    # 安装科学计算和可视化工具
    pip install scipy
    pip install matplotlib
    pip install seaborn
    
    # 安装机器学习工具
    pip install scikit-learn
    pip install tensorflow
    pip install torch
    pip install transformers
    
    # 退出虚拟环境
    deactivate
EOF
    
    log_success "机器学习包安装完成"
}

# 配置虚拟环境
configure_venv() {
    log_info "配置虚拟环境..."
    
    # 切换到用户目录
    cd "${USER_HOME}"
    
    # 以用户身份运行命令
    su - $CURRENT_USER << EOF
    # 设置环境变量
    export PYENV_ROOT="\$HOME/.pyenv"
    export PATH="\$PYENV_ROOT/bin:\$PATH"
    eval "\$(pyenv init --path)"
    eval "\$(pyenv init -)"
    
    # 使用主虚拟环境安装virtualenvwrapper
    source \$HOME/.venv/bin/activate
    pip install virtualenvwrapper
    
    # 创建必要的目录
    mkdir -p \$HOME/.virtualenvs
    mkdir -p \$HOME/projects
    
    deactivate
EOF
    
    # 备份现有配置
    manage_config "backup" "${USER_HOME}/.zshrc"
    
    # 配置virtualenvwrapper
    cat >> "${USER_HOME}/.zshrc" << EOF
# virtualenvwrapper配置
export WORKON_HOME=\$HOME/.virtualenvs
export PROJECT_HOME=\$HOME/projects
export VIRTUALENVWRAPPER_PYTHON=\$HOME/.venv/bin/python
source \$HOME/.venv/bin/virtualenvwrapper.sh
EOF
    
    # 确保virtualenvwrapper.sh存在
    if [ ! -f "${USER_HOME}/.venv/bin/virtualenvwrapper.sh" ]; then
        log_error "virtualenvwrapper.sh 不存在，虚拟环境配置可能不完整"
        return 1
    fi
    
    # 设置正确的权限
    chown -R $CURRENT_USER:$CURRENT_USER "${USER_HOME}/.venv"
    chmod +x "${USER_HOME}/.venv/bin/virtualenvwrapper.sh"
    
    log_success "虚拟环境配置完成"
}

# 主函数
main() {
    # 检查root权限
    check_root
    
    # 获取用户信息
    get_user_info
    
    # 设置日志
    setup_logging
    
    log_info "开始安装Python开发环境..."
    
    # 安装系统依赖
    install_dependencies
    
    # 安装pyenv
    install_pyenv
    
    # 安装Python
    install_python
    
    # 配置pip
    configure_pip
    
    # 安装基础Python包
    install_python_packages
    
    # 配置虚拟环境
    configure_venv
    
    # 询问是否安装机器学习包
    read -p "是否安装机器学习相关包？这些包体积较大，安装时间较长。(y/N) " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        install_ml_packages
    fi
    
    log_success "Python开发环境安装完成！"
    log_info "请运行 'source ~/.zshrc' 或重新打开终端以使环境变量生效。"
}

# 设置清理函数
trap cleanup EXIT

# 执行主函数
main