#cloud-config
#
# Benchmark VM provisioning for microsoft/typeagent-py
#
# Structured RAG library (ingest, index, query) -- Python 3.13, uv-based build.
# Clones from KRRT7 fork, installs via uv sync (lockfile-based).
#
# Usage:
#   az vm create ... --custom-data infra/cloud-init.yaml
#
# VM: Azure Standard_D2s_v5 (2 vCPU, 8 GB RAM, non-burstable)
#   Non-burstable ensures consistent CPU -- no thermal throttling or turbo variability.

package_update: true
packages:
  - git
  - build-essential
  - curl
  - wget
  - jq

write_files:
  # --- Benchmark: import time ---
  - path: /home/azureuser/bench/bench_import.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail

      PYTHON=~/typeagent/.venv/bin/python
      OUTDIR="$HOME/results/import"
      mkdir -p "$OUTDIR"

      echo "=== typeagent import benchmark ==="
      echo ""

      # E2E import time (direct venv to avoid uv run overhead)
      hyperfine --warmup 5 --min-runs 30 --shell=none \
        --export-json "$OUTDIR/import.json" \
        "$PYTHON -c 'import typeagent'"

      # Per-module breakdown
      $PYTHON -X importtime -c 'import typeagent' 2>"$OUTDIR/importtime_raw.txt"
      sort -t'|' -k1 -rn "$OUTDIR/importtime_raw.txt" | head -30 > "$OUTDIR/importtime_top30.txt"

      echo ""
      echo "Top 30 imports by self time:"
      cat "$OUTDIR/importtime_top30.txt"

  # --- Benchmark: offline E2E test suite ---
  - path: /home/azureuser/bench/bench_tests.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail

      PYTHON=~/typeagent/.venv/bin/python
      OUTDIR="$HOME/results/tests"
      mkdir -p "$OUTDIR"

      cd ~/typeagent

      echo "=== typeagent offline E2E test benchmark ==="
      echo ""

      # Run offline tests with timing (direct venv to avoid uv run overhead)
      $PYTHON -m pytest \
        tests/test_incremental_index.py \
        tests/test_add_messages_with_indexing.py \
        tests/test_podcast_incremental.py \
        tests/test_sqlite_indexes.py \
        tests/test_query.py \
        --durations=0 -v 2>&1 | tee "$OUTDIR/test_output.txt"

      # Hyperfine on the full offline suite
      hyperfine --warmup 2 --min-runs 10 --shell=none \
        --export-json "$OUTDIR/test_suite.json" \
        "$PYTHON -m pytest tests/test_incremental_index.py tests/test_add_messages_with_indexing.py tests/test_podcast_incremental.py tests/test_sqlite_indexes.py tests/test_query.py -q"

  # --- Benchmark: all baselines (runs import + tests) ---
  - path: /home/azureuser/bench/bench_baseline.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail

      echo "=== Running all baseline benchmarks ==="
      echo ""

      bash ~/bench/bench_import.sh
      echo ""
      bash ~/bench/bench_tests.sh

      echo ""
      echo "=== All baselines complete ==="
      echo "Results in ~/results/"
      ls -R ~/results/

  # --- Benchmark: A/B branch comparison ---
  - path: /home/azureuser/bench/bench_compare.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      BRANCH="${1:?Usage: bench_compare.sh <branch-or-commit>}"
      TS=$(date +%Y%m%d-%H%M%S)
      OUTDIR="$HOME/results/${BRANCH//\//-}-${TS}"
      mkdir -p "$OUTDIR"

      cd ~/typeagent
      git fetch origin
      git checkout "$BRANCH"

      # Rebuild after switching branches
      export PATH="$HOME/.local/bin:$PATH"
      uv sync

      PYTHON=~/typeagent/.venv/bin/python

      echo "=== Benchmarking branch: $BRANCH ==="

      # Import time (direct venv to avoid uv run overhead)
      hyperfine --warmup 5 --min-runs 30 --shell=none \
        --export-json "$OUTDIR/import.json" \
        "$PYTHON -c 'import typeagent'"

      # Offline test suite
      hyperfine --warmup 2 --min-runs 10 --shell=none \
        --export-json "$OUTDIR/test_suite.json" \
        "$PYTHON -m pytest tests/test_incremental_index.py tests/test_add_messages_with_indexing.py tests/test_podcast_incremental.py tests/test_sqlite_indexes.py tests/test_query.py -q"

      echo ""
      echo "Results saved to $OUTDIR/"
      ls -la "$OUTDIR/"

  # --- Benchmark: side-by-side two branches ---
  - path: /home/azureuser/bench/bench_ab.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      BASE="${1:?Usage: bench_ab.sh <base-branch> <opt-branch>}"
      OPT="${2:?Usage: bench_ab.sh <base-branch> <opt-branch>}"

      echo "=== A/B comparison: $BASE vs $OPT ==="
      bash ~/bench/bench_compare.sh "$BASE"
      bash ~/bench/bench_compare.sh "$OPT"

      echo ""
      echo "Compare results in ~/results/"
      ls ~/results/

  # --- Setup script (runs once via runcmd) ---
  - path: /home/azureuser/setup.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      export PATH="$HOME/.local/bin:$PATH"

      echo "=== Cloning typeagent ==="
      git clone https://github.com/KRRT7/typeagent-py.git ~/typeagent
      cd ~/typeagent

      echo "=== Installing toolchain and building ==="
      curl -LsSf https://astral.sh/uv/install.sh | sh
      export PATH="$HOME/.local/bin:$PATH"
      uv sync --python 3.13

      echo "=== Creating results directory ==="
      mkdir -p ~/results

      echo "=== Verifying installation ==="
      uv run python -c 'import typeagent; print("OK")'

      echo "=== Running baseline benchmarks ==="
      bash ~/bench/bench_baseline.sh

      echo "=== Done ==="

runcmd:
  - wget -q https://github.com/sharkdp/hyperfine/releases/download/v1.19.0/hyperfine_1.19.0_amd64.deb -O /tmp/hyperfine.deb
  - dpkg -i /tmp/hyperfine.deb
  - su - azureuser -c 'bash /home/azureuser/setup.sh'
