name: Driver VM Linux

on:
  workflow_call:
    inputs:
      cargo-version:
        required: true
        type: string
      image-tag:
        required: true
        type: string
      checkout-ref:
        required: true
        type: string

permissions:
  contents: read
  packages: read

defaults:
  run:
    shell: bash

jobs:
  download-kernel-runtime:
    name: Download Kernel Runtime
    runs-on: build-amd64
    timeout-minutes: 10
    container:
      image: ghcr.io/nvidia/openshell/ci:latest
      credentials:
        username: ${{ github.actor }}
        password: ${{ secrets.GITHUB_TOKEN }}
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs['checkout-ref'] }}

      - name: Download Linux runtime tarballs
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          set -euo pipefail
          mkdir -p runtime-artifacts

          for platform in linux-aarch64 linux-x86_64; do
            asset="vm-runtime-${platform}.tar.zst"
            echo "Downloading ${asset}..."
            asset_url=$(curl -fsSL \
              -H "Accept: application/vnd.github+json" \
              -H "Authorization: Bearer ${GH_TOKEN}" \
              "https://api.github.com/repos/${GITHUB_REPOSITORY}/releases/tags/vm-dev" \
              | jq -r --arg asset "$asset" '.assets[] | select(.name == $asset) | .browser_download_url' \
              | head -n1)
            if [ -z "$asset_url" ]; then
              echo "::error::No ${asset} asset found on vm-dev release"
              exit 1
            fi
            curl -fL -o "runtime-artifacts/${asset}" "$asset_url"
          done

          ls -lah runtime-artifacts/

      - name: Verify downloads
        run: |
          set -euo pipefail
          for platform in linux-aarch64 linux-x86_64; do
            test -f "runtime-artifacts/vm-runtime-${platform}.tar.zst"
          done

      - name: Upload runtime artifacts
        uses: actions/upload-artifact@v4
        with:
          name: driver-vm-kernel-runtime-tarballs
          path: runtime-artifacts/vm-runtime-*.tar.zst
          retention-days: 1

  build-rootfs:
    name: Build Rootfs (${{ matrix.arch }})
    strategy:
      matrix:
        include:
          - arch: arm64
            runner: build-arm64
            guest_arch: aarch64
          - arch: amd64
            runner: build-amd64
            guest_arch: x86_64
    runs-on: ${{ matrix.runner }}
    timeout-minutes: 30
    container:
      image: ghcr.io/nvidia/openshell/ci:latest
      credentials:
        username: ${{ github.actor }}
        password: ${{ secrets.GITHUB_TOKEN }}
      options: --privileged
      volumes:
        - /var/run/docker.sock:/var/run/docker.sock
    env:
      MISE_GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
      OPENSHELL_IMAGE_TAG: ${{ inputs['image-tag'] }}
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs['checkout-ref'] }}

      - name: Mark workspace safe for git
        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"

      - name: Log in to GHCR
        run: echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u "${{ github.actor }}" --password-stdin

      - name: Install tools
        run: mise install --locked

      - name: Install zstd
        run: apt-get update && apt-get install -y --no-install-recommends zstd && rm -rf /var/lib/apt/lists/*

      - name: Build base rootfs tarball
        run: |
          set -euo pipefail
          crates/openshell-vm/scripts/build-rootfs.sh \
            --base \
            --arch ${{ matrix.guest_arch }} \
            target/rootfs-build

          mkdir -p target/vm-runtime-compressed
          tar -C target/rootfs-build -cf - . \
            | zstd -19 -T0 -o target/vm-runtime-compressed/rootfs.tar.zst

      - name: Upload rootfs artifact
        uses: actions/upload-artifact@v4
        with:
          name: driver-vm-rootfs-${{ matrix.arch }}
          path: target/vm-runtime-compressed/rootfs.tar.zst
          retention-days: 1

  build-driver-vm-linux:
    name: Build Driver VM (Linux ${{ matrix.arch }})
    needs: [download-kernel-runtime, build-rootfs]
    strategy:
      matrix:
        include:
          - arch: arm64
            runner: build-arm64
            target: aarch64-unknown-linux-gnu
            platform: linux-aarch64
          - arch: amd64
            runner: build-amd64
            target: x86_64-unknown-linux-gnu
            platform: linux-x86_64
    runs-on: ${{ matrix.runner }}
    timeout-minutes: 30
    container:
      image: ghcr.io/nvidia/openshell/ci:latest
      credentials:
        username: ${{ github.actor }}
        password: ${{ secrets.GITHUB_TOKEN }}
      options: --privileged
    env:
      MISE_GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
      SCCACHE_MEMCACHED_ENDPOINT: ${{ vars.SCCACHE_MEMCACHED_ENDPOINT }}
      OPENSHELL_IMAGE_TAG: ${{ inputs['image-tag'] }}
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs['checkout-ref'] }}
          fetch-depth: 0

      - name: Mark workspace safe for git
        run: git config --global --add safe.directory "$GITHUB_WORKSPACE"

      - name: Fetch tags
        run: git fetch --tags --force

      - name: Install tools
        run: mise install --locked

      - name: Cache Rust target and registry
        uses: Swatinem/rust-cache@779680da715d629ac1d338a641029a2f4372abb5 # v2
        with:
          shared-key: driver-vm-linux-${{ matrix.arch }}
          cache-directories: .cache/sccache
          cache-targets: "true"

      - name: Install zstd
        run: apt-get update && apt-get install -y --no-install-recommends zstd && rm -rf /var/lib/apt/lists/*

      - name: Download kernel runtime tarball
        uses: actions/download-artifact@v4
        with:
          name: driver-vm-kernel-runtime-tarballs
          path: runtime-download/

      - name: Download rootfs tarball
        uses: actions/download-artifact@v4
        with:
          name: driver-vm-rootfs-${{ matrix.arch }}
          path: rootfs-download/

      - name: Stage compressed runtime for embedding
        run: |
          set -euo pipefail
          COMPRESSED_DIR="${PWD}/target/vm-runtime-compressed"
          mkdir -p "$COMPRESSED_DIR"

          EXTRACT_DIR=$(mktemp -d)
          zstd -d "runtime-download/vm-runtime-${{ matrix.platform }}.tar.zst" --stdout \
            | tar -xf - -C "$EXTRACT_DIR"

          for file in "$EXTRACT_DIR"/*; do
            [ -f "$file" ] || continue
            name=$(basename "$file")
            [ "$name" = "provenance.json" ] && continue
            zstd -19 -f -q -T0 -o "${COMPRESSED_DIR}/${name}.zst" "$file"
          done

          cp rootfs-download/rootfs.tar.zst "${COMPRESSED_DIR}/rootfs.tar.zst"
          ls -lah "$COMPRESSED_DIR"

      - name: Scope workspace to driver-vm crates
        run: |
          set -euo pipefail
          sed -i 's|members = \["crates/\*"\]|members = ["crates/openshell-driver-vm", "crates/openshell-core"]|' Cargo.toml

      - name: Patch workspace version
        if: ${{ inputs['cargo-version'] != '' }}
        run: |
          set -euo pipefail
          sed -i -E '/^\[workspace\.package\]/,/^\[/{s/^version[[:space:]]*=[[:space:]]*".*"/version = "'"${{ inputs['cargo-version'] }}"'"/}' Cargo.toml

      - name: Build openshell-driver-vm
        run: |
          set -euo pipefail
          OPENSHELL_VM_RUNTIME_COMPRESSED_DIR="${PWD}/target/vm-runtime-compressed" \
            mise x -- cargo build --release -p openshell-driver-vm

      - name: sccache stats
        if: always()
        run: mise x -- sccache --show-stats

      - name: Package binary
        run: |
          set -euo pipefail
          mkdir -p artifacts
          tar -czf "artifacts/openshell-driver-vm-${{ matrix.target }}.tar.gz" \
            -C target/release openshell-driver-vm

      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
          name: driver-vm-linux-${{ matrix.arch }}
          path: artifacts/*.tar.gz
          retention-days: 5