name: Build Single Tinygrad Model and Push

on:
  workflow_dispatch:
    inputs:
      build_model_ref:
        description: 'Branch to use for build-model workflow'
        required: false
        default: 'master-new'
        type: string
      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. 3 for recompiled3)'
        required: true
        type: number
      json_version:
        description: 'driving_models version number to update (e.g. 5 for driving_models_v5.json)'
        required: true
        type: number
      model_folder:
        description: 'Model folder'
        required: true
        type: choice
        options:
          - Simple Plan Models
          - TR Models
          - DTR Models
          - Custom Merge Models
          - FOF series models
          - Other
      custom_model_folder:
        description: 'Custom model folder name (if "Other" selected)'
        required: false
        type: string
      generation:
        description: 'Model generation'
        required: false
        type: number
      version:
        description: 'Minimum selector version'
        required: false
        type: number

jobs:
  build-single:
    runs-on: ubuntu-latest
    env:
      RECOMPILED_DIR: recompiled${{ github.event.inputs.recompiled_dir }}
      JSON_FILE: docs/docs/driving_models_v${{ github.event.inputs.json_version }}.json
    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
          echo "checkout models/${RECOMPILED_DIR}"
          git sparse-checkout set --no-cone models/${RECOMPILED_DIR}
          git checkout main
          cd ..

      - 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: Validate recompiled dir and JSON version
        run: |
          if [ ! -d "gitlab_docs/models/$RECOMPILED_DIR" ]; then
            echo "Recompiled dir $RECOMPILED_DIR does not exist in GitLab repo"
            exit 1
          fi
          if [ ! -f "$JSON_FILE" ]; then
            echo "JSON file $JSON_FILE does not exist!"
            exit 1
          fi

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

      - name: Build model
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          START_TIME=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
          gh workflow run sunnypilot-build-model.yaml \
            --repo sunnypilot/sunnypilot \
            --ref "${{ github.event.inputs.build_model_ref }}" \
            -f upstream_branch="${{ github.event.inputs.upstream_branch }}" \
            -f custom_name="${{ github.event.inputs.custom_name }}"

          for i in {1..24}; do
            RUN_ID=$(gh run list --repo sunnypilot/sunnypilot --workflow=sunnypilot-build-model.yaml --branch="${{ github.event.inputs.build_model_ref }}" --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 "Could not find the triggered workflow run."
            exit 1
          fi

          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')
          echo "Run concluded with: $CONCLUSION"
          if [[ "$CONCLUSION" != "success" ]]; then
            echo "Workflow run failed with conclusion: $CONCLUSION"
            exit 1
          fi

      - name: Download and extract model artifact
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          ARTIFACT_DIR="gitlab_docs/models/$RECOMPILED_DIR"
          RUN_ID=$(gh run list --repo sunnypilot/sunnypilot --workflow=sunnypilot-build-model.yaml --branch="${{ github.event.inputs.build_model_ref }}" --limit=1 --json databaseId --jq '.[0].databaseId')
          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"
          fi
          echo "Done processing $ARTIFACT_NAME"

      - name: Push recompiled dir to GitLab
        env:
          GITLAB_SSH_PRIVATE_KEY: ${{ secrets.GITLAB_SSH_PRIVATE_KEY }}
        run: |
          cd gitlab_docs
          git checkout main
          for d in models/"$RECOMPILED_DIR"/*/; do
            git sparse-checkout add "$d"
          done
          git add models/"$RECOMPILED_DIR"
          git config --global user.name "GitHub Action"
          git config --global user.email "action@github.com"
          git commit -m "Update $RECOMPILED_DIR with new/updated model from build-single-tinygrad-model" || echo "No changes to commit"
          git push origin main

      - name: Run json_parser.py to update JSON
        run: |
          FOLDER="${{ github.event.inputs.model_folder }}"
          if [ "$FOLDER" = "Other" ]; then
            FOLDER="${{ github.event.inputs.custom_model_folder }}"
          fi
          ARGS=""
          [ -n "$FOLDER" ] && ARGS="$ARGS --model-folder \"$FOLDER\""
          [ -n "${{ github.event.inputs.generation }}" ] && ARGS="$ARGS --generation \"${{ github.event.inputs.generation }}\""
          [ -n "${{ github.event.inputs.version }}" ] && ARGS="$ARGS --version \"${{ github.event.inputs.version }}\""
          eval python3 docs/json_parser.py \
            --json-path "$JSON_FILE" \
            --recompiled-dir "gitlab_docs/models/$RECOMPILED_DIR" \
            $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