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

# Some hardened environments (e.g., newer glibc / gVisor) refuse to dlopen()
# shared objects that request an executable stack. Certain PyTorch builds can
# ship with this flag set on libtorch_cpu.so. Clearing it fixes torch import.

python - <<'PY'
import os, site, glob

paths = []
for fn in (getattr(site, "getsitepackages", None), getattr(site, "getusersitepackages", None)):
    if fn is None:
        continue
    try:
        p = fn()
        if isinstance(p, str):
            paths.append(p)
        else:
            paths.extend(list(p))
    except Exception:
        pass

targets = []
for p in paths:
    targets += glob.glob(os.path.join(p, "torch", "lib", "libtorch_cpu.so"))
    targets += glob.glob(os.path.join(p, "torch", "lib", "libtorch_python.so"))

print("\n".join(sorted(set(targets))))
PY

TARGETS="$(python - <<'PY'
import os, site, glob

paths = []
for fn in (getattr(site, "getsitepackages", None), getattr(site, "getusersitepackages", None)):
    if fn is None:
        continue
    try:
        p = fn()
        if isinstance(p, str):
            paths.append(p)
        else:
            paths.extend(list(p))
    except Exception:
        pass

targets = []
for p in paths:
    targets += glob.glob(os.path.join(p, "torch", "lib", "libtorch_cpu.so"))
    targets += glob.glob(os.path.join(p, "torch", "lib", "libtorch_python.so"))

print(" ".join(sorted(set(targets))))
PY
)"

if [[ -z "${TARGETS}" ]]; then
  echo "postBuild: no torch .so targets found yet (torch may not be installed)."
  exit 0
fi

for so in ${TARGETS}; do
  if [[ -f "${so}" ]]; then
    echo "postBuild: clearing execstack on ${so}"
    # If it fails (e.g. already clear), continue.
    patchelf --clear-execstack "${so}" || true
  fi
done