# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This workflow handles full testing.
name: Full Tests

on:
  # Allows running this workflow manually from the Actions tab
  workflow_dispatch:

  pull_request_review:
    types: [submitted]
  push:
    branches:
      - main
    paths:
      - "src/**"
      - "tests/**"
      - ".github/workflows/**"
      - "pyproject.toml"
      - "Makefile"

permissions:
  contents: read

# Sets up the environment variables
env:
  UV_VERSION: "0.8.0"
  PYTHON_VERSION: "3.10"
  DOCKER_IMAGE_NAME: huggingface/lerobot-gpu

# Ensures that only the latest action is built, canceling older runs.
concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
  cancel-in-progress: true

jobs:

  # This job runs the E2E tests + pytest with all extras
  # It runs everytime a PR is approved or a push to main
  full-tests:
    name: Full Tests
    runs-on: ubuntu-latest
    if: |
      (github.event_name == 'pull_request_review' && github.event.review.state == 'approved') ||
      github.event_name == 'push' ||
      github.event_name == 'workflow_dispatch'
    env:
      MUJOCO_GL: egl
    steps:
      - uses: actions/checkout@v4
        with:
          lfs: true
          persist-credentials: false

      - name: Install apt dependencies
        run: |
          sudo apt-get update && sudo apt-get install -y build-essential \
          git curl libglib2.0-0 libegl1-mesa-dev ffmpeg libusb-1.0-0-dev \
          speech-dispatcher libgeos-dev portaudio19-dev

      - name: Setup uv and Python
        uses: astral-sh/setup-uv@v6 # zizmor: ignore[unpinned-uses]
        with:
          enable-cache: true
          version: ${{ env.UV_VERSION }}
          python-version: ${{ env.PYTHON_VERSION }}

      - name: Install lerobot with all extras
        run: uv sync --all-extras --no-extra groot # TODO(Steven): Make flash-attn optional

      - name: Run pytest (all extras)
        run: uv run pytest tests -vv --maxfail=10

      - name: Run end-to-end tests
        run: uv run make test-end-to-end

  # This job builds a GPU enabled image for testing
  # It runs everytime a PR is approved or a push to main
  # TODO(Steven): For now we skip this job for community PRs
  build-and-push-docker:
    name: Build and Push Docker
    runs-on:
      group: aws-general-8-plus
    if: |
      (github.event_name == 'pull_request_review' && github.event.review.state == 'approved' && github.event.pull_request.head.repo.fork == false) ||
      github.event_name == 'push' ||
      github.event_name == 'workflow_dispatch'
    outputs:
      image_tag: ${{ steps.set_tag.outputs.image_tag }}
    env:
      GITHUB_EVENT_NAME: ${{ github.event_name }}
      GITHUB_REF: ${{ github.ref }}
      GITHUB_PR_NUMBER: ${{ github.event.pull_request.number }}
    steps:
      - name: Set Docker image tag
        id: set_tag
        run: |
          if [[ "${GITHUB_EVENT_NAME}" == "push" ]]; then
            TAG="${DOCKER_IMAGE_NAME}:latest"
          elif [[ -n "${GITHUB_PR_NUMBER}" ]]; then
            TAG="${DOCKER_IMAGE_NAME}:pr-${GITHUB_PR_NUMBER}"
          else
            TAG="${DOCKER_IMAGE_NAME}:pr-${GITHUB_REF##*/}"
          fi
          echo "image_tag=$TAG" >> $GITHUB_OUTPUT
      - name: Install Git LFS
        run: |
          sudo apt-get update
          sudo apt-get install git-lfs
          git lfs install
      - uses: actions/checkout@v4
        with:
          lfs: true
          persist-credentials: false
      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
        with:
          cache-binary: false
      - name: Login to Docker Hub
        uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
        with:
          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
      - name: Build and push Docker image
        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
        with:
          context: .
          file: ./docker/Dockerfile.internal
          push: true
          tags: ${{ steps.set_tag.outputs.image_tag }}

  # This job runs pytest with all extras in a GPU enabled host
  # It runs everytime a test image is created
  gpu-tests:
    name: GPU Tests
    needs: [build-and-push-docker]
    runs-on:
      group: aws-g6-4xlarge-plus
    env:
      HF_HOME: /home/user_lerobot/.cache/huggingface
      HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot
      TORCH_HOME: /home/user_lerobot/.cache/torch
      TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
    container:
      image: ${{ needs.build-and-push-docker.outputs.image_tag }} # zizmor: ignore[unpinned-images]
      options: --gpus all --shm-size "16gb"
      credentials:
        username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
        password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
    defaults:
      run:
        shell: bash
        working-directory: /lerobot
    steps:
      - name: Run pytest on GPU
        run: pytest tests -vv --maxfail=10
      - name: Run end-to-end tests
        run: make test-end-to-end

  # This job deletes the test image recently created
  # It runs everytime after the gpu-tests have finished
  delete-pr-image:
    name: Delete PR Image
    needs: [gpu-tests, build-and-push-docker]
    if: always() && ((github.event.review.state == 'approved') || (github.event_name == 'workflow_dispatch')) && needs.build-and-push-docker.result == 'success'
    runs-on: ubuntu-latest
    steps:
      - name: Get Docker Hub Token and Delete Image
        # zizmor: ignore[template-injection]
        run: |
          IMAGE_NAME=$(echo "${{ needs.build-and-push-docker.outputs.image_tag }}" | cut -d':' -f1)
          IMAGE_TAG=$(echo "${{ needs.build-and-push-docker.outputs.image_tag }}" | cut -d':' -f2)

          echo "Attempting to delete image: $IMAGE_NAME:$IMAGE_TAG"

          TOKEN=$(curl -s -H "Content-Type: application/json" \
                       -X POST \
                       -d '{"username": "${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}", "password": "${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}"}' \
                       https://hub.docker.com/v2/users/login/ | jq -r .token)

          if [ "$TOKEN" == "null" ] || [ -z "$TOKEN" ]; then
            echo "::error::Failed to get Docker Hub token."
            exit 1
          fi

          HTTP_RESPONSE=$(curl -s -o /dev/null -w "%{http_code}" \
                               -H "Authorization: JWT ${TOKEN}" \
                               -X DELETE \
                               https://hub.docker.com/v2/repositories/${IMAGE_NAME}/tags/${IMAGE_TAG}/)

          if [ "$HTTP_RESPONSE" -eq 204 ]; then
            echo "Successfully deleted Docker image tag: $IMAGE_NAME:$IMAGE_TAG"
          else
            echo "::error::Failed to delete Docker image. HTTP status: $HTTP_RESPONSE"
            exit 1
          fi

# TODO(Steven): Check dockerimages pull in ubuntu