"""Run a short code snippet in an ephemeral working directory.

This is the *second* track of the code-preview feature and deliberately not the
first. `render_preview` (synapse/tools.py) executes nothing server-side: it
validates markup and the chat UI renders it inside an opaque-origin iframe.
That model fits HTML/SVG/JSX and cannot fit C, Rust or Erlang, which need a real
toolchain. So those go through here instead, and the result is shown as terminal
output rather than a rendered document.

WHAT THIS IS NOT
----------------
Not a security sandbox. Snippets run as the current user on the host. What this
module actually provides is *containment by limits*, layered:

  1. Consent      run_snippet is an ACTION tool, so it is withheld entirely
                  unless `action_tool_policy` is "ask"/"allow" — and on "ask"
                  every call waits for the user's Approve/Deny in chat.
  2. Screening    `critique` rejects the obvious-abuse shapes (sockets, process
                  spawning, absolute paths) before anything is written to disk.
  3. Isolation    cwd is a fresh temp dir that is deleted afterwards; HOME and
                  TMPDIR point at it; the environment is scrubbed to a small
                  allowlist.
  4. Limits       wall-clock timeout, RLIMIT_CPU/AS/FSIZE/NPROC on POSIX,
                  truncated output.
  5. Network      on Linux, `unshare -rn` when unprivileged user namespaces are
                  available (probed once, see `_net_isolation`). Nowhere else —
                  macOS and Windows get no network isolation at all.

Layer 2 is a tripwire, not a boundary: arbitrary Python can evade any regex
trivially. It exists to catch a model that reaches for `requests` by habit, not
an adversary. The load-bearing layers are 1 and 3–5. Anyone wanting a real
boundary should run this under a container or a VM; that is a deployment
decision this module does not make for them.
"""
from __future__ import annotations

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

# Wall clock. Compilation gets its own, larger budget: rustc on a cold cache
# routinely spends longer than any snippet is allowed to *run*, and killing a
# compile at the run timeout would look like the code hung when it never started.
RUN_TIMEOUT = 5.0
COMPILE_TIMEOUT = 25.0

MAX_SOURCE = 100_000      # chars of source accepted
MAX_OUTPUT = 20_000       # chars of stdout/stderr returned, per stream
MAX_STDIN = 10_000

# Applied to the run step only. A compiler legitimately needs more address space
# than a snippet does and forks a linker, so imposing these on the compile step
# breaks the toolchain rather than containing the snippet.
_MEM_BYTES = 512 * 1024 * 1024
_MAX_PROCS = 64
_MAX_FILE_BYTES = 8 * 1024 * 1024

# RLIMIT_NPROC is counted per real UID, not per run. An absolute "512" is
# therefore 512 minus whatever the user already has — enough on a quiet host,
# zero on a busy macOS desktop. Erlang needs relative headroom for BEAM's
# boot process tree; other languages keep the absolute _MAX_PROCS (blocking
# fork is the intent there).
_ERLANG_PROC_HEADROOM = 256


# ---------------------------------------------------------------------------
# Static screening
# ---------------------------------------------------------------------------

def _shared_issues(source: str) -> list[str]:
    stripped = source.strip()
    if len(stripped) < 8:
        return ["source is empty or too short to run."]
    issues: list[str] = []
    # Only flag URLs inside string literals. A citation in a comment
    # (`/* see https://… */`) is not a network reach and bouncing it costs a
    # useless retry round — this layer is a tripwire, not a parser.
    if re.search(r"""['"][^'"]*https?://[^'"]*['"]""", source):
        issues.append(
            "source contains an http(s) URL string — the runner has no network "
            "access. Inline the data you need."
        )
    if re.search(r"\b(TODO|FIXME|your code here|implement this)\b", source, re.I):
        issues.append("source still contains a placeholder — send the finished code.")
    # An absolute path is either reaching outside the ephemeral dir or is a
    # machine-specific guess that will not exist. Relative paths are fine: cwd
    # is the temp dir and goes away with it.
    # Windows: match `C:\…` and `C:\\…` (raw / escaped). Unix: common roots.
    if re.search(
        r"""['"](?:/(?:etc|home|root|usr|var|proc|sys)/|[A-Za-z]:[/\\])""",
        source,
    ):
        issues.append(
            "source references an absolute filesystem path. The snippet runs in a "
            "throwaway directory — use relative paths, or inline the data."
        )
    return issues


def _deny(source: str, rules: list[tuple[str, str]]) -> list[str]:
    return [msg for pattern, msg in rules if re.search(pattern, source)]


_NO_NET = "the runner has no network access"
_NO_SPAWN = "the runner does not allow spawning other processes"

_PY_RULES = [
    (r"\b(?:import|from)\s+(?:socket|ssl|ftplib|smtplib|telnetlib|urllib|http)\b",
     f"networking module imported — {_NO_NET}."),
    (r"\bimport\s+(?:requests|httpx|aiohttp|urllib3)\b",
     f"HTTP client imported — {_NO_NET}."),
    (r"\b(?:import|from)\s+(?:subprocess|multiprocessing)\b",
     f"subprocess/multiprocessing imported — {_NO_SPAWN}."),
    (r"\bos\.(?:system|popen|exec[lv]|fork|spawn|kill)\b",
     f"os process call — {_NO_SPAWN}."),
    (r"\b(?:import|from)\s+ctypes\b",
     "ctypes imported — the runner does not allow native calls."),
]

_C_RULES = [
    (r"#\s*include\s*<(?:sys/socket|netinet/|arpa/|netdb)",
     f"socket header included — {_NO_NET}."),
    (r"\b(?:system|popen|fork|execv?[lpe]*)\s*\(",
     f"process call — {_NO_SPAWN}."),
]

_RUST_RULES = [
    (r"\bstd::net\b", f"std::net used — {_NO_NET}."),
    (r"\bstd::process::(?:Command|abort)\b", f"std::process::Command used — {_NO_SPAWN}."),
]

_ERL_RULES = [
    (r"\b(?:gen_tcp|gen_udp|httpc|inets|ssl)\b", f"networking module used — {_NO_NET}."),
    (r"\bos:cmd\b", f"os:cmd/1 used — {_NO_SPAWN}."),
]


def _entry_point(pattern: str, message: str):
    """Most of these languages fail with a linker/loader error rather than a
    useful one when the entry point is missing, so name it up front."""
    def check(source: str) -> list[str]:
        return [] if re.search(pattern, source) else [message]
    return check


def _critique(rules: list[tuple[str, str]], entry=None):
    def check(source: str) -> list[str]:
        issues = _shared_issues(source)
        if issues and issues[0].startswith("source is empty"):
            return issues  # nothing else is worth saying about an empty body
        issues += _deny(source, rules)
        if entry:
            issues += entry(source)
        return issues
    return check


# ---------------------------------------------------------------------------
# Drivers
# ---------------------------------------------------------------------------

def _source_name(default: str):
    return lambda _source: default


def _erlang_source_name(source: str) -> str:
    """escript reads the same file two different ways, chosen by extension: a
    `.erl` file is compiled as a module (needs -module/-export), anything else
    is a plain script (needs only main/1). Models write both, so let the source
    pick its own filename instead of forcing one dialect."""
    return "main.erl" if re.search(r"^\s*-module\s*\(", source, re.M) else "main.escript"


def _erlang_write_source(source: str) -> str:
    """OTP 28+ escript rejects a shebang-less `.escript` with 'Premature end of
    file'. Module-form `.erl` files do not need one. Inject only when missing so
    a model that already wrote `#!/usr/bin/env escript` is left alone."""
    if _erlang_source_name(source).endswith(".escript"):
        stripped = source.lstrip()
        if not stripped.startswith("#!"):
            return "#!/usr/bin/env escript\n" + source
    return source


# The one place that says which languages can be executed. Each entry owns that
# language's screening, toolchain probe and argv. The tool schema's `lang` enum,
# the capability line in the system prompt and the dispatch below are all derived
# from these keys rather than repeating them.
#
# The frontend keeps a matching registry (RUN_LANGS in
# interface/web/src/preview/run-langs.js) because the two sides need different
# things per language — this side executes, that side labels and displays — and
# neither should depend on the other at runtime. tests/test_tools.py asserts the
# key sets stay equal, so drift fails the check gate instead of silently showing
# a run result with the wrong language on it.
RUN_LANGS: dict[str, dict] = {
    "python": {
        "summary": "Python script (stdlib only)",
        "tool": lambda: sys.executable,
        "install": "Python is bundled with NexusOS; this should not happen.",
        "source_name": _source_name("main.py"),
        # -I is isolated mode: ignores PYTHON* env vars, the user site-packages
        # dir and the script's own directory on sys.path.
        "compile": None,
        "run": lambda tool, src, _exe: [tool, "-I", src],
        "critique": _critique(_PY_RULES),
    },
    "c": {
        "summary": "single-file C program (C11, libm linked)",
        "tool": lambda: shutil.which("cc") or shutil.which("gcc") or shutil.which("clang"),
        "install": "install a C compiler (clang or gcc)",
        "source_name": _source_name("main.c"),
        "compile": lambda cc, src, exe: [cc, "-std=c11", "-O0", "-Wall", "-o", exe, src, "-lm"],
        "run": lambda _tool, _src, exe: [exe],
        "critique": _critique(_C_RULES, _entry_point(
            r"\bmain\s*\(", "no main() — a C program needs `int main(void)`.")),
    },
    "cpp": {
        "summary": "single-file C++ program (C++17)",
        "tool": lambda: shutil.which("c++") or shutil.which("g++") or shutil.which("clang++"),
        "install": "install a C++ compiler (clang++ or g++)",
        "source_name": _source_name("main.cpp"),
        "compile": lambda cc, src, exe: [cc, "-std=c++17", "-O0", "-Wall", "-o", exe, src],
        "run": lambda _tool, _src, exe: [exe],
        "critique": _critique(_C_RULES, _entry_point(
            r"\bmain\s*\(", "no main() — a C++ program needs `int main()`.")),
    },
    "rust": {
        "summary": "single-file Rust program (2021 edition, std only)",
        "tool": lambda: shutil.which("rustc"),
        "install": "install Rust (https://rustup.rs)",
        "source_name": _source_name("main.rs"),
        # Debug build: -O roughly triples compile time for snippets that run for
        # milliseconds either way.
        "compile": lambda cc, src, exe: [cc, "--edition", "2021", "-o", exe, src],
        "run": lambda _tool, _src, exe: [exe],
        "critique": _critique(_RUST_RULES, _entry_point(
            r"\bfn\s+main\s*\(", "no main() — a Rust program needs `fn main()`.")),
    },
    "erlang": {
        "summary": "escript program with a main/1 entry point",
        "tool": lambda: shutil.which("escript"),
        "install": "install Erlang/OTP (provides escript)",
        "source_name": _erlang_source_name,
        "compile": None,
        "run": lambda tool, src, _exe: [tool, src],
        "critique": _critique(_ERL_RULES, _entry_point(
            r"\bmain\s*\(", "no main/1 — escript calls `main(Args)`.")),
    },
}

# Spellings a model reaches for that aren't the registry key. Resolved before
# lookup so `c++` and `py` work without doubling the tool schema's enum.
ALIASES = {
    "c++": "cpp", "cc": "c", "py": "python", "python3": "python",
    "rs": "rust", "erl": "erlang", "escript": "erlang",
}


def resolve_lang(lang: str) -> str:
    key = (lang or "").strip().lower()
    return ALIASES.get(key, key)


def lang_prose() -> str:
    """'python, c or rust' — the runnable languages as a phrase for prompts."""
    names = list(RUN_LANGS)
    if len(names) < 2:
        return names[0] if names else ""
    return f"{', '.join(names[:-1])} or {names[-1]}"


# ---------------------------------------------------------------------------
# Execution
# ---------------------------------------------------------------------------

# Environment variables the child keeps. Everything else is dropped: the snippet
# has no business seeing API keys, proxy settings or the user's shell config, and
# PYTHON*/LD_* in particular would let ambient config change how it runs.
_ENV_KEEP = ("PATH", "LANG", "LC_ALL", "TERM", "SYSTEMROOT", "COMSPEC")

# rustc is usually a rustup shim, and a shim with HOME rewritten cannot find its
# own toolchain. Passing these through is what makes Rust work at all here; they
# point at read-only toolchain data, not at anything the snippet should write.
_ENV_KEEP_BY_LANG = {"rust": ("RUSTUP_HOME", "CARGO_HOME", "RUSTUP_TOOLCHAIN")}


def _child_env(lang: str, workdir: Path) -> dict:
    env = {k: os.environ[k] for k in _ENV_KEEP if k in os.environ}
    for k in _ENV_KEEP_BY_LANG.get(lang, ()):
        if k in os.environ:
            env[k] = os.environ[k]
    if lang == "rust" and "RUSTUP_HOME" not in env:
        # HOME is about to be rewritten, so resolve rustup's default location
        # against the real home while we still know it.
        default = Path.home() / ".rustup"
        if default.is_dir():
            env["RUSTUP_HOME"] = str(default)
            env.setdefault("CARGO_HOME", str(Path.home() / ".cargo"))
    env["HOME"] = str(workdir)
    env["TMPDIR"] = str(workdir)
    env.setdefault("LC_ALL", "C.UTF-8")
    return env


def _user_process_count() -> int | None:
    """How many processes this UID already owns. None when we cannot count
    (Windows, or psutil missing) — callers must not invent an absolute cap."""
    if os.name == "nt" or not hasattr(os, "getuid"):
        return None
    try:
        import psutil  # optional; process extra
    except ImportError:
        return None
    uid = os.getuid()
    n = 0
    for proc in psutil.process_iter(["uids"]):
        try:
            uids = proc.info.get("uids")
            if uids is not None and getattr(uids, "real", None) == uid:
                n += 1
        except (psutil.Error, TypeError, AttributeError):
            continue
    return n


def _max_procs_for(lang: str) -> int | None:
    """Soft RLIMIT_NPROC for this run, or None to leave the limit unset.

    Erlang: current per-UID count + headroom (RLIMIT_NPROC is UID-scoped).
    Everything else: the absolute _MAX_PROCS tripwire against fork bombs.
    """
    if lang == "erlang":
        current = _user_process_count()
        if current is None:
            return None
        return current + _ERLANG_PROC_HEADROOM
    return _MAX_PROCS


def _limits(constrain_memory: bool, max_procs: int | None = None):
    """preexec_fn applying POSIX rlimits, or None where they don't exist.

    RLIMIT_CPU is a backstop for the wall-clock timeout: a snippet that ignores
    SIGTERM still loses the CPU. The memory and process caps are skipped for
    compilation — see _MEM_BYTES.

    `max_procs=None` with constrain_memory means "do not set RLIMIT_NPROC"
    (used when we cannot compute a relative Erlang ceiling). Passing an int
    always sets it.
    """
    if sys.platform == "win32":
        return None
    try:
        import resource
    except ImportError:  # pragma: no cover - POSIX only
        return None

    cpu = int(COMPILE_TIMEOUT if not constrain_memory else RUN_TIMEOUT) + 1
    wanted = [("RLIMIT_CPU", cpu), ("RLIMIT_FSIZE", _MAX_FILE_BYTES), ("RLIMIT_CORE", 0)]
    if constrain_memory:
        wanted.append(("RLIMIT_AS", _MEM_BYTES))
        # Distinguish "caller omitted" (use default) from "explicitly skip"
        # by requiring the kw to be passed — see _spawn.
        if max_procs is not None:
            wanted.append(("RLIMIT_NPROC", max_procs))

    def apply():  # runs in the forked child, between fork and exec
        # Every limit is set independently and failure is swallowed. Which of
        # these exist, and which can be lowered, varies by platform (macOS has
        # no usable RLIMIT_AS, RLIMIT_NPROC is absent on some POSIX systems) —
        # and an exception raised here does not "skip a limit", it aborts the
        # spawn entirely. Partial limits are the right failure mode; no run at
        # all is not.
        for name, soft in wanted:
            which = getattr(resource, name, None)
            if which is None:
                continue
            try:
                _, hard = resource.getrlimit(which)
                if hard != resource.RLIM_INFINITY:
                    soft = min(soft, hard)
                resource.setrlimit(which, (soft, hard))
            except (ValueError, OSError):
                continue

    return apply


_net_isolation_cache: list | None = None


def _net_isolation() -> list[str]:
    """argv prefix that drops the child into an empty network namespace, or [].

    Linux only, and only where unprivileged user namespaces are enabled — which
    is a kernel/distro setting we can't change and shouldn't fail over. Probed
    once and cached; an empty list means the run simply has host networking, and
    callers must not treat this as a guarantee either way.
    """
    global _net_isolation_cache
    if _net_isolation_cache is not None:
        return _net_isolation_cache
    _net_isolation_cache = []
    if sys.platform.startswith("linux") and shutil.which("unshare"):
        try:
            probe = subprocess.run(
                ["unshare", "-rn", "true"],
                capture_output=True, timeout=5,
            )
            if probe.returncode == 0:
                _net_isolation_cache = ["unshare", "-rn"]
        except (OSError, subprocess.SubprocessError):
            pass
    return _net_isolation_cache


def _clip(raw: bytes) -> str:
    text = raw.decode("utf-8", errors="replace")
    if len(text) <= MAX_OUTPUT:
        return text
    return text[:MAX_OUTPUT] + f"\n... [truncated at {MAX_OUTPUT} characters]"


def _spawn(argv: list[str], workdir: Path, env: dict, timeout: float,
           stdin: str = "", constrain_memory: bool = True,
           max_procs: int | None = _MAX_PROCS):
    """Spawn a child. `max_procs` defaults to `_MAX_PROCS`; pass `None` to skip
    RLIMIT_NPROC entirely (Erlang, when a relative ceiling cannot be computed)."""
    return subprocess.run(
        argv,
        cwd=str(workdir),
        env=env,
        input=stdin.encode("utf-8"),
        capture_output=True,
        timeout=timeout,
        preexec_fn=_limits(constrain_memory, max_procs=max_procs),
    )


def run(lang: str, source: str, stdin: str = "") -> dict:
    """Compile (if needed) and run `source`. Blocking — call from a thread.

    Returns {ok, lang, stage, exit_code, stdout, stderr, error}. `ok` is False
    only when the snippet could not be run at all (missing toolchain, compile
    error, timeout); a program that runs and exits non-zero is a successful run
    with a non-zero exit_code, because its stderr is the answer the user wants.
    """
    key = resolve_lang(lang)
    entry = RUN_LANGS.get(key)
    if entry is None:
        return {"ok": False, "lang": lang, "stage": "lang",
                "error": f"cannot run {lang!r} — use {lang_prose()}."}

    tool = entry["tool"]()
    if not tool:
        return {"ok": False, "lang": key, "stage": "toolchain",
                "error": f"no toolchain for {key} on this machine — {entry['install']}. "
                         "Show the code instead of running it."}

    with tempfile.TemporaryDirectory(prefix="nexus-run-") as tmp:
        workdir = Path(tmp)
        body = _erlang_write_source(source) if key == "erlang" else source
        src = workdir / entry["source_name"](source)
        src.write_text(body, encoding="utf-8")
        exe = str(workdir / ("program.exe" if sys.platform == "win32" else "program"))
        env = _child_env(key, workdir)
        max_procs = _max_procs_for(key)

        if entry["compile"]:
            try:
                built = _spawn(entry["compile"](tool, str(src), exe), workdir, env,
                               COMPILE_TIMEOUT, constrain_memory=False)
            except subprocess.TimeoutExpired:
                return {"ok": False, "lang": key, "stage": "compile",
                        "error": f"compilation timed out after {COMPILE_TIMEOUT:g}s."}
            except OSError as e:
                return {"ok": False, "lang": key, "stage": "compile",
                        "error": f"could not start the compiler: {e}"}
            if built.returncode != 0:
                return {"ok": False, "lang": key, "stage": "compile",
                        "exit_code": built.returncode,
                        "stdout": _clip(built.stdout), "stderr": _clip(built.stderr),
                        "error": "compilation failed — read stderr, fix the source, "
                                 "and call run_snippet again."}

        argv = _net_isolation() + entry["run"](tool, str(src), exe)
        try:
            done = _spawn(argv, workdir, env, RUN_TIMEOUT, stdin=stdin[:MAX_STDIN],
                          max_procs=max_procs)
        except subprocess.TimeoutExpired as e:
            return {"ok": False, "lang": key, "stage": "run",
                    "stdout": _clip(e.stdout or b""), "stderr": _clip(e.stderr or b""),
                    "error": f"the program did not finish within {RUN_TIMEOUT:g}s — "
                             "it is probably looping. Bound the work and try again."}
        except OSError as e:
            return {"ok": False, "lang": key, "stage": "run",
                    "error": f"could not start the program: {e}"}

    stdout = _clip(done.stdout)
    stderr = _clip(done.stderr)
    # BEAM failing to fork under RLIMIT_NPROC looks like a snippet bug if we
    # report ok=True. Call it out as a runner limit so the model does not keep
    # rewriting a correct program.
    if (
        key == "erlang"
        and done.returncode != 0
        and "Resource temporarily unavailable" in stderr
    ):
        return {
            "ok": False,
            "lang": key,
            "stage": "run",
            "exit_code": done.returncode,
            "stdout": stdout,
            "stderr": stderr,
            "error": (
                "Erlang could not start under the process limit (host already has "
                "many processes for this user). Free some processes and retry, or "
                "show the code instead of running it."
            ),
        }

    return {
        "ok": True,
        "lang": key,
        "stage": "run",
        "exit_code": done.returncode,
        "stdout": stdout,
        "stderr": stderr,
    }


def critique(lang: str, source: str) -> list[str]:
    """Static screening for one snippet. See the module docstring on what this
    is worth: a tripwire against habitual network/process reaches, not a
    boundary."""
    key = resolve_lang(lang)
    entry = RUN_LANGS.get(key)
    if entry is None:
        return [f"cannot run {lang!r} — use {lang_prose()}."]
    if len(source) > MAX_SOURCE:
        return [f"source is over {MAX_SOURCE} characters — send a single focused snippet."]
    return entry["critique"](source)
