#!/usr/bin/env bash
set -euo pipefail

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

# clang-format executable: prefer the CLANG_FORMAT env var, otherwise look it up on PATH / in the pip package
find_clang_format() {
    if [ -n "${CLANG_FORMAT:-}" ]; then
        echo "$CLANG_FORMAT"
        return
    fi
    if command -v clang-format >/dev/null 2>&1; then
        command -v clang-format
        return
    fi
    # the binary shipped by pip install clang-format (not necessarily on PATH)
    local pip_bin
    pip_bin=$(python -c "import clang_format, os; print(os.path.join(os.path.dirname(clang_format.__file__), 'data', 'bin', 'clang-format'))" 2>/dev/null || true)
    if [ -n "$pip_bin" ] && [ -x "$pip_bin" ]; then
        echo "$pip_bin"
        return
    fi
    echo "ERROR: clang-format not found. Install it or set CLANG_FORMAT=<path>." >&2
    exit 1
}

CLANG_FORMAT_BIN=$(find_clang_format)

# file extensions
EXTENSIONS=("*.cpp" "*.cc" "*.h" "*.hpp" "*.cu" "*.cuh")

# build the -name arguments for find
NAME_ARGS=()
for i in "${!EXTENSIONS[@]}"; do
    if [ "$i" -gt 0 ]; then
        NAME_ARGS+=("-o")
    fi
    NAME_ARGS+=("-name" "${EXTENSIONS[$i]}")
done

FILES=$(find "$CSRC_DIR" \( "${NAME_ARGS[@]}" \))

if [ -z "$FILES" ]; then
    echo "No files found to format."
    exit 0
fi

FILE_COUNT=$(echo "$FILES" | wc -l)

# mode: --check only checks and does not modify
if [ "${1:-}" = "--check" ]; then
    echo "Checking clang-format for $FILE_COUNT files in csrc/..."
    FAILED=0
    while IFS= read -r file; do
        if ! "$CLANG_FORMAT_BIN" --dry-run --Werror "$file" 2>/dev/null; then
            echo "  FAIL: $file"
            FAILED=$((FAILED + 1))
        fi
    done <<< "$FILES"
    if [ "$FAILED" -gt 0 ]; then
        echo "$FAILED file(s) need formatting."
        exit 1
    else
        echo "All files are correctly formatted."
        exit 0
    fi
else
    echo "Formatting $FILE_COUNT files in csrc/..."
    echo "$FILES" | xargs "$CLANG_FORMAT_BIN" -i
    echo "Done."
fi