#!/bin/bash

# 使用当前用户的家目录环境变量
HOME_PATH="$HOME"
MAMBA_ROOT_PREFIX="$HOME/micromamba"
JUPYTERHUB_ENV_NAME="jupyterhub_env"
MICROMAMBA_INSTALL_SCRIPT="./micromamba/install.sh"
CONDARC_PATH="$HOME/.condarc"

# 检查 micromamba 是否安装
check_micromamba_installed() {
    if [[ -f "$MAMBA_ROOT_PREFIX/bin/micromamba" ]]; then
        echo "micromamba 已安装。"
        return 0
    else
        echo "micromamba 未安装。"
        return 1
    fi
}

# 安装 micromamba
install_micromamba() {
    echo "开始安装 micromamba..."
    bash "$MICROMAMBA_INSTALL_SCRIPT" --silent-install
    source "$HOME/.bashrc"
}

# 创建并激活 jupyterhub_env
create_and_activate_env() {
    local env_name=$JUPYTERHUB_ENV_NAME
    local env_dir="$MAMBA_ROOT_PREFIX/envs"

    echo "检查 $env_name 环境是否存在..."
    if [ -d "$env_dir/$env_name" ]; then
        echo "环境 $env_name 已存在, 跳过。如果需要删除 使用命令： micromamba remove -n $env_name --all --yes 再次执行"
    else
        echo "环境 $env_name 不存在，正在创建并激活..."
        micromamba create -n jupyterhub_env -c conda-forge nodejs jupyterhub jupyterlab notebook jupyter-rsession-proxy ipykernel git jupyterlab-language-pack-zh-CN jupyterlab-git jupyterlab-system-monitor jupyter_nbextensions_configurator jupyter_contrib_nbextensions jupyterlab-unfold jupyterlab-variableinspector jupyterlab_widgets jupyterlab-drawio jupyterlab-spreadsheet-editor jupyterlab-cell-flash jedi-language-server jupyterlab_code_formatter jupyterlab-spellchecker jupyterlab_vim nbresuse ipydrawio jedi ipympl black isort theme-darcula ipywidgets tensorboard jupyterlab_latex jupyter_bokeh  autopep8 xeus-python jupyterlab-lsp python-lsp-server nglview mdtraj --yes
    fi
}

# 配置 JupyterHub
configure_jupyterhub() {
    echo "配置 JupyterHub..."
    mkdir -p "$HOME/.jupyterhub"
    cp "./config/jupyterhub_config.yml" "$HOME/.jupyterhub/jupyterhub_config.yml"
}

# 配置 .condarc 文件
configure_condarc() {
    echo "配置 .condarc 文件..."
    cat <<EOF > "$CONDARC_PATH"
ssl_verify: true
channels:
  - defaults
  - https://levinthal:paradox@conda.graylab.jhu.edu
EOF
}

# 修改 JupyterHub 配置文件中的端口号
modify_jupyterhub_port() {
    local config_path="$HOME/.jupyterhub/jupyterhub_config.yml"
    local new_port=$1

    # 使用 sed 来修改配置文件中的端口号
    sed -i "s/^c\.JupyterHub\.port = .*/c.JupyterHub.port = $new_port/" "$config_path"
    echo "JupyterHub 端口已更改为 $new_port。"
}

# 设置 JupyterHub 服务开机启动
setup_systemd_service() {
    local service_file="/etc/systemd/system/jupyterhub.service"
    local user_name=$(whoami)
    echo "启动权限：$user_name"

    cat <<EOF | sudo tee "$service_file"
[Unit]
Description=JupyterHub Service
After=network.target

[Service]
User=$user_name
ExecStart=$HOME/.jupyterhub/jupyterhub_entrypoint.sh

[Install]
WantedBy=multi-user.target
EOF

    # 创建启动脚本
    cat <<EOF | tee "$HOME/.jupyterhub/jupyterhub_entrypoint.sh"
#!/bin/bash
exec jupyterhub -f $HOME/.jupyterhub/jupyterhub_config.yml
EOF

    chmod +x "$HOME/.jupyterhub/jupyterhub_entrypoint.sh"

    # 重新加载 systemd 并启用服务
    sudo systemctl daemon-reload
    sudo systemctl enable jupyterhub.service
    echo "JupyterHub 服务已设置为开机启动。"
}

# 显示菜单并处理用户输入
show_menu() {
    echo "请选择操作："
    echo "1) 安装 micromamba"
    echo "2) 配置 JupyterHub"
    echo "3) 设置 JupyterHub 服务开机启动"
    echo "q) 退出"
    read -p "输入选项（1/2/3/q）: " choice
    case "$choice" in
        1) install_micromamba_and_create_env ;;
        2) configure_jupyterhub_and_condarc ;;
        3) setup_systemd_service ;;
        q) exit 0 ;;
        *) echo "无效选项，请重新输入！" ;;
    esac
}

# 安装 micromamba 并创建环境
install_micromamba_and_create_env() {
    if ! check_micromamba_installed; then
        install_micromamba
    fi
    create_and_activate_env
    echo "Micromamba 安装并激活 jupyterhub_env 完成。"
}

# 配置 JupyterHub 和 .condarc 文件
configure_jupyterhub_and_condarc() {
    configure_jupyterhub
    configure_condarc
    modify_jupyterhub_port 9000  # 更改端口为 9000，可根据需要修改
    echo "JupyterHub 和 .condarc 文件配置完成。"
}

# 主函数
main() {
    while true; do
        show_menu
    done
}

# 执行主函数
main