#!/bin/bash

# 定义仓库基础地址
repo_base_url="http://1.94.193.239/debrepo/agiros/ubuntu2204lts/"

# 系统架构检测
arch=$(dpkg --print-architecture 2>/dev/null || uname -m)
if [[ "$arch" == "armhf" ]]; then
    ros_arch="armhf"
else
    ros_arch="amd64"
fi

# 系统版本检测
os_codename=$(lsb_release -cs 2>/dev/null || echo "jammy")

# 公钥相关地址
gpg_key_url="${repo_base_url}agiros.gpg"
ros_gpg_key_urls=(
    "https://gitee.com/ohhuo/rosdistro/raw/master/ros.asc"
    "https://raw.githubusercontent.com/ros/rosdistro/master/ros.asc"
)

# 临时文件路径
temp_gpg_key_file="agiros.gpg"
ros_temp_gpg_key="ros.asc"
keyring_file="/usr/share/keyrings/agiros.gpg"
ros_keyring_file="/usr/share/keyrings/ros-archive-keyring.gpg"

# 权限检查函数
check_permission() {
    if [[ $EUID -ne 0 ]]; then
        echo "错误: 此操作需要管理员权限，请使用sudo运行脚本"
        exit 1
    fi
}

# 替换系统源为清华源（含ROS源）
replace_sources_with_tuna() {
    backup_file="/etc/apt/sources.list.backup"
    echo "正在备份系统源到 $backup_file"
    cp /etc/apt/sources.list "$backup_file" 2>/dev/null || {
        echo "备份系统源失败，将直接替换"
    }
    
    cat << EOF | tee /etc/apt/sources.list > /dev/null
# 默认注释了源码镜像以提高 apt update 速度，如有需要可自行取消注释
deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ $os_codename main restricted universe multiverse
# deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ $os_codename main restricted universe multiverse
deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ $os_codename-updates main restricted universe multiverse
# deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ $os_codename-updates main restricted universe multiverse
deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ $os_codename-backports main restricted universe multiverse
# deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ $os_codename-backports main restricted universe multiverse
deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ $os_codename-security main restricted universe multiverse
# deb-src https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ $os_codename-security main restricted universe multiverse
EOF

    # 添加清华ROS源
    ros_source_file="/etc/apt/sources.list.d/ros-fish.list"
    echo "deb [arch=$ros_arch] http://mirrors.tuna.tsinghua.edu.cn/ros2/ubuntu/ $os_codename main" | tee -a "$ros_source_file" > /dev/null
    
    if [[ $? -ne 0 ]]; then
        echo "替换系统源失败"
        exit 1
    fi
    echo "系统源已成功替换为清华源"
}

# 下载AGIROS公钥
download_agiros_key() {
    echo "正在下载AGIROS公钥..."
    if ! wget -qO "$temp_gpg_key_file" "$gpg_key_url"; then
        echo "下载AGIROS公钥失败"
        exit 1
    fi
    if ! mv "$temp_gpg_key_file" "$keyring_file"; then
        echo "移动AGIROS公钥到 $keyring_file 失败"
        if [[ -f "$temp_gpg_key_file" ]]; then
            rm "$temp_gpg_key_file"
        fi
        exit 1
    fi
    echo "AGIROS公钥已安装"
}

# 下载ROS公钥
download_ros_key() {
    #echo "正在下载ROS公钥..."
    fastest_key=""
    for url in "${ros_gpg_key_urls[@]}"; do
        if wget -q --spider "$url"; then
            fastest_key="$url"
            break
        fi
    done
    
    if [[ -z "$fastest_key" ]]; then
        #echo "下载ROS公钥失败，将尝试从keyserver获取"
        if ! apt-key adv --keyserver keyserver.ubuntu.com --recv-keys F42ED6FBAB17C654; then
            #echo "从keyserver获取ROS公钥失败"
            return 1
        fi
        #echo "已从keyserver获取ROS公钥"
        return 0
    fi
    
    #echo "已选择最快公钥服务: $fastest_key"
    if ! wget -qO "$ros_temp_gpg_key" "$fastest_key"; then
        #echo "下载ROS公钥失败"
        return 1
    fi
    
    if [[ -f "$ros_keyring_file" ]]; then
        rm "$ros_keyring_file"
    fi
    if ! mv "$ros_temp_gpg_key" "$ros_keyring_file"; then
        #echo "移动ROS公钥到 $ros_keyring_file 失败"
        if [[ -f "$ros_temp_gpg_key" ]]; then
            rm "$ros_temp_gpg_key"
        fi
        return 1
    fi
    #echo "ROS公钥已安装"
    return 0
}

# 添加AGIROS仓库源
add_agiros_source() {
    repo_list_file="/etc/apt/sources.list.d/agiros.list"
    repo_source="deb [signed-by=$keyring_file] ${repo_base_url} $os_codename main"
    
    if ! echo "$repo_source" | tee "$repo_list_file" > /dev/null; then
        local error_message=$(echo "$repo_source" | tee "$repo_list_file" 2>&1)
        echo "写入AGIROS仓库源文件失败: $error_message"
        exit 1
    fi
    echo "AGIROS仓库源已添加"
}

# 更新软件包列表
update_package_list() {
    echo "正在更新软件包列表..."
    if ! apt update; then
        echo "更新软件包列表失败"
        exit 1
    fi
    echo "软件包列表更新完成"
}

# 安装软件包
install_package() {
    install_mode="$1"
    
    case $install_mode in
        1)
            package_name="agiros-loong-ros-base"
            ;;
        2)
            package_name="agiros-loong-desktop"
            ;;
        3)
            package_name="agiros-loong-desktop-full"
            ;;
        4)
            package_name="agiros-loong-*"
            ;;
        *)
            echo "无效的安装模式选择"
            exit 1
            ;;
    esac
    
    echo "开始安装 $package_name ..."
    if ! apt install -y "$package_name"; then
        echo "安装软件包 $package_name 失败，尝试修复依赖..."
        
        # 尝试使用aptitude解决依赖问题
        if ! aptitude install -y "$package_name"; then
            echo "依赖修复失败，请手动处理"
            exit 1
        fi
    fi
    echo "软件包安装完成"
}

# 配置环境变量
configure_environment_variable() {
    env_var_dir="/opt/agiros/loong"
    bashrc_path="$HOME/.bashrc"
    
    # 添加环境变量
    env_var_key="export PATH=\$env_var_dir:\$PATH"
    if ! grep -qF "$env_var_key" "$bashrc_path"; then
        if ! echo -e "\n$env_var_key\n" >> "$bashrc_path"; then
            echo "写入环境变量失败"
            exit 1
        fi
        echo "环境变量已成功添加到 ~/.bashrc"
    else
        echo "环境变量已存在于 ~/.bashrc 中"
    fi
    
    # 添加AGIROS Loong环境配置
    setup_bash="source /opt/agiros/loong/setup.bash"
    if ! grep -qF "$setup_bash" "$bashrc_path"; then
        if ! echo "$setup_bash" >> "$bashrc_path"; then
            echo "添加AGIROS Loong环境配置到 ~/.bashrc 失败"
            exit 1
        fi
        echo "已将AGIROS Loong环境配置添加到 ~/.bashrc"
        # 在当前会话中加载配置
        source /opt/agiros/loong/setup.bash
        echo "已在当前会话中加载AGIROS Loong环境配置"
    else
        echo "AGIROS Loong环境配置已存在于 ~/.bashrc 中"
    fi
    
    # 提示用户刷新环境
    echo "请执行 'source ~/.bashrc' 或重启终端使环境变量生效"
}

# 显示菜单
show_menu() {
    echo "欢迎使用AGIROS一键安装脚本"
    echo "本脚本将配置清华源，并支持多种安装模式"
    echo "请选择操作："
    echo "0. 仅配置AGIROS环境"
    echo "1. 最小安装模式"
    echo "2. 桌面版基础安装"
    echo "3. 桌面版完整安装"
    echo "4. 全量安装"
    echo "5. 退出"
}

# 主程序
check_permission

show_menu
read -p "请输入选项(0/1/2/3/4/5): " choice

case $choice in
    0)
        read -p "是否更换系统源为清华源（首次使用建议选是） [y/N]: " replace_choice
        replace_choice=${replace_choice:-N}
        if [[ "$replace_choice" == "y" || "$replace_choice" == "Y" ]]; then
            replace_sources_with_tuna
        fi
        
        download_agiros_key
        download_ros_key
        add_agiros_source
        update_package_list
        
        echo "AGIROS环境配置完成"
        exit 0
        ;;
    1|2|3|4)
        read -p "是否更换系统源为清华源（首次使用建议选是） [y/N]: " replace_choice
        replace_choice=${replace_choice:-N}
        if [[ "$replace_choice" == "y" || "$replace_choice" == "Y" ]]; then
            replace_sources_with_tuna
        else
            # 未更换系统源时单独添加ROS源
            ros_source_file="/etc/apt/sources.list.d/ros-fish.list"
            echo "deb [arch=$ros_arch] http://mirrors.tuna.tsinghua.edu.cn/ros/ubuntu/ $os_codename main" | tee "$ros_source_file" > /dev/null
            echo "deb [arch=$ros_arch] http://mirrors.tuna.tsinghua.edu.cn/ros2/ubuntu/ $os_codename main" | tee -a "$ros_source_file" > /dev/null
        fi
        
        download_agiros_key
        download_ros_key
        add_agiros_source
        update_package_list
        
        install_package "$choice"
        configure_environment_variable
        
        echo "安装和配置完成"
        exit 0
        ;;
    5)
        echo "退出脚本"
        exit 0
        ;;
    *)
        echo "无效的选项，脚本退出"
        exit 1
        ;;
esac