name: Build All Tinygrad Models and Push to GitLab

on:
  workflow_dispatch:
    inputs:
      branch:
        description: 'Branch to run workflow from'
        required: false
        default: 'master-new'
        type: string

jobs:
  setup:
    runs-on: ubuntu-latest
    outputs:
      json_file: ${{ steps.get-json.outputs.json_file }}
    steps:
      - name: Checkout docs repo
        uses: actions/checkout@v4
        with:
          repository: sunnypilot/sunnypilot-docs
          ref: gh-pages
          path: docs
          ssh-key: ${{ secrets.CI_SUNNYPILOT_DOCS_PRIVATE_KEY }}

      - name: Get next JSON version to use
        id: get-json
        run: |
          cd docs/docs
          latest=$(ls driving_models_v*.json | sed -E 's/.*_v([0-9]+)\.json/\1/' | sort -n | tail -1)
          next=$((latest+1))
          json_file="driving_models_v${next}.json"
          cp "driving_models_v${latest}.json" "$json_file"
          echo "json_file=$json_file" >> $GITHUB_OUTPUT

      - name: Upload context for next jobs
        uses: actions/upload-artifact@v4
        with:
          name: context
          path: docs

  build-all:
    runs-on: ubuntu-latest
    needs: setup
    env:
      JSON_FILE: docs/docs/${{ needs.setup.outputs.json_file }}
    steps:
      - name: Set up SSH
        uses: webfactory/ssh-agent@v0.9.0
        with:
          ssh-private-key: ${{ secrets.GITLAB_SSH_PRIVATE_KEY }}

      - name: Add GitLab.com SSH key to known_hosts
        run: |
          mkdir -p ~/.ssh
          ssh-keyscan -H gitlab.com >> ~/.ssh/known_hosts

      - name: Clone GitLab docs repo
        env:
          GIT_SSH_COMMAND: 'ssh -o UserKnownHostsFile=~/.ssh/known_hosts'
        run: |
          echo "Cloning GitLab"
          git clone --depth 1 --filter=tree:0 --sparse git@gitlab.com:sunnypilot/public/docs.sunnypilot.ai2.git gitlab_docs
          cd gitlab_docs
          git checkout main
          cd ..

      - name: Set next recompiled dir
        id: set-recompiled
        run: |
          cd gitlab_docs/models
          latest_dir=$(ls -d recompiled* 2>/dev/null | sed -E 's/recompiled([0-9]+)/\1/' | sort -n | tail -1)
          if [[ -z "$latest_dir" ]]; then
            next_dir=1
          else
            next_dir=$((latest_dir+1))
          fi
          recompiled_dir="recompiled${next_dir}"
          mkdir -p "$recompiled_dir"
          echo "RECOMPILED_DIR=$recompiled_dir" >> $GITHUB_ENV

      - name: Download context
        uses: actions/download-artifact@v4
        with:
          name: context
          path: .

      - name: Install dependencies
        run: |
          sudo apt-get update
          sudo apt-get install -y jq gh

      - name: Build all tinygrad models
        id: trigger-builds
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          set -e
          > triggered_run_ids.txt
          BRANCH="${{ github.event.inputs.branch }}"
          jq -c '.bundles[] | select(.runner=="tinygrad")' "$JSON_FILE" | while read -r bundle; do
            ref=$(echo "$bundle" | jq -r '.ref')
            display_name=$(echo "$bundle" | jq -r '.display_name' | sed 's/ ([^)]*)//g')
            is_20hz=$(echo "$bundle" | jq -r '.is_20hz')
            echo "Triggering build for: $display_name ($ref) [20Hz: $is_20hz]"
            START_TIME=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
            gh workflow run sunnypilot-build-model.yaml \
              --repo sunnypilot/sunnypilot \
              --ref "$BRANCH" \
              -f upstream_branch="$ref" \
              -f custom_name="$display_name" \
              -f is_20hz="$is_20hz"
            for i in {1..24}; do
              RUN_ID=$(gh run list --repo sunnypilot/sunnypilot --workflow=sunnypilot-build-model.yaml --branch="$BRANCH" --created ">$START_TIME" --limit=1 --json databaseId --jq '.[0].databaseId')
              if [ -n "$RUN_ID" ]; then
                break
              fi
              sleep 5
            done
            if [ -z "$RUN_ID" ]; then
              echo "ould not find the triggered workflow run for $display_name ($ref)"
              exit 1
            fi
            echo "$RUN_ID" >> triggered_run_ids.txt
          done

      - name: Wait for all model builds to finish
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          set -e
          SUCCESS_RUNS=()
          FAILED_RUNS=()
          declare -A RUN_ID_TO_NAME

          while read -r RUN_ID; do
            ARTIFACT_NAME=$(gh api repos/sunnypilot/sunnypilot/actions/runs/$RUN_ID/artifacts --jq '.artifacts[] | select(.name | startswith("model-")) | .name' || echo "unknown")
            RUN_ID_TO_NAME["$RUN_ID"]="$ARTIFACT_NAME"
          done < triggered_run_ids.txt

          while read -r RUN_ID; do
            echo "Watching run ID: $RUN_ID"
            gh run watch "$RUN_ID" --repo sunnypilot/sunnypilot
            CONCLUSION=$(gh run view "$RUN_ID" --repo sunnypilot/sunnypilot --json conclusion --jq '.conclusion')
            ARTIFACT_NAME="${RUN_ID_TO_NAME[$RUN_ID]}"
            echo "Run $RUN_ID ($ARTIFACT_NAME) concluded with: $CONCLUSION"
            if [[ "$CONCLUSION" == "success" ]]; then
              SUCCESS_RUNS+=("$RUN_ID")
            else
              FAILED_RUNS+=("$RUN_ID")
            fi
          done < triggered_run_ids.txt

          if [[ ${#SUCCESS_RUNS[@]} -eq 0 ]]; then
            echo "All model builds failed. Aborting."
            exit 1
          fi

          if [[ ${#FAILED_RUNS[@]} -gt 0 ]]; then
            echo "WARNING: The following model builds failed:"
            for RUN_ID in "${FAILED_RUNS[@]}"; do
              echo "- $RUN_ID (${RUN_ID_TO_NAME[$RUN_ID]})"
            done
            echo "You may want to rerun these models manually."
          fi

      - name: Download and extract all model artifacts
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          ARTIFACT_DIR="gitlab_docs/models/$RECOMPILED_DIR"
          SUCCESS_RUNS=()
          while read -r RUN_ID; do
            CONCLUSION=$(gh run view "$RUN_ID" --repo sunnypilot/sunnypilot --json conclusion --jq '.conclusion')
            if [[ "$CONCLUSION" == "success" ]]; then
              SUCCESS_RUNS+=("$RUN_ID")
            fi
          done < triggered_run_ids.txt

          for RUN_ID in "${SUCCESS_RUNS[@]}"; do
            ARTIFACT_NAME=$(gh api repos/sunnypilot/sunnypilot/actions/runs/$RUN_ID/artifacts --jq '.artifacts[] | select(.name | startswith("model-")) | .name')
            echo "Downloading artifact: $ARTIFACT_NAME from run: $RUN_ID"
            mkdir -p "$ARTIFACT_DIR/$ARTIFACT_NAME"
            echo "Created directory: $ARTIFACT_DIR/$ARTIFACT_NAME"
            gh run download "$RUN_ID" --repo sunnypilot/sunnypilot -n "$ARTIFACT_NAME" --dir "$ARTIFACT_DIR/$ARTIFACT_NAME"
            echo "Downloaded artifact zip(s) to: $ARTIFACT_DIR/$ARTIFACT_NAME"
            ZIP_PATH=$(find "$ARTIFACT_DIR/$ARTIFACT_NAME" -type f -name '*.zip' | head -n1)
            if [ -n "$ZIP_PATH" ]; then
              echo "Unzipping $ZIP_PATH to $ARTIFACT_DIR/$ARTIFACT_NAME"
              unzip -o "$ZIP_PATH" -d "$ARTIFACT_DIR/$ARTIFACT_NAME"
              rm -f "$ZIP_PATH"
              echo "Unzipped and removed $ZIP_PATH"
            else
              echo "No zip file found in $ARTIFACT_DIR/$ARTIFACT_NAME (This is NOT an error)."
            fi
            echo "Done processing $ARTIFACT_NAME"
          done

      - name: Push recompiled dir to GitLab
        env:
          GITLAB_SSH_PRIVATE_KEY: ${{ secrets.GITLAB_SSH_PRIVATE_KEY }}
        run: |
          cd gitlab_docs
          git checkout main
          mkdir -p models/"$(basename $RECOMPILED_DIR)"
          git add models/"$(basename $RECOMPILED_DIR)"
          git config --global user.name "GitHub Action"
          git config --global user.email "action@github.com"
          git commit -m "Add $(basename $RECOMPILED_DIR) from build-all-tinygrad-models"
          git push origin main

      - name: Run json_parser.py to update JSON
        run: |
          python3 docs/json_parser.py \
            --json-path "$JSON_FILE" \
            --recompiled-dir "gitlab_docs/models/$RECOMPILED_DIR"

      - name: Push updated JSON to GitHub docs repo
        run: |
          cd docs
          git config --global user.name "GitHub Action"
          git config --global user.email "action@github.com"
          git checkout gh-pages
          git add docs/"$(basename $JSON_FILE)"
          git commit -m "Update $(basename $JSON_FILE) after recompiling models" || echo "No changes to commit"
          git push origin gh-pages