setup.py

#!/usr/bin/env python3
"""
setup.py — Post-bootstrap configuration for managed-python.

Pure stdlib. No external dependencies.

Invoked by install.sh / install.ps1 after uv and the venv have been
created. Handles everything that does not require platform-specific
shell syntax: bin/ wrappers, env.sh, env.ps1, distro.toml copy, and
optional shell profile update.
"""

from __future__ import annotations

import argparse
import os
import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

_IS_WINDOWS = sys.platform == "win32"
_QUIET = False

# Default package cooldown: ignore distributions uploaded in the last 24 hours.
# Malicious releases are usually spotted and yanked within hours, so waiting a
# day is cheap insurance against being patient zero.
_DEFAULT_COOLDOWN = "P1D"

# "No cooldown" — a zero-length duration, which uv accepts. Note that uv's
# documented "off" value (false) is only valid in uv.toml, NOT as an env var.
_COOLDOWN_OFF = "P0D"


# ── TOML (minimal key=value reader, no deps) ─────────────────────────────────

def _toml_get(path: Path, key: str) -> str:
    for line in path.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if stripped.startswith(key) and "=" in stripped:
            _, _, raw = stripped.partition("=")
            return raw.split("#")[0].strip().strip('"')
    raise KeyError(f"{key!r} not found in {path}")


# ── Output ────────────────────────────────────────────────────────────────────

def _banner(msg: str) -> None:
    if _QUIET:
        return
    width = len(msg) + 4  # 2 spaces padding each side
    print(f"┌{'─' * width}┐")
    print(f"│  {msg}  │")
    print(f"└{'─' * width}┘")


def _ok(msg: str) -> None:
    if _QUIET:
        return
    print(f"  \u2713 {msg}")


def _info(msg: str) -> None:
    if _QUIET:
        return
    print(f"  \u2139 {msg}")


def _warn(msg: str) -> None:
    print(f"  \u26a0 {msg}")


def _step(msg: str) -> None:
    if _QUIET:
        return
    print(f"\n==> {msg}")


# ── PATH detection ────────────────────────────────────────────────────────────

def _path_decision(bin_dir: Path) -> tuple[bool, list[str]]:
    """Decide whether to add bin_dir to PATH in generated env files.

    Args:
        bin_dir: The managed-python bin/ directory to conditionally add to PATH.

    Returns:
        Tuple of (add_to_path, note_lines) where add_to_path is True when the
        bin directory should be prepended to PATH, and note_lines is a list of
        human-readable notes explaining the decision.
    """
    def _real_executable(path: str | None) -> str | None:
        """Return path only if it is a non-empty file (filters Windows App Execution Alias stubs)."""
        if not path:
            return None
        try:
            if Path(path).stat().st_size > 0:
                return path
        except OSError:
            return None
        return None

    python_found = _real_executable(shutil.which("python")) or _real_executable(shutil.which("python3"))
    uv_found = _real_executable(shutil.which("uv"))

    if python_found and uv_found:
        if _IS_WINDOWS:
            path_hint = f'To use managed versions: $env:PATH = "{bin_dir};" + $env:PATH'
        else:
            path_hint = f'To use managed versions: export PATH="{bin_dir}:$PATH"'
        return False, [
            "python and uv already on PATH — PATH not modified",
            path_hint,
        ]

    notes: list[str] = []
    if python_found:
        notes.append(f"system python found at {python_found} — will be shadowed by managed version")
    if uv_found:
        notes.append(f"system uv found at {uv_found} — will be shadowed by managed version")
    return True, notes


# ── Package cooldown ──────────────────────────────────────────────────────────

def _is_zero_cooldown(cooldown: str) -> bool:
    """Return True when the cooldown is a zero-length window, i.e. no cooldown.

    uv accepts several spellings of "no window" (P0D, PT0H, "0 days") and they
    all mean the same thing, so they must all be reported the same way. A date
    or timestamp can never be all-zeros, so no special-casing is needed.

    Args:
        cooldown: A value already validated by _validate_cooldown.

    Returns:
        True when every numeric component of the value is zero.
    """
    digits = re.findall(r"\d+", cooldown)
    return bool(digits) and all(int(d) == 0 for d in digits)


def _validate_cooldown(prefix: Path, cooldown: str) -> str | None:
    """Ask uv itself whether it accepts the cooldown value.

    Delegating beats hand-rolling a grammar: uv is the only consumer of this
    value, so its parser is the only opinion that matters. Runs fully offline
    against an empty requirements list, so it is cheap and network-free.

    Args:
        prefix: The install prefix containing the bootstrapped uv binary.
        cooldown: The candidate --cooldown value.

    Returns:
        uv's error message when the value is rejected, otherwise None. Also
        returns None when uv cannot be run at all, so a broken validator never
        blocks an install.
    """
    uv_bin = prefix / ("uv.exe" if _IS_WINDOWS else "uv")
    if not uv_bin.exists():
        return None

    with tempfile.TemporaryDirectory() as tmp:
        try:
            proc = subprocess.run(
                [str(uv_bin), "pip", "compile", "--offline", "--quiet", "--no-cache",
                 "--exclude-newer", cooldown, "-", "-o", str(Path(tmp) / "out.txt")],
                input="", capture_output=True, text=True, timeout=60,
            )
        except (OSError, subprocess.SubprocessError):
            return None

    for line in proc.stderr.splitlines():
        if "invalid value" in line:
            return line.strip()
    return None


def _cooldown_lines(cooldown: str, comment: str, assign: str) -> list[str]:
    """Render the UV_EXCLUDE_NEWER block for a generated env file.

    When the cooldown is disabled the assignment is emitted commented-out, so
    the file still documents the decision without touching the environment.

    The hint text avoids angle brackets on purpose: env.bat is consumed by CMD,
    where < and > are redirection operators, and relying on ":: " label
    semantics to suppress them is not a bet worth taking.

    Args:
        cooldown: The cooldown duration, e.g. "P1D". Any zero-length value
            (P0D, PT0H, "0 days") means disabled.
        comment: The comment marker for the target format ("#" or "::").
        assign: A format string with one {} placeholder for the value, e.g.
            'export UV_EXCLUDE_NEWER="{}"'.

    Returns:
        The lines to append to the generated env file.
    """
    lines = [
        f"{comment} Package cooldown - ignore distributions uploaded within this window.",
        f"{comment} Applies to uv/uvx resolution. Bypass for an urgent patch with:",
        f"{comment}   uv pip install --exclude-newer {_COOLDOWN_OFF} PACKAGE",
    ]
    if _is_zero_cooldown(cooldown):
        lines.append(f"{comment} Cooldown disabled at install time (--cooldown {cooldown}).")
        lines.append(f"{comment} {assign.format(_DEFAULT_COOLDOWN)}")
    else:
        lines.append(assign.format(cooldown))
    lines.append("")
    return lines


# ── bin/ wrappers ─────────────────────────────────────────────────────────────

def _symlink(link: Path, target: Path) -> None:
    if link.is_symlink() or link.exists():
        link.unlink()
    link.symlink_to(target)


def _create_bin(prefix: Path) -> None:
    _step("Creating bin/ wrappers")
    bin_dir = prefix / "bin"
    bin_dir.mkdir(parents=True, exist_ok=True)

    if _IS_WINDOWS:
        uv_exe  = prefix / "uv.exe"
        uvx_exe = prefix / "uvx.exe"
        venv_py = prefix / "venv" / "Scripts" / "python.exe"
        (bin_dir / "python.cmd").write_text(f'@"{venv_py}" %*\n', encoding="utf-8")
        (bin_dir / "uv.cmd").write_text(f'@"{uv_exe}" %*\n', encoding="utf-8")
        (bin_dir / "uvx.cmd").write_text(f'@"{uvx_exe}" %*\n', encoding="utf-8")
        _ok("bin\\python.cmd")
        _ok("bin\\uv.cmd")
        _ok("bin\\uvx.cmd")
    else:
        _symlink(bin_dir / "python", Path("../venv/bin/python"))
        _symlink(bin_dir / "uv",     Path("../uv"))
        _symlink(bin_dir / "uvx",    Path("../uvx"))
        _ok("bin/python \u2192 ../venv/bin/python")
        _ok("bin/uv \u2192 ../uv")
        _ok("bin/uvx \u2192 ../uvx")


# ── env.sh ────────────────────────────────────────────────────────────────────

def _to_sh_path(p: Path) -> str:
    """Return a POSIX path string suitable for use in env.sh.

    On Windows, converts drive-letter paths (C:/...) to Git Bash POSIX paths
    (/c/...) so that env.sh can be sourced from Git Bash.
    """
    s = p.as_posix()
    if _IS_WINDOWS and len(s) >= 2 and s[1] == ":":
        s = "/" + s[0].lower() + s[2:]
    return s


def _write_env_sh(prefix: Path, uv_env: str, uvx_env: str, python_env: str,
                  distro_version: str, cooldown: str, *, isolated: bool = False) -> None:
    _step("Writing env.sh")

    if _IS_WINDOWS:
        uv_bin  = prefix / "uv.exe"
        uvx_bin = prefix / "uvx.exe"
        venv_py = prefix / "venv" / "Scripts" / "python.exe"
    else:
        uv_bin  = prefix / "uv"
        uvx_bin = prefix / "uvx"
        venv_py = prefix / "venv" / "bin" / "python"
    bin_dir = prefix / "bin"

    lines: list[str] = [
        f"# managed-python v{distro_version} \u2014 generated by setup.py",
        "# Do not edit manually; re-run install.sh to regenerate",
        "",
        "# Env vars (always set \u2014 these are the reliable contract)",
        f'export {uv_env}="{_to_sh_path(uv_bin)}"',
        f'export {uvx_env}="{_to_sh_path(uvx_bin)}"',
        f'export {python_env}="{_to_sh_path(venv_py)}"',
        "",
    ]
    lines += _cooldown_lines(cooldown, "#", 'export UV_EXCLUDE_NEWER="{}"')

    if _IS_WINDOWS:
        # bin/ wrappers on Windows are .cmd files — not usable in Git Bash.
        # Users should invoke the env vars directly (e.g. "$REDMATTER_UV").
        lines.append("# bin/ wrappers are .cmd files — PATH not modified for Git Bash")
        (prefix / "env.sh").write_text("\n".join(lines) + "\n", encoding="utf-8")
        _ok("env.sh")
        _info("Git Bash users can source this file")
        return

    if isolated:
        add_to_path = True
        notes: list[str] = ["--isolated: always adding bin/ to PATH"]
    else:
        add_to_path, notes = _path_decision(bin_dir)

    if add_to_path:
        lines += [f"# {n}" for n in notes]
        lines.append(f'export PATH="{bin_dir}:$PATH"')
    else:
        lines += [f"# {n}" for n in notes]

    (prefix / "env.sh").write_text("\n".join(lines) + "\n", encoding="utf-8")
    _ok("env.sh")
    for note in notes:
        (_warn if "shadowed" in note else _info)(note)


# ── env.ps1 ───────────────────────────────────────────────────────────────────

def _write_env_ps1(prefix: Path, uv_env: str, uvx_env: str, python_env: str,
                   distro_version: str, cooldown: str, *, isolated: bool = False) -> None:
    _step("Writing env.ps1")

    uv_exe  = prefix / ("uv.exe" if _IS_WINDOWS else "uv")
    uvx_exe = prefix / ("uvx.exe" if _IS_WINDOWS else "uvx")
    venv_py = prefix / "venv" / ("Scripts" if _IS_WINDOWS else "bin") / (
        "python.exe" if _IS_WINDOWS else "python"
    )
    bin_dir = prefix / "bin"

    if isolated:
        add_to_path = True
        notes = ["--isolated: always adding bin/ to PATH"]
    else:
        add_to_path, notes = _path_decision(bin_dir)

    lines: list[str] = [
        f"# managed-python v{distro_version} -- generated by setup.py",
        "# Do not edit manually; re-run install to regenerate",
        "",
        "# Env vars (always set -- these are the reliable contract)",
        f'$env:{uv_env} = "{uv_exe}"',
        f'$env:{uvx_env} = "{uvx_exe}"',
        f'$env:{python_env} = "{venv_py}"',
    ]
    if _IS_WINDOWS:
        lines.append('$env:PYTHONUTF8 = "1"')
    lines.append("")
    lines += _cooldown_lines(cooldown, "#", '$env:UV_EXCLUDE_NEWER = "{}"')

    if add_to_path:
        lines += [f"# {n}" for n in notes]
        lines.append(f'$env:PATH = "{bin_dir};" + $env:PATH')
    else:
        lines += [f"# {n}" for n in notes]
        lines.append(f'# $env:PATH = "{bin_dir};" + $env:PATH')

    (prefix / "env.ps1").write_text("\n".join(lines) + "\n", encoding="utf-8")
    _ok("env.ps1")


# ── env.bat ───────────────────────────────────────────────────────────────────

def _write_env_bat(prefix: Path, uv_env: str, uvx_env: str, python_env: str,
                   distro_version: str, cooldown: str, *, isolated: bool = False) -> None:
    if not _IS_WINDOWS:
        return

    uv_exe  = prefix / "uv.exe"
    uvx_exe = prefix / "uvx.exe"
    venv_py = prefix / "venv" / "Scripts" / "python.exe"
    bin_dir = prefix / "bin"

    if isolated:
        add_to_path = True
        notes = ["--isolated: always adding bin/ to PATH"]
    else:
        add_to_path, notes = _path_decision(bin_dir)

    lines: list[str] = [
        f"@echo off",
        f":: managed-python v{distro_version} -- generated by setup.py",
        ":: Do not edit manually; re-run install to regenerate",
        "",
        ":: Env vars (always set -- these are the reliable contract)",
        f'SET {uv_env}={uv_exe}',
        f'SET {uvx_env}={uvx_exe}',
        f'SET {python_env}={venv_py}',
        'SET PYTHONUTF8=1',
        "",
    ]
    lines += _cooldown_lines(cooldown, "::", "SET UV_EXCLUDE_NEWER={}")

    if add_to_path:
        lines += [f":: {n}" for n in notes]
        lines.append(f'SET PATH={bin_dir};%PATH%')
    else:
        lines += [f":: {n}" for n in notes]

    (prefix / "env.bat").write_text("\r\n".join(lines) + "\r\n", encoding="utf-8")
    _ok("env.bat")


# ── distro.toml ───────────────────────────────────────────────────────────────

def _write_installed_distro_toml(
    script_dir: Path,
    prefix: Path,
    min_python: str,
    uv_env: str,
    uvx_env: str,
    python_env: str,
    shell_profile: bool,
    cooldown: str,
    isolated: bool = False,
) -> None:
    """
    Write distro.toml to the prefix, appending an [install] section that
    records the options used. Enables inspection and replay of the install.
    """
    source = (script_dir / "distro.toml").read_text(encoding="utf-8").rstrip()
    install_section = (
        f"\n\n[install]\n"
        f'prefix       = "{prefix.as_posix()}"\n'
        f'python       = "{min_python}"\n'
        f'uv_env       = "{uv_env}"\n'
        f'uvx_env      = "{uvx_env}"\n'
        f'python_env   = "{python_env}"\n'
        f'cooldown     = "{cooldown}"\n'
        f"shell_profile = {'true' if shell_profile else 'false'}\n"
        f"isolated      = {'true' if isolated else 'false'}\n"
    )
    (prefix / "distro.toml").write_text(source + install_section, encoding="utf-8")


# ── Shell profile ─────────────────────────────────────────────────────────────

def _update_shell_profile(prefix: Path) -> None:
    _step("Updating shell profile")

    if _IS_WINDOWS:
        env_ps1 = prefix / "env.ps1"
        _info(f'Add to your PowerShell profile:  . "{env_ps1}"')
        return

    env_sh      = prefix / "env.sh"
    source_line = f'source "{env_sh}"'

    shell = os.environ.get("SHELL", "")
    if "zsh" in shell:
        rc = Path.home() / ".zshrc"
    elif "bash" in shell:
        rc = Path.home() / ".bashrc"
    elif (Path.home() / ".zshrc").exists():
        rc = Path.home() / ".zshrc"
    elif (Path.home() / ".bashrc").exists():
        rc = Path.home() / ".bashrc"
    else:
        _warn(f'Could not detect shell rc. Add manually:  source "{env_sh}"')
        return

    existing = rc.read_text(encoding="utf-8") if rc.exists() else ""
    if source_line in existing:
        _ok(f"Shell profile already configured: {rc}")
        return

    with rc.open("a", encoding="utf-8") as fh:
        fh.write(f"\n# managed-python\n{source_line}\n")
    _ok(f"Appended to {rc}")
    _info(f'Restart your shell or run:  source "{env_sh}"')


# ── Main ──────────────────────────────────────────────────────────────────────

def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Post-bootstrap configuration for managed-python (stdlib, no deps).",
    )
    p.add_argument("--prefix",       required=True,                help="Install prefix")
    p.add_argument("--python",       required=True, dest="python_version")
    p.add_argument("--env-prefix",   dest="env_prefix",           help="Common prefix for env var names (e.g. REDMATTER → REDMATTER_UV, REDMATTER_UVX, REDMATTER_PYTHON)")
    p.add_argument("--uv-env",       dest="uv_env",               help="Env var name for uv path")
    p.add_argument("--uvx-env",      dest="uvx_env",              help="Env var name for uvx path")
    p.add_argument("--python-env",   dest="python_env",           help="Env var name for python path")
    p.add_argument("--cooldown",     default=_DEFAULT_COOLDOWN,
                   help=f"Ignore packages uploaded within this window "
                        f"(default: {_DEFAULT_COOLDOWN}). Accepts a duration "
                        f"(P1D, '2 days', PT12H), a date, or a timestamp. "
                        f"Use {_COOLDOWN_OFF} to disable.")
    p.add_argument("--shell-profile", action="store_true", dest="shell_profile")
    p.add_argument("--isolated",     action="store_true", dest="isolated")
    p.add_argument("--quiet", "-q", action="store_true", dest="quiet",
                   help="Suppress all output except warnings")
    args = p.parse_args()

    args.cooldown = args.cooldown.strip()

    individual = [args.uv_env, args.uvx_env, args.python_env]
    if args.env_prefix:
        if any(individual):
            p.error("--env-prefix cannot be combined with --uv-env, --uvx-env, or --python-env")
        args.uv_env     = f"{args.env_prefix}_UV"
        args.uvx_env    = f"{args.env_prefix}_UVX"
        args.python_env = f"{args.env_prefix}_PYTHON"
    else:
        missing = [flag for flag, val in (("--uv-env", args.uv_env), ("--uvx-env", args.uvx_env), ("--python-env", args.python_env)) if not val]
        if missing:
            p.error(f"the following arguments are required: {', '.join(missing)} (or use --env-prefix)")

    return args


def main() -> None:
    global _QUIET
    args       = _parse_args()
    _QUIET     = args.quiet
    prefix     = Path(args.prefix).expanduser().resolve()
    script_dir = Path(__file__).parent.resolve()

    distro_version = _toml_get(script_dir / "distro.toml", "version")

    # Validate before writing anything — a bad cooldown would otherwise poison
    # every uv invocation of everyone who sources the generated env files.
    rejection = _validate_cooldown(prefix, args.cooldown)
    if rejection:
        print(f"ERROR: --cooldown {args.cooldown!r} was rejected by uv:\n  {rejection}", file=sys.stderr)
        sys.exit(1)

    _create_bin(prefix)
    _write_env_sh(prefix, args.uv_env, args.uvx_env, args.python_env, distro_version,
                  args.cooldown, isolated=args.isolated)
    _write_env_ps1(prefix, args.uv_env, args.uvx_env, args.python_env, distro_version,
                   args.cooldown, isolated=args.isolated)
    _write_env_bat(prefix, args.uv_env, args.uvx_env, args.python_env, distro_version,
                   args.cooldown, isolated=args.isolated)
    _write_installed_distro_toml(
        script_dir, prefix, args.python_version, args.uv_env, args.uvx_env, args.python_env, args.shell_profile,
        args.cooldown, isolated=args.isolated,
    )

    _step("Package cooldown")
    if _is_zero_cooldown(args.cooldown):
        _warn(f"disabled (--cooldown {args.cooldown}) - fresh releases will be installed immediately")
    else:
        _ok(f"UV_EXCLUDE_NEWER={args.cooldown} - packages uploaded within this window are ignored")
        _info(f"Urgent patch? uv pip install --exclude-newer {_COOLDOWN_OFF} PACKAGE")

    if args.shell_profile:
        _update_shell_profile(prefix)

    if not _QUIET:
        print()
        _banner("Install complete!")
        print()
        if _IS_WINDOWS:
            print(f'  . "{prefix / "env.ps1"}"')
            print(f'  or (PowerShell, if scripts are restricted):')
            print(f'    Invoke-Expression (Get-Content "{prefix / "env.ps1"}" -Raw)')
            print(f'  or (CMD):      call "{prefix / "env.bat"}"')
            print(f'  or (Git Bash): source "{_to_sh_path(prefix / "env.sh")}"')
        else:
            print(f'  source "{prefix / "env.sh"}"')
        print()
        if _IS_WINDOWS:
            print(f'  Then:  $env:{args.python_env} /path/to/script.py')
            print(f'         $env:{args.uv_env} run --project /path/to/app script.py')
            print(f'         $env:{args.uvx_env} ruff --version')
        else:
            print(f'  Then:  "${args.python_env}" /path/to/script.py')
            print(f'         "${args.uv_env}" run --project /path/to/app script.py')
            print(f'         "${args.uvx_env}" ruff --version')
        print()


if __name__ == "__main__":
    main()