#cloud-config
package_update: true
packages:
  - git
  - build-essential
  - curl
  - wget
  - jq

write_files:
  - path: /home/azureuser/bench/bench_import.sh
    owner: azureuser:azureuser
    permissions: "0755"
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      VENV_PYTHON="$HOME/rich/.venv/bin/python"
      echo "=== Rich overall import time ==="
      hyperfine --warmup 3 --min-runs 30 --shell=none \
        "$VENV_PYTHON -c 'import rich'"

  - path: /home/azureuser/bench/bench_module.sh
    owner: azureuser:azureuser
    permissions: "0755"
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      VENV_PYTHON="$HOME/rich/.venv/bin/python"
      echo "=== Per-module import time ==="
      hyperfine --warmup 3 --min-runs 20 --shell=none \
        -n 'rich (top-level)' "$VENV_PYTHON -c 'import rich'" \
        -n 'rich.console.Console' "$VENV_PYTHON -c 'from rich.console import Console'" \
        -n 'rich.logging.RichHandler' "$VENV_PYTHON -c 'from rich.logging import RichHandler'" \
        -n 'rich.traceback.Traceback' "$VENV_PYTHON -c 'from rich.traceback import Traceback'" \
        -n 'rich.print_json' "$VENV_PYTHON -c 'from rich import print_json'" \
        -n 'rich.syntax.Syntax' "$VENV_PYTHON -c 'from rich.syntax import Syntax'" \
        -n 'rich.pretty' "$VENV_PYTHON -c 'import rich.pretty'" \
        -n 'rich.markdown.Markdown' "$VENV_PYTHON -c 'from rich.markdown import Markdown'"

  - path: /home/azureuser/bench/bench_importtime.py
    owner: azureuser:azureuser
    permissions: "0755"
    content: |
      #!/usr/bin/env python3
      """Parse python -X importtime output and produce a sorted breakdown."""
      import subprocess
      import sys
      import re
      import os

      def parse_importtime(stderr_lines):
          pattern = re.compile(
              r"import time:\s+(\d+)\s+\|\s+(\d+)\s+\|\s+(\s*)([\w.]+)"
          )
          results = []
          for line in stderr_lines:
              m = pattern.match(line)
              if m:
                  self_us = int(m.group(1))
                  cumul_us = int(m.group(2))
                  indent = len(m.group(3)) // 2
                  module = m.group(4)
                  results.append((module, self_us, cumul_us, indent))
          return results

      def main():
          target = sys.argv[1] if len(sys.argv) > 1 else "import rich"
          venv_python = os.path.expanduser("~/rich/.venv/bin/python")

          proc = subprocess.run(
              [venv_python, "-X", "importtime", "-c", target],
              capture_output=True, text=True
          )
          entries = parse_importtime(proc.stderr.splitlines())
          entries.sort(key=lambda e: e[1], reverse=True)

          print(f"{'Module':<50} {'Self (us)':>12} {'Cumul (us)':>12} {'Depth':>6}")
          print("-" * 82)
          for mod, self_us, cumul_us, depth in entries[:40]:
              print(f"{mod:<50} {self_us:>12,} {cumul_us:>12,} {depth:>6}")

          if len(sys.argv) > 2:
              with open(sys.argv[2], "w") as f:
                  f.write("module\tself_us\tcumul_us\tdepth\n")
                  for mod, self_us, cumul_us, depth in entries:
                      f.write(f"{mod}\t{self_us}\t{cumul_us}\t{depth}\n")
              print(f"\nTSV written to {sys.argv[2]}")

      if __name__ == "__main__":
          main()

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

      cd ~/rich
      git checkout "$BRANCH"
      export PATH="$HOME/.local/bin:$PATH"
      uv pip install -e .

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

      hyperfine --warmup 3 --min-runs 30 --shell=none \
        --export-json "$OUTDIR/import.json" \
        "$VENV_PYTHON -c 'import rich'"

      hyperfine --warmup 3 --min-runs 20 --shell=none \
        --export-json "$OUTDIR/modules.json" \
        -n 'console' "$VENV_PYTHON -c 'from rich.console import Console'" \
        -n 'logging' "$VENV_PYTHON -c 'from rich.logging import RichHandler'" \
        -n 'traceback' "$VENV_PYTHON -c 'from rich.traceback import Traceback'" \
        -n 'syntax' "$VENV_PYTHON -c 'from rich.syntax import Syntax'" \
        -n 'markdown' "$VENV_PYTHON -c 'from rich.markdown import Markdown'"

      python3 ~/bench/bench_importtime.py "import rich" "$OUTDIR/importtime.tsv"

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

  - path: /home/azureuser/setup_rich.sh
    owner: azureuser:azureuser
    permissions: "0755"
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      export PATH="$HOME/.local/bin:$PATH"

      echo "=== Installing uv ==="
      curl -LsSf https://astral.sh/uv/install.sh | sh

      echo "=== Installing Python 3.12 and 3.13 ==="
      uv python install 3.12 3.13

      echo "=== Cloning Rich ==="
      git clone https://github.com/Textualize/rich.git ~/rich

      echo "=== Creating venv and installing Rich ==="
      cd ~/rich
      uv venv --python 3.13
      uv pip install -e .

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

      echo "=== Done ==="
      ~/rich/.venv/bin/python -c "import rich; print(f'Rich {rich.__version__} installed')"

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_rich.sh'
