padz 0.8.1

#!/usr/bin/env bash
# Generate LOC report from cargo warloc output
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
cd "$PROJECT_ROOT"

FORMAT="term"
COMPARE_BRANCH=""
WORKSPACES=()
ALL_WORKSPACES=false
TMPFILES=()
ORIGINAL_BRANCH=""

# Function to extract workspace members from Cargo.toml using Python to avoid fragile parsing
get_workspace_members() {
    python3 - "$PROJECT_ROOT/Cargo.toml" <<'PY'
import sys
from pathlib import Path

try:  # Python 3.11+
    import tomllib  # type: ignore[attr-defined]
except ModuleNotFoundError:  # pragma: no cover
    import tomli as tomllib  # type: ignore[assignment]

cargo_toml = Path(sys.argv[1])
with cargo_toml.open("rb") as fh:
    data = tomllib.load(fh)

members = data.get("workspace", {}).get("members", [])
for member in members:
    member = member.strip()
    if member:
        print(member)
PY
}

mapfile -t VALID_WORKSPACES < <(get_workspace_members)

if [[ ${#VALID_WORKSPACES[@]} -eq 0 ]]; then
  echo "No workspaces found in Cargo.toml" >&2
  exit 1
fi

DEFAULT_WORKSPACE="${VALID_WORKSPACES[0]}"

cleanup() {
  for tmp in "${TMPFILES[@]}"; do
    if [[ -f "$tmp" ]]; then
      rm -f "$tmp"
    fi
  done

  if [[ -n "$ORIGINAL_BRANCH" ]]; then
    local current_branch
    current_branch="$(git rev-parse --abbrev-ref HEAD 2>/dev/null || true)"
    if [[ -n "$current_branch" && "$current_branch" != "$ORIGINAL_BRANCH" ]]; then
      git checkout "$ORIGINAL_BRANCH" >/dev/null 2>&1 || true
    fi
  fi
}
trap cleanup EXIT

usage() {
  echo "Usage: $0 [--format term|json] [--compare-branch <branch>] [--all | <workspace>...]" >&2
  echo "Available workspaces: ${VALID_WORKSPACES[*]}" >&2
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --format)
      FORMAT="$2"
      shift 2
      ;;
    --format=*)
      FORMAT="${1#*=}"
      shift
      ;;
    --compare-branch)
      COMPARE_BRANCH="$2"
      shift 2
      ;;
    --compare-branch=*)
      COMPARE_BRANCH="${1#*=}"
      shift
      ;;
    --all)
      ALL_WORKSPACES=true
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      if [[ "$1" == -* ]]; then
        echo "Unknown argument: $1" >&2
        usage
        exit 1
      else
        WORKSPACES+=("$1")
        shift
      fi
      ;;
  esac
done

if [[ "$FORMAT" != "term" && "$FORMAT" != "json" ]]; then
  echo "Invalid format: $FORMAT" >&2
  usage
  exit 1
fi

if ! $ALL_WORKSPACES; then
  if [[ ${#WORKSPACES[@]} -eq 0 ]]; then
    WORKSPACES=("$DEFAULT_WORKSPACE")
  fi

  for ws in "${WORKSPACES[@]}"; do
    if ! [[ " ${VALID_WORKSPACES[*]} " =~ " $ws " ]]; then
      echo "Invalid workspace: $ws" >&2
      usage
      exit 1
    fi
  done
fi

generate_tmpfile() {
  local tmpfile
  tmpfile=$(mktemp)
  TMPFILES+=("$tmpfile")
  cargo warloc --by-file > "$tmpfile"
  echo "$tmpfile"
}

render_report() {
  local tmpfile="$1"
  local ws_args=()
  if ! $ALL_WORKSPACES; then
    ws_args+=("--workspaces" "${WORKSPACES[@]}")
  fi

  if [[ "$FORMAT" == "term" ]]; then
    python3 "$SCRIPT_DIR/_rust_loc.py" --format term "$tmpfile" "${ws_args[@]}" | npx prettier --parser markdown
  else
    python3 "$SCRIPT_DIR/_rust_loc.py" --format json "$tmpfile" "${ws_args[@]}"
  fi
}

collect_summary() {
  local tmpfile="$1"
  local ws_args=()
  if ! $ALL_WORKSPACES; then
    ws_args+=("--workspaces" "${WORKSPACES[@]}")
  fi
  python3 "$SCRIPT_DIR/_rust_loc.py" --format json "$tmpfile" "${ws_args[@]}"
}

print_comparison() {
  local base_json="$1"
  local compare_json="$2"
  local base_branch="$3"
  local compare_branch="$4"

  python3 - "$base_json" "$compare_json" "$base_branch" "$compare_branch" <<'PY'
import json
import sys

base = json.loads(sys.argv[1])
other = json.loads(sys.argv[2])
base_branch = sys.argv[3]
other_branch = sys.argv[4]

base_summary = base['summary']
other_summary = other['summary']

diff_code = other_summary['total_code'] - base_summary['total_code']
diff_tests = other_summary['total_tests'] - base_summary['total_tests']

print()
print(f"Comparison against branch '{other_branch}':")
print(f"  {base_branch}:   code {base_summary['total_code']}, tests {base_summary['total_tests']}")
print(f"  {other_branch}: code {other_summary['total_code']}, tests {other_summary['total_tests']}")
print(f"  Δ (other - {base_branch}): code {diff_code:+}, tests {diff_tests:+}")
PY
}

if [[ -n "$COMPARE_BRANCH" ]]; then
  if [[ -n "$(git status --porcelain)" ]]; then
    echo "Cannot compare branches with a dirty working tree. Please commit or stash changes first." >&2
    exit 1
  fi
fi

initial_tmp=$(generate_tmpfile)

if [[ "$FORMAT" == "term" ]]; then
  render_report "$initial_tmp"
  current_summary_json=$(collect_summary "$initial_tmp")
else
  current_summary_json=$(render_report "$initial_tmp")
  printf '%s\n' "$current_summary_json"
fi

if [[ -n "$COMPARE_BRANCH" ]]; then
  ORIGINAL_BRANCH="$(git rev-parse --abbrev-ref HEAD)"
  base_branch_name="$ORIGINAL_BRANCH"

  git checkout "$COMPARE_BRANCH" >/dev/null 2>&1

  compare_tmp=$(generate_tmpfile)
  compare_summary_json=$(collect_summary "$compare_tmp")

  git checkout "$ORIGINAL_BRANCH" >/dev/null 2>&1
  ORIGINAL_BRANCH=""

  print_comparison "$current_summary_json" "$compare_summary_json" "$base_branch_name" "$COMPARE_BRANCH"
fi