#!/bin/bash

# ROS Podman 容器启动脚本 — 用于 rosbag 可视化
# 用法:
#   ./start_ros_container.sh                    # 启动交互式 shell
#   ./start_ros_container.sh play <bag_name>    # 自动播放 bag + 启动 rviz

set -euo pipefail

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; CYAN='\033[0;36m'; NC='\033[0m'

# ─── 配置 ───────────────────────────────────────────────────────────────────
# 使用镜像名而非 ID，更可读；优先用本地 custom 镜像，fallback 到官方
CONTAINER_IMAGE="docker.io/osrf/ros:noetic-desktop-full"
# 如果本地有 custom 镜像，优先用它（可自行修改）
if podman image exists localhost/custom-ros-noetic:v1.0 2>/dev/null; then
    CONTAINER_IMAGE="localhost/custom-ros-noetic:v1.0"
fi

WORKSPACE_DIR="$(pwd)"
MOUNT_POINT="/workspace"
BAG_DIR="$WORKSPACE_DIR/bags"
CONTAINER_NAME="ros_$(date +%s)"

# ─── 前置检查 ───────────────────────────────────────────────────────────────
if ! command -v podman &> /dev/null; then
    echo -e "${RED}错误: 未安装 podman${NC}"; exit 1
fi

if ! podman image exists "$CONTAINER_IMAGE"; then
    echo -e "${YELLOW}正在拉取镜像 $CONTAINER_IMAGE ...${NC}"
    podman pull "$CONTAINER_IMAGE"
fi

# ─── X11 配置 ────────────────────────────────────────────────────────────────
# 你的 DISPLAY=localhost:10.0，容器内需要能访问宿主机 X11 socket
# 方案: host 网络 + 挂载 X11 socket
if [ -z "${DISPLAY:-}" ]; then
    echo -e "${YELLOW}警告: DISPLAY 未设置，GUI 应用可能无法启动${NC}"
fi

# 临时允许本地 X11 连接（仅当前用户）
xhost +SI:localuser:$USER 2>/dev/null || true

# ─── 构建 podman 参数 ──────────────────────────────────────────────────────
ENV_ARGS=(
    -e DISPLAY="$DISPLAY"
    -e QT_X11_NO_MITSHM=1
    -e ROS_DISTRO=noetic
    -e ROS_MASTER_URI=http://localhost:11311
    -e TZ=Asia/Shanghai
)

VOLUME_ARGS=(
    -v "$WORKSPACE_DIR:$MOUNT_POINT:rw"
    -v "/tmp/.X11-unix:/tmp/.X11-unix:rw"
)

# host 网络模式: 容器直接使用宿主机网络栈，X11 和 ros master 通信都正常
NET_ARGS=(--network=host)

# ─── 打印配置 ────────────────────────────────────────────────────────────────
echo -e "${GREEN}========================================${NC}"
echo -e "${GREEN}ROS Podman 容器 — rosbag 可视化${NC}"
echo -e "${GREEN}========================================${NC}"
echo -e "  镜像:       ${CYAN}$CONTAINER_IMAGE${NC}"
echo -e "  工作目录:   $WORKSPACE_DIR → $MOUNT_POINT"
echo -e "  DISPLAY:    $DISPLAY"
echo -e "  网络模式:   host"
echo -e "  Bag 目录:   $BAG_DIR"
echo -e "${GREEN}========================================${NC}"

# ─── 执行 ────────────────────────────────────────────────────────────────────
# 根据参数决定模式
if [ "${1:-}" = "play" ] && [ -n "${2:-}" ]; then
    BAG_FILE="$BAG_DIR/${2}_davis_with_gt.bag"
    if [ ! -f "$BAG_FILE" ]; then
        echo -e "${RED}错误: bag 文件不存在: $BAG_FILE${NC}"
        echo -e "${YELLOW}可用的 bag 文件:${NC}"
        ls -1 "$BAG_DIR"/*.bag 2>/dev/null | xargs -n1 basename | sed 's/_davis_with_gt.bag//'
        exit 1
    fi

    echo -e "${GREEN}启动 rosbag play + rviz 模式${NC}"
    echo -e "  Bag: ${CYAN}$(basename "$BAG_FILE")${NC}"
    echo -e "${YELLOW}容器启动后会自动播放 bag 并打开 rviz${NC}"
    echo -e "${YELLOW}按 Ctrl+C 停止${NC}"
    echo -e "${GREEN}========================================${NC}"

    podman run -it --rm --name "$CONTAINER_NAME" \
        "${VOLUME_ARGS[@]}" \
        "${ENV_ARGS[@]}" \
        "${NET_ARGS[@]}" \
        "$CONTAINER_IMAGE" \
        /bin/bash -c "
            set -e
            source /opt/ros/noetic/setup.bash
            echo '启动 roscore...'
            roscore &
            ROSCORE_PID=\$!
            sleep 2

            echo '播放 bag: ${2}_davis_with_gt.bag'
            rosbag play --clock \"$MOUNT_POINT/bags/${2}_davis_with_gt.bag\" &
            BAG_PID=\$!
            sleep 1

            echo '启动 rviz...'
            rviz -d \"$MOUNT_POINT/visualize/default.rviz\" 2>/dev/null || rviz

            # 等待 bag 播放完成
            wait \$BAG_PID 2>/dev/null || true
            kill \$ROSCORE_PID 2>/dev/null || true
        "
else
    # 交互式 shell 模式
    echo -e "${GREEN}启动交互式 shell${NC}"
    echo -e "${YELLOW}提示: 进入后执行以下命令播放 bag${NC}"
    echo -e "  source /opt/ros/noetic/setup.bash"
    echo -e "  roscore &"
    echo -e "  rosbag play --clock $MOUNT_POINT/bags/indoor_forward_3_davis_with_gt.bag"
    echo -e "  rviz"
    echo -e "${GREEN}========================================${NC}"

    podman run -it --rm --name "$CONTAINER_NAME" \
        "${VOLUME_ARGS[@]}" \
        "${ENV_ARGS[@]}" \
        "${NET_ARGS[@]}" \
        "$CONTAINER_IMAGE" \
        /bin/bash -c "
            source /opt/ros/noetic/setup.bash
            cd $MOUNT_POINT
            exec /bin/bash -l
        "
fi

# ─── 清理 ────────────────────────────────────────────────────────────────────
xhost -SI:localuser:$USER 2>/dev/null || true
echo -e "${GREEN}容器已退出${NC}"