FROM nvgear/ros-2:latest

# Accept build argument for username
ARG USERNAME
ARG USERID
ARG HOME_DIR
ARG WORKTREE_NAME

# Create user with the same name as host
RUN if [ "$USERID" != "0" ]; then \
        useradd -m -u ${USERID} -s /bin/bash ${USERNAME} && \
        echo "${USERNAME} ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers && \
        # Add user to video and render groups for GPU access
        usermod -a -G video,render ${USERNAME} || true; \
    fi

# Copy .bashrc with color settings before switching user
COPY --chown=${USERNAME}:${USERNAME} docker/.bashrc ${HOME_DIR}/.bashrc

# Install Manus udev rules
COPY --chown=${USERNAME}:${USERNAME} docker/70-manus-hid.rules /etc/udev/rules.d/70-manus-hid.rules

# Copy tmux configuration
COPY --chown=${USERNAME}:${USERNAME} docker/.tmux.conf ${HOME_DIR}/.tmux.conf

# Switch to user
USER ${USERNAME}

# Install tmux plugin manager and uv in parallel
RUN git clone https://github.com/tmux-plugins/tpm ${HOME_DIR}/.tmux/plugins/tpm & \
    curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=${HOME_DIR}/.cargo/bin sh & \
    wait

# Install tmux plugins automatically
RUN ${HOME_DIR}/.tmux/plugins/tpm/bin/install_plugins || true

# Add uv to PATH
ENV PATH="${HOME_DIR}/.cargo/bin:$PATH"
ENV UV_PYTHON=${HOME_DIR}/venv/bin/python

# Create venv
RUN uv venv --python 3.10 ${HOME_DIR}/venv

# Install hardware-specific packages (x86 only - not available on ARM64/Orin)
USER root
COPY --chown=${USERNAME}:${USERNAME} gr00t_wbc/control/teleop/device/pico/XRoboToolkit_PC_Service_1.0.0_ubuntu_22.04_amd64.deb ${HOME_DIR}/XRoboToolkit_PC_Service_1.0.0_ubuntu_22.04_amd64.deb
COPY --chown=${USERNAME}:${USERNAME} gr00t_wbc/control/teleop/device/pico/roboticsservice_1.0.0.0_arm64.deb ${HOME_DIR}/roboticsservice_1.0.0.0_arm64.deb

RUN if [ "$(dpkg --print-architecture)" = "amd64" ]; then \
        # Ultra Leap setup
        wget -qO - https://repo.ultraleap.com/keys/apt/gpg | gpg --dearmor | tee /etc/apt/trusted.gpg.d/ultraleap.gpg && \
        echo 'deb [arch=amd64] https://repo.ultraleap.com/apt stable main' | tee /etc/apt/sources.list.d/ultraleap.list && \
        apt-get update && \
        echo "yes" | DEBIAN_FRONTEND=noninteractive apt-get install -y ultraleap-hand-tracking libhidapi-dev && \
        # Space Mouse udev rules
        echo 'KERNEL=="hidraw*", SUBSYSTEM=="hidraw", MODE="0664", GROUP="plugdev"' > /etc/udev/rules.d/99-hidraw-permissions.rules && \
        usermod -aG plugdev ${USERNAME}; \
        # Pico setup
        apt-get install -y xdg-utils && \
        dpkg -i ${HOME_DIR}/XRoboToolkit_PC_Service_1.0.0_ubuntu_22.04_amd64.deb; \
    else \
        echo "Skipping x86-only hardware packages on $(dpkg --print-architecture)"; \
    fi

USER ${USERNAME}
# Install hardware Python packages (x86 only) with caching
RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
    if [ "$(dpkg --print-architecture)" = "amd64" ]; then \
        # Ultra Leap Python bindings
        git clone https://github.com/ultraleap/leapc-python-bindings ${HOME_DIR}/leapc-python-bindings && \
        cd ${HOME_DIR}/leapc-python-bindings && \
        UV_CONCURRENT_DOWNLOADS=8 uv pip install -r requirements.txt && \
        MAKEFLAGS="-j$(nproc)" ${HOME_DIR}/venv/bin/python -m build leapc-cffi && \
        uv pip install leapc-cffi/dist/leapc_cffi-0.0.1.tar.gz && \
        uv pip install -e leapc-python-api && \
        # Space Mouse Python package
        uv pip install pyspacemouse && \
        # Pico Python bindings
        git clone https://github.com/XR-Robotics/XRoboToolkit-PC-Service-Pybind.git ${HOME_DIR}/XRoboToolkit-PC-Service-Pybind && \
        cd ${HOME_DIR}/XRoboToolkit-PC-Service-Pybind && \
        uv pip install setuptools pybind11 && \
        sed -i "s|pip install|uv pip install|g" setup_ubuntu.sh && \
        sed -i "s|pip uninstall|uv pip uninstall|g" setup_ubuntu.sh && \
        sed -i "s|python setup.py install|${HOME_DIR}/venv/bin/python setup.py install|g" setup_ubuntu.sh && \
        bash setup_ubuntu.sh; \
    fi

# Install Python dependencies using uv with caching
RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
    UV_CONCURRENT_DOWNLOADS=8 uv pip install --upgrade pip ipython jupyter notebook debugpy


# Copy entire project to the workspace directory where it will be mounted at runtime
# NOTE: The build context must be the project root for this to work
# Use dynamic worktree name to match runtime mount path
COPY --chown=${USERNAME}:${USERNAME} . ${HOME_DIR}/Projects/${WORKTREE_NAME}

# Install Python dependencies inside the venv with caching - split into separate commands
RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
    UV_CONCURRENT_DOWNLOADS=8 uv pip install \
    -e ${HOME_DIR}/Projects/${WORKTREE_NAME}/external_dependencies/unitree_sdk2_python

# Unlike pip, uv downloads LFS files by default. There's a bug in uv that causes LFS files
# to fail to download (https://github.com/astral-sh/uv/issues/3312). So we need to set
# UV_GIT_LFS=1 to prevent uv from downloading LFS files.
# Install project dependencies (VLA extras only on x86) with caching
RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
    GIT_LFS_SKIP_SMUDGE=1 UV_CONCURRENT_DOWNLOADS=8 uv pip install -e "${HOME_DIR}/Projects/${WORKTREE_NAME}[dev]"

# Clone and install robosuite with specific branch
RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
    git clone https://github.com/xieleo5/robosuite.git ${HOME_DIR}/robosuite && \
    cd ${HOME_DIR}/robosuite && \
    git checkout leo/support_g1_locomanip && \
    UV_CONCURRENT_DOWNLOADS=8 uv pip install -e .

# Install gr00trobocasa
RUN --mount=type=cache,target=${HOME_DIR}/.cache/uv,uid=${USERID},gid=${USERID} \
    UV_CONCURRENT_DOWNLOADS=8 uv pip install -e ${HOME_DIR}/Projects/${WORKTREE_NAME}/gr00t_wbc/dexmg/gr00trobocasa

# Configure bash environment with virtual environment and ROS2 setup
RUN echo "source ${HOME_DIR}/venv/bin/activate" >> ${HOME_DIR}/.bashrc && \
    echo "source /opt/ros/humble/setup.bash" >> ${HOME_DIR}/.bashrc && \
    echo "export ROS_LOCALHOST_ONLY=1" >> ${HOME_DIR}/.bashrc

# Default command (can be overridden at runtime)
CMD ["/bin/bash"]