name: Build Single Tinygrad Model and Push

on:
  workflow_call:
    inputs:
      upstream_branch:
        description: 'Upstream commit to build from'
        required: true
        type: string
      custom_name:
        description: 'Custom name for the model (no date, only name)'
        required: false
        type: string
      recompiled_dir:
        description: 'Existing recompiled directory number (e.g. 1 for recompiled1)'
        required: true
        type: string
      json_version:
        description: 'driving_models version number to update (e.g. 18 for driving_models_v18.json)'
        required: true
        type: string
      artifact_suffix:
        description: 'Suffix for artifact name'
        required: false
        type: string
        default: ''
      is_20hz:
        description: 'Is this a 20Hz model'
        required: false
        type: boolean
        default: true
      target_hardware:
        description: 'Hardware target to compile for (qcom or chestnut)'
        required: false
        type: string
        default: 'qcom'
      hf_repo:
        description: 'Hugging Face dataset repository (e.g. sunnypilot/sunnypilot_models_v1)'
        required: false
        type: string
        default: 'sunnypilot/sunnypilot_models_v1'
      set_min_version:
        description: 'Minimum selector version'
        required: false
        type: string
      tinygrad_ref:
        description: 'Tinygrad reference'
        required: false
        type: string
  workflow_dispatch:
    inputs:
      upstream_branch:
        description: 'Upstream commit to build from'
        required: true
        type: string
      custom_name:
        description: 'Custom name for the model (no date, only name)'
        required: false
        type: string
      is_20hz:
        description: 'Is this a 20Hz model'
        required: false
        type: boolean
        default: true
      recompiled_dir:
        description: 'Existing recompiled directory number (e.g. 3 for recompiled3)'
        required: true
        type: string
      json_version:
        description: 'driving_models version number to update (e.g. 5 for driving_models_v5.json)'
        required: true
        type: string
      model_folder:
        description: 'Model folder'
        type: choice
        default: 'None'
        options:
          - None
          - Master Models
          - Release Models
          - 2026 World Models
          - 2026 Deep RL Models
          - Custom Merge Models
          - Other
      custom_model_folder:
        description: 'Custom model folder name (if "Other" selected)'
        required: false
        type: string
      generation:
        description: 'Model generation'
        required: false
        type: string
      version:
        description: 'Minimum selector version'
        required: false
        type: string
      target_hardware:
        description: 'Hardware target to compile for'
        required: false
        type: choice
        default: 'qcom'
        options:
          - qcom
          - chestnut
      hf_repo:
        description: 'Hugging Face dataset repository'
        required: false
        type: string
        default: 'sunnypilot/sunnypilot_models_v1'
env:
  RECOMPILED_DIR: recompiled${{ inputs.recompiled_dir }}
  JSON_FILE: docs/docs/driving_models_${{ inputs.target_hardware == 'chestnut' && 'chestnut_v' || 'v' }}${{ inputs.json_version }}.json

jobs:
  build_model:
    uses: ./.github/workflows/sunnypilot-build-model.yaml
    with:
      upstream_branch: ${{ inputs.upstream_branch }}
      custom_name: ${{ inputs.custom_name || inputs.upstream_branch }}
      is_20hz: ${{ inputs.is_20hz }}
      artifact_suffix: ${{ inputs.artifact_suffix }}
      target_hardware: ${{ inputs.target_hardware }}
    secrets: inherit

  publish_model:
    if: ${{ !cancelled() && needs.build_model.result == 'success' }}
    concurrency:
      group: hf-push-${{ inputs.recompiled_dir }}
      cancel-in-progress: false
    needs: build_model
    runs-on: ubuntu-latest
    permissions:
      id-token: write
      contents: write
    steps:
      - name: Checkout docs repo
        uses: actions/checkout@v4
        with:
          repository: sunnypilot/sunnypilot-models
          ref: gh-pages
          path: docs
          ssh-key: ${{ secrets.CI_SUNNYPILOT_DOCS_PRIVATE_KEY }}

      - name: Install huggingface_hub
        run: pip install --upgrade "huggingface_hub>=0.22.0"

      - name: Validate hf_repo and JSON version
        env:
          HF_OIDC_RESOURCE: datasets/${{ inputs.hf_repo }}
        run: |
          if [ ! -f "$JSON_FILE" ]; then
            echo "JSON file $JSON_FILE does not exist!"
            exit 1
          fi
          python3 -c "
          import sys
          from huggingface_hub import HfApi
          try:
              api = HfApi()
              api.repo_info(repo_id=sys.argv[1], repo_type='dataset')
              print(f'Success: Repo {sys.argv[1]} exists.')
          except Exception as e:
              print('HF validation failed:', e)
              sys.exit(1)
          " "${{ inputs.hf_repo }}"

      - name: Download artifact name file
        uses: actions/download-artifact@v4
        with:
          name: artifact-name-${{ inputs.custom_name || inputs.upstream_branch }}
          path: artifact_name

      - name: Read artifact name
        id: read-artifact-name
        run: |
          ARTIFACT_NAME=$(cat artifact_name/artifact_name.txt)
          echo "artifact_name=$ARTIFACT_NAME" >> $GITHUB_OUTPUT

      - name: Download and extract model artifact
        uses: actions/download-artifact@v4
        with:
          name: ${{ steps.read-artifact-name.outputs.artifact_name }}
          path: output

      - name: Create models folder
        env:
          ARTIFACT_NAME: ${{ steps.read-artifact-name.outputs.artifact_name }}
        run: |
          mkdir -p "local_models/${RECOMPILED_DIR}/${ARTIFACT_NAME}"
          cp -r output/* "local_models/${RECOMPILED_DIR}/${ARTIFACT_NAME}/"
          rm -f "local_models/${RECOMPILED_DIR}/${ARTIFACT_NAME}/artifact_name.txt"

      - name: Upload to Hugging Face
        env:
          HF_OIDC_RESOURCE: datasets/${{ inputs.hf_repo }}
          ARTIFACT_NAME: ${{ steps.read-artifact-name.outputs.artifact_name }}
        run: |
          hf upload ${{ inputs.hf_repo }} \
            output/ \
            "models/${RECOMPILED_DIR}/${ARTIFACT_NAME}/" \
            --repo-type=dataset

      - name: Pull gh-pages
        run: |
          cd docs
          git pull origin gh-pages

      - name: Run json_parser.py to update JSON
        run: |
          FOLDER="${{ inputs.model_folder }}"
          if [ "$FOLDER" = "Other" ]; then
            FOLDER="${{ inputs.custom_model_folder }}"
          fi
          ARGS=""
          if [ "$FOLDER" != "None" ] && [ -n "$FOLDER" ]; then
            ARGS="$ARGS --model-folder \"$FOLDER\""
          fi
          [ -n "${{ inputs.generation }}" ] && ARGS="$ARGS --generation \"${{ inputs.generation }}\""
          [ -n "${{ inputs.version }}" ] && ARGS="$ARGS --version \"${{ inputs.version }}\""
          [ -n "${{ inputs.set_min_version }}" ] && ARGS="$ARGS --set-min-version \"${{ inputs.set_min_version }}\""
          [ -n "${{ inputs.tinygrad_ref }}" ] && ARGS="$ARGS --tinygrad-ref \"${{ inputs.tinygrad_ref }}\""
          eval python3 docs/json_parser.py \
            --json-path "$JSON_FILE" \
            --recompiled-dir "local_models/$RECOMPILED_DIR" \
            --sort-by-date \
            $ARGS

      - 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 model" || echo "No changes to commit"
          git push origin gh-pages