"""Programs the kernel puts first on PATH for the model's jobs, each in its
own directory under $ALBEDO_HOME/shims, prepended only where the program it
needs is on the job's PATH. They are written where the kernel runs (only .py
files travel to a remote kernel). robot-docs/shims.md has the measurements.

`nix` (needs nix) runs albedo_nix.py with the kernel's Python: a flake in a jj
workspace is read from its commit instead of copied into the store, and
`nix develop -c` reuses a cached dev shell. ALBEDO_PLAIN_NIX=1 runs nix as is.

`cargo` (needs mbx, jdx's mr-boxington) runs cargo through mbx, so every
checkout of a project shares compiled work: a second checkout of
native/render builds in 2 s instead of 32 s. ALBEDO_NO_MBX=1 runs plain cargo.
For mbx's process only, SDKROOT goes, because mbx names the macOS linker with
`xcrun --sdk $SDKROOT`, which rejects a path such as a nix shell's, and an
unnamed linker leaves every link, build script, and proc macro uncached.
Inside a nix shell (ALBEDO_NIX_ENV, set by albedo_nix.py) CC and CXX go too
when they are the shell's cc and c++ anyway (a shell that picks another
compiler keeps it), so build scripts' C compiles through mbx, which an
explicit CC hides it from, and HOST_CFLAGS names the shell's
compiler-wrapper flags, which no mbx key covers: without it a changed shell
could be served another checkout's stale object.
"""

from __future__ import annotations

import os
import shutil
import sys
from pathlib import Path

CARGO = r"""#!/bin/sh
# albedo's cargo: mbx when it is installed (see albedo_shims.py), else the next
# cargo on PATH.
here=$(cd "$(dirname "$0")" && pwd)
# whether compiler $1 is the $2 on PATH, which a build script runs without it
default() {
  a=$(command -v "$1" 2>/dev/null) && b=$(command -v "$2" 2>/dev/null) || return 1
  a=$(realpath "$a" 2>/dev/null) && b=$(realpath "$b" 2>/dev/null) || return 1
  [ "$a" = "$b" ]
}
if [ -z "${ALBEDO_NO_MBX:-}" ] && mbx=$(command -v mbx); then
  if [ -n "${ALBEDO_NIX_ENV:-}" ]; then
    HOST_CFLAGS="${HOST_CFLAGS:-${CFLAGS:-}} -DALBEDO_NIX_ENV_$ALBEDO_NIX_ENV"
    HOST_CXXFLAGS="${HOST_CXXFLAGS:-${CXXFLAGS:-}} -DALBEDO_NIX_ENV_$ALBEDO_NIX_ENV"
    export HOST_CFLAGS HOST_CXXFLAGS
    if default "${CC:-}" cc; then unset CC; fi
    if default "${CXX:-}" c++; then unset CXX; fi
  fi
  MBX_CARGO_SHIM_MODE=1 MBX_CARGO_SHIM_PATH="$0" exec env -u SDKROOT "$mbx" "$@"
fi
PATH=$(printf %s "$PATH" | tr : '\n' | grep -vxF "$here" | paste -sd: -)
exec cargo "$@"
"""

NIX = r"""#!/bin/sh
# albedo's nix: albedo_nix.py with the kernel's Python, else the next nix on PATH.
here=$(cd "$(dirname "$0")" && pwd)
if [ -z "${ALBEDO_PLAIN_NIX:-}" ] && [ -x "${ALBEDO_SHIMS_PYTHON:-}" ]; then
  exec "$ALBEDO_SHIMS_PYTHON" -E -s "$ALBEDO_SHIMS_LIB/albedo_nix.py" "$here" "$@"
fi
PATH=$(printf %s "$PATH" | tr : '\n' | grep -vxF "$here" | paste -sd: -)
exec nix "$@"
"""

# directory (named for the program a shim needs on PATH) -> (shim, script)
SHIMS = {"nix": ("nix", NIX), "mbx": ("cargo", CARGO)}
# where a program lives when the process that started the daemon had a thin
# PATH (a GUI or launchd start never ran the user's profile)
FALLBACKS = {"nix": ("/nix/var/nix/profiles/default/bin", "~/.nix-profile/bin")}
_written: set[Path] = set()


def _with_fallbacks(entries: list[str]) -> list[str]:
    """`entries`, plus the directory providing each program of FALLBACKS that
    PATH lacks. It goes last, so it never shadows what the user put on PATH, and
    only the first directory that holds the program is added, not a whole profile."""
    for program, directories in FALLBACKS.items():
        if shutil.which(program, path=os.pathsep.join(entries)):
            continue
        for directory in map(os.path.expanduser, directories):
            if os.access(os.path.join(directory, program), os.X_OK):
                entries = [*entries, directory]
                break
    return entries


def _directory(needs: str) -> Path:
    """The directory of the shim that needs `needs`, written once per kernel."""
    home = os.environ.get("ALBEDO_HOME") or os.path.expanduser("~/.albedo")
    directory = Path(home, "shims", needs)
    if directory not in _written:
        name, script = SHIMS[needs]
        directory.mkdir(parents=True, exist_ok=True)
        shim = directory / name
        if not shim.is_file() or shim.read_text() != script:
            temporary = directory / f".{name}.{os.getpid()}"
            temporary.write_text(script)
            temporary.chmod(0o755)
            temporary.replace(shim)
        _written.add(directory)
    return directory


def prepend(environment: dict[str, str] | None) -> dict[str, str] | None:
    """Return the job environment with only this kernel's shims before real tools.
    Remove inherited shim generations even when their tools are absent."""
    base = environment or os.environ
    path = base.get("PATH", "")
    entries = path.split(os.pathsep)
    # Remove inherited generations before adding this kernel's shims.
    entries = [
        entry
        for entry in entries
        if not (Path(entry).name in SHIMS and Path(entry).parent.name == "shims")
    ]
    path = os.pathsep.join(_with_fallbacks(entries))
    shims = [
        str(_directory(needs))
        for needs in SHIMS
        if shutil.which(needs, path=path) is not None
    ]
    if not shims:
        if environment is None and path == base.get("PATH", ""):
            return None
        return {**base, "PATH": path}
    return {
        **base,
        "PATH": os.pathsep.join([*shims, path]),
        "ALBEDO_SHIMS_PYTHON": sys.executable,
        "ALBEDO_SHIMS_LIB": str(Path(__file__).resolve().parent),
    }