#!/bin/zsh

set -euo pipefail

REAL_XCRUN=/usr/bin/xcrun
METAL_BIN=/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.3.7003.10.6SAP8k/Metal.xctoolchain/usr/bin/metal
MODULE_CACHE=/tmp/mlx-metal-module-cache
LOG_FILE=/tmp/mlx-xcrun.log

args=("$@")
print -r -- "${(j: :)args}" >> "$LOG_FILE"

if (( ${#args[@]} >= 3 )) && [[ "${args[1]}" == "-sdk" || "${args[1]}" == "--sdk" ]] && [[ "${args[3]}" == "metal" ]]; then
    metal_args=("${args[4,-1]}")
    if (( ${#metal_args[@]} == 1 )) && [[ "${metal_args[1]}" == *" "* ]]; then
        metal_args=(${=metal_args[1]})
    fi
    metal_arg_string="${(j: :)metal_args}"
    if [[ "$metal_arg_string" == *" -E "* || "$metal_arg_string" == "-E "* || "$metal_arg_string" == *" -E" ]] && [[ "$metal_arg_string" == *" -P "* || "$metal_arg_string" == "-P "* || "$metal_arg_string" == *" -P" ]] && [[ "${metal_args[-1]}" == "-" ]]; then
        print '400'
        exit 0
    fi
    mkdir -p "$MODULE_CACHE"
    exec "$METAL_BIN" -fmodules-cache-path="$MODULE_CACHE" "${metal_args[@]}"
fi

exec "$REAL_XCRUN" "$@"