name: Build and push all tinygrad models

on:
  workflow_dispatch:
    inputs:
      set_min_version:
        description: 'Minimum selector version required for the models (see helpers.py or readme.md)'
        required: true
        type: string
      target_hardware:
        description: 'Hardware target to compile for (qcom or chestnut)'
        required: true
        type: choice
        default: 'qcom'
        options:
          - qcom
          - chestnut
      hf_repo:
        description: 'Hugging Face dataset repository'
        required: false
        type: string
        default: 'sunnypilot/sunnypilot_models_v1'

jobs:
  setup:
    runs-on: ubuntu-latest
    outputs:
      json_version: ${{ steps.get-json.outputs.json_version }}
      recompiled_dir: ${{ steps.create-recompiled-dir.outputs.recompiled_dir }}
      json_file: ${{ steps.get-json.outputs.json_file }}
      model_matrix: ${{ steps.set-matrix.outputs.model_matrix }}
      tinygrad_ref: ${{ steps.get-tinygrad-ref.outputs.tinygrad_ref }}
    steps:
      - name: Checkout sunnypilot repo
        uses: actions/checkout@v4
        with:
          repository: sunnypilot/sunnypilot
          path: sunnypilot
          submodules: recursive

      - name: Get tinygrad_repo ref
        id: get-tinygrad-ref
        run: |
          cd sunnypilot
          export PYTHONPATH=$(pwd)
          ref=$(python3 openpilot/sunnypilot/models/tinygrad_ref.py)
          echo "tinygrad_ref=$ref" >> $GITHUB_OUTPUT
          echo "tinygrad_ref is $ref"

      - name: Checkout docs repo (sunnypilot-models, gh-pages)
        uses: actions/checkout@v4
        with:
          repository: sunnypilot/sunnypilot-models
          ref: gh-pages
          path: docs
          ssh-key: ${{ secrets.CI_SUNNYPILOT_DOCS_PRIVATE_KEY }}

      - name: Get next JSON version to use (from GitHub docs repo)
        id: get-json
        run: |
          cd docs/docs
          PREFIX="driving_models_${{ inputs.target_hardware == 'chestnut' && 'chestnut_' || '' }}v"
          latest=$(ls ${PREFIX}*.json | sed -E "s/${PREFIX}([0-9]+)\.json/\1/" | sort -n | tail -1)
          next=$((latest+1))
          json_file="${PREFIX}${next}.json"
          cp "${PREFIX}${latest}.json" "$json_file"
          echo "json_file=docs/docs/$json_file" >> $GITHUB_OUTPUT
          echo "json_version=$((next+0))" >> $GITHUB_OUTPUT
          echo "SRC_JSON_FILE=docs/docs/${PREFIX}${latest}.json" >> $GITHUB_ENV

      - name: Extract tinygrad models
        id: set-matrix
        working-directory: docs/docs
        run: |
          jq -c '[.bundles[] | select(.runner=="tinygrad") | {ref, display_name: (.display_name | gsub(" \\([^)]*\\)"; "")), is_20hz}]' "$(basename "${SRC_JSON_FILE}")" > matrix.json
          echo "model_matrix=$(cat matrix.json)" >> $GITHUB_OUTPUT

      - name: Get next recompiled dir number
        id: create-recompiled-dir
        env:
          HF_REPO: ${{ github.event.inputs.hf_repo }}
        run: |
          pip install huggingface_hub
          recompiled_dir=$(python3 -c "
          from huggingface_hub import HfApi
          import re, sys
          api = HfApi()
          files = api.list_repo_files(repo_id=sys.argv[1], repo_type='dataset')
          dirs = [re.search(r'models/recompiled([0-9]+)', f) for f in files]
          nums = [int(m.group(1)) for m in dirs if m]
          print(max(nums) + 1)
          " "$HF_REPO")
          echo "recompiled_dir=$recompiled_dir" >> $GITHUB_OUTPUT

      - name: Push new JSON to GitHub docs repo
        run: |
          cd docs
          git pull origin gh-pages
          git add docs/"$(basename ${{ steps.get-json.outputs.json_file }})"
          git config --global user.name "GitHub Action"
          git config --global user.email "action@github.com"
          git commit -m "Add new ${{ steps.get-json.outputs.json_file }} for build-all" || echo "No changes to commit"
          git push origin gh-pages

  get_and_build:
    needs: [setup]
    strategy:
      matrix:
        model: ${{ fromJson(needs.setup.outputs.model_matrix) }}
      fail-fast: false
    uses: ./.github/workflows/build-single-tinygrad-model.yaml
    with:
      upstream_branch: ${{ matrix.model.ref }}
      custom_name: ${{ matrix.model.display_name }}
      is_20hz: ${{ matrix.model.is_20hz }}
      recompiled_dir: ${{ needs.setup.outputs.recompiled_dir }}
      json_version: ${{ needs.setup.outputs.json_version }}
      target_hardware: ${{ github.event.inputs.target_hardware }}
      hf_repo: ${{ github.event.inputs.hf_repo }}
      set_min_version: ${{ github.event.inputs.set_min_version }}
      tinygrad_ref: ${{ needs.setup.outputs.tinygrad_ref }}
    secrets: inherit

  retry_failed_models:
    needs: [setup, get_and_build]
    runs-on: ubuntu-latest
    if: ${{ !cancelled() && needs.setup.result == 'success' && (needs.get_and_build.result == 'success' || needs.get_and_build.result == 'failure') }}
    outputs:
      retry_matrix: ${{ steps.set-retry-matrix.outputs.retry_matrix }}
    steps:
      - uses: actions/download-artifact@v4
        with:
          pattern: artifact-name-*
          path: output
        continue-on-error: true

      - id: set-retry-matrix
        run: |
          echo '${{ needs.setup.outputs.model_matrix }}' > matrix.json
          built=(); while IFS= read -r line; do [ -n "$line" ] && built+=("$line"); done < <(
            find output -maxdepth 1 -name 'artifact-name-*' -printf "%f\n" 2>/dev/null | sed -E 's/^artifact-name-//' | awk '{gsub(/^ +| +$/, ""); print}'
          )
          jq -c --argjson built "$(printf '%s\n' "${built[@]}" | jq -R . | jq -s .)" \
            'map(select(.display_name as $n | ($built | index($n | gsub("^ +| +$"; "")) | not)))' matrix.json > retry_matrix.json
          echo "retry_matrix=$(cat retry_matrix.json)" >> $GITHUB_OUTPUT

  retry_get_and_build:
    needs: [setup, get_and_build, retry_failed_models]
    if: ${{ !cancelled() && needs.retry_failed_models.result == 'success' && needs.retry_failed_models.outputs.retry_matrix != '[]' && needs.retry_failed_models.outputs.retry_matrix != '' }}
    strategy:
      matrix:
        model: ${{ fromJson(needs.retry_failed_models.outputs.retry_matrix) }}
      fail-fast: false
    uses: ./.github/workflows/build-single-tinygrad-model.yaml
    with:
      upstream_branch: ${{ matrix.model.ref }}
      custom_name: ${{ matrix.model.display_name }}
      is_20hz: ${{ matrix.model.is_20hz }}
      recompiled_dir: ${{ needs.setup.outputs.recompiled_dir }}
      json_version: ${{ needs.setup.outputs.json_version }}
      target_hardware: ${{ github.event.inputs.target_hardware }}
      artifact_suffix: -retry
      hf_repo: ${{ github.event.inputs.hf_repo }}
      set_min_version: ${{ github.event.inputs.set_min_version }}
      tinygrad_ref: ${{ needs.setup.outputs.tinygrad_ref }}
    secrets: inherit