#!/usr/bin/env python3
"""Lightweight sandboxed execution for code-generation evals (HumanEval, ClassEval).

Replaces the usual Docker-per-suite setup with a cached venv plus whatever
process isolation the host can give us without root:

    tier 1  bubblewrap   read-only root, tmpfs /tmp and $HOME, no network
    tier 2  unshare -rn  no network (home still visible)
    tier 3  rlimits only last resort, still capped on cpu/memory/procs/files

Every tier also applies rlimits and runs the candidate program in a throwaway
working directory as its own process group, so a runaway solution is killed
along with anything it spawned.
"""

import json
import os
import re
import resource
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import venv
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Sequence

DEFAULT_CACHE_ROOT = Path(
    os.environ.get("LLAMA_EVAL_CACHE", Path.home() / ".cache" / "llama-eval")
)

LIMIT_AS_BYTES = 4 * 1024 * 1024 * 1024
LIMIT_CPU_SECONDS_SLACK = 5
LIMIT_NPROC = 256
LIMIT_FSIZE_BYTES = 64 * 1024 * 1024


@dataclass
class ExecResult:
    passed: bool
    status: str
    stdout: str = ""
    stderr: str = ""
    returncode: Optional[int] = None
    duration: float = 0.0

    def summary(self, max_chars: int = 400) -> str:
        """Short human-readable reason, for grader logs."""
        if self.passed:
            return "ok"
        if self.status == "timeout":
            return "timeout"
        tail = (self.stderr or self.stdout).strip()
        if not tail:
            return f"{self.status} (rc={self.returncode})"
        lines = [ln for ln in tail.splitlines() if ln.strip()]
        reason = lines[-1] if lines else tail
        return reason[:max_chars]


def resolve_python(spec: Optional[str]) -> Optional[str]:
    """Turn '3.13' or a path into a usable interpreter, or None if unavailable.

    Suites can need an older interpreter than the one running this tool, because
    a dependency has no wheel (or no working source build) for the newest
    CPython. Accepts an explicit path, a bare version, or an executable name.
    """
    if not spec:
        return None

    candidate = Path(spec).expanduser()
    if candidate.exists() and candidate.is_file():
        return str(candidate)

    uv = shutil.which("uv")
    if uv and re.fullmatch(r"\d+\.\d+(\.\d+)?", spec):
        try:
            r = subprocess.run([uv, "python", "find", spec],
                               capture_output=True, text=True, timeout=120)
            path = r.stdout.strip()
            if r.returncode == 0 and path and Path(path).exists():
                return path
        except Exception:
            pass

    return shutil.which(spec) or shutil.which(f"python{spec}")


def _detect_isolation(probe_python: Optional[str] = None,
                      ro_binds: Sequence[Path] = ()) -> str:
    """Pick the strongest isolation tier this host supports, once.

    The probe must run the *same* interpreter the suite will, with the same
    binds. Probing a bare sys.executable silently downgrades the tier whenever
    this tool is itself run from a venv under $HOME -- which is the usual case,
    since it needs `datasets`. bwrap tmpfs-mounts the home directory, the
    interpreter disappears, the probe fails, and the sandbox falls back to a
    tier that does not hide $HOME at all. Nothing about the results looks wrong
    when that happens, so it has to be got right here.
    """
    py = probe_python or sys.executable
    probe = "import sys; sys.exit(0)"

    if shutil.which("bwrap"):
        try:
            r = subprocess.run(
                _bwrap_argv(Path.cwd(), [py, "-I", "-c", probe], ro_binds=ro_binds),
                capture_output=True, timeout=20,
            )
            if r.returncode == 0:
                return "bwrap"
        except Exception:
            pass

    if shutil.which("unshare"):
        for tier, argv in (
            ("unshare", ["unshare", "--user", "--map-current-user", "--net"]),
            ("unshare-root", ["unshare", "-rn"]),
        ):
            try:
                r = subprocess.run(
                    [*argv, py, "-I", "-c", probe],
                    capture_output=True, timeout=20,
                )
                if r.returncode == 0:
                    return tier
            except Exception:
                continue

    return "rlimit"


UNSHARE_ARGV = {
    "unshare": ["unshare", "--user", "--map-current-user", "--net"],
    "unshare-root": ["unshare", "-rn"],
}


def _under(path: Path, parent: str) -> bool:
    return bool(parent) and (str(path) == parent
                             or str(path).startswith(parent.rstrip("/") + os.sep))


def _bwrap_argv(workdir: Path, inner: Sequence[str],
                ro_binds: Sequence[Path] = ()) -> List[str]:
    home = os.path.expanduser("~")
    argv = [
        "bwrap",
        "--ro-bind", "/", "/",
        "--dev", "/dev",
        "--proc", "/proc",
        "--tmpfs", "/tmp",
    ]
    if home and home != "/":
        argv += ["--tmpfs", home]
    for path in ro_binds:
        if _under(path, home):
            argv += ["--ro-bind", str(path), str(path)]
    argv += [
        "--bind", str(workdir), str(workdir),
        "--chdir", str(workdir),
        "--unshare-all",
        "--die-with-parent",
        "--new-session",
        "--",
    ]
    return argv + list(inner)


def _nproc_ceiling() -> int:
    """A process cap that bounds a fork bomb without tripping on existing load.

    RLIMIT_NPROC is enforced per real UID across the whole system, not per
    process tree, so it has to sit above whatever the user is already running.
    """
    try:
        current = len(os.listdir("/proc"))
    except OSError:
        current = 512
    return max(current * 2, 1024)


def _apply_rlimits(timeout: float, with_nproc: bool,
                   as_bytes: Optional[int] = LIMIT_AS_BYTES):
    def _preexec():
        os.setsid()
        cpu = int(timeout) + LIMIT_CPU_SECONDS_SLACK
        limits = [
            (resource.RLIMIT_CPU, (cpu, cpu)),
            (resource.RLIMIT_FSIZE, (LIMIT_FSIZE_BYTES, LIMIT_FSIZE_BYTES)),
            (resource.RLIMIT_CORE, (0, 0)),
        ]
        if as_bytes:
            limits.append((resource.RLIMIT_AS, (as_bytes, as_bytes)))
        if with_nproc:
            ceiling = _nproc_ceiling()
            limits.append((resource.RLIMIT_NPROC, (ceiling, ceiling)))
        for res, lim in limits:
            try:
                resource.setrlimit(res, lim)
            except (ValueError, OSError):
                pass
    return _preexec


class Sandbox:
    """A cached venv plus an isolated runner for untrusted candidate programs."""

    def __init__(
        self,
        name: str,
        packages: Optional[Sequence[str]] = None,
        cache_root: Optional[Path] = None,
        base_python: Optional[str] = None,
        isolation: Optional[str] = None,
        provision: Optional[Sequence[str]] = None,
        env: Optional[Dict[str, str]] = None,
        address_space_limit: Optional[int] = LIMIT_AS_BYTES,
    ):
        self.name = name
        self.packages = sorted(set(packages or []))
        self.provision = list(provision or [])
        self.env = dict(env or {})
        # None drops the RLIMIT_AS cap entirely -- see _apply_rlimits
        self.address_space_limit = address_space_limit
        self.cache_root = Path(cache_root or DEFAULT_CACHE_ROOT)
        self.base_python = base_python or sys.executable
        self.venv_dir = self.cache_root / f"{name}-venv"
        self._python: Optional[Path] = None
        self._isolation = isolation
        self._base_prefix: List[Path] = []

    @property
    def python(self) -> Path:
        if self._python is None:
            raise RuntimeError("Sandbox.ensure() must be called first")
        return self._python

    @property
    def isolation(self) -> str:
        if self._isolation is None:
            if self._python:
                py = str(self._python)
                roots = [self.venv_dir, *(self._base_prefix or [])]
            else:
                py = sys.executable
                roots = self._detect_base_prefix(Path(py))
            self._isolation = _detect_isolation(py, roots)
        return self._isolation

    def _stamp_path(self) -> Path:
        return self.venv_dir / ".llama-eval-stamp.json"

    def _stamp_matches(self) -> bool:
        try:
            stamp = json.loads(self._stamp_path().read_text())
        except Exception:
            return False
        return (
            stamp.get("packages") == self.packages
            and stamp.get("py") == self._base_version()
            and stamp.get("provision") == self.provision
        )

    def _resolved_env(self) -> Dict[str, str]:
        return {k: v.replace("{venv}", str(self.venv_dir)) for k, v in self.env.items()}

    def _base_version(self) -> str:
        r = subprocess.run(
            [self.base_python, "-c",
             "import sys;print('.'.join(map(str,sys.version_info[:3])))"],
            capture_output=True, text=True,
        )
        return r.stdout.strip()

    def ensure(self, force: bool = False, quiet: bool = False) -> Path:
        """Create/refresh the venv. Idempotent and cheap when already valid."""
        py = self.venv_dir / "bin" / "python"
        if force and self.venv_dir.exists():
            shutil.rmtree(self.venv_dir)

        if py.exists() and self._stamp_matches():
            self._python = py
            self._base_prefix = self._detect_base_prefix(py)
            return py

        if not quiet:
            pkgs = ", ".join(self.packages) if self.packages else "stdlib only"
            print(f"[sandbox] provisioning venv {self.venv_dir} ({pkgs})")

        self.cache_root.mkdir(parents=True, exist_ok=True)
        if self.venv_dir.exists():
            shutil.rmtree(self.venv_dir)

        uv = shutil.which("uv")
        if uv:
            subprocess.run(
                [uv, "venv", "--python", self.base_python, str(self.venv_dir)],
                check=True, capture_output=True,
            )
        elif Path(self.base_python).resolve() == Path(sys.executable).resolve():
            venv.EnvBuilder(with_pip=bool(self.packages), clear=True).create(self.venv_dir)
        else:
            subprocess.run(
                [self.base_python, "-m", "venv", str(self.venv_dir)],
                check=True, capture_output=True,
            )

        if self.packages:
            if uv:
                cmd = [uv, "pip", "install", "--python", str(py), *self.packages]
            else:
                cmd = [str(py), "-m", "pip", "install", "--quiet",
                       "--disable-pip-version-check", *self.packages]
            r = subprocess.run(cmd, capture_output=True, text=True)
            if r.returncode != 0:
                raise RuntimeError(
                    f"failed to install {len(self.packages)} package(s) into "
                    f"{self.venv_dir}:\n{(r.stderr or r.stdout)[-3000:]}"
                )

        for i, snippet in enumerate(self.provision, 1):
            if not quiet:
                print(f"[sandbox] provisioning step {i}/{len(self.provision)}")
            prov_env = dict(os.environ)
            prov_env["SANDBOX_VENV"] = str(self.venv_dir)
            prov_env.update(self._resolved_env())
            r = subprocess.run([str(py), "-c", snippet], env=prov_env,
                               capture_output=True, text=True, timeout=1800)
            if r.returncode != 0:
                raise RuntimeError(
                    f"provisioning step {i} failed:\n{(r.stderr or r.stdout)[-3000:]}"
                )

        self._stamp_path().write_text(json.dumps({
            "packages": self.packages,
            "py": self._base_version(),
            "provision": self.provision,
        }))
        self._python = py
        self._base_prefix = self._detect_base_prefix(py)
        if not quiet:
            print(f"[sandbox] ready ({self.isolation} isolation)")
        if self.isolation == "unshare-root":
            print("Warning: this host's unshare cannot keep the caller's uid, so "
                  "candidate code runs as root inside the namespace. Tests that "
                  "assert a permission denial will not see one. Install "
                  "bubblewrap (bwrap) for correct results.")
        return py

    @staticmethod
    def _detect_base_prefix(py: Path) -> List[Path]:
        """Every directory that must stay visible for the venv to start.

        A venv is only symlinks into its base interpreter, so the whole install
        tree has to survive the tmpfs that hides $HOME -- and binding just the
        resolved tree is not enough. Version-managed interpreters (uv, pyenv)
        are commonly reached through an unversioned alias symlink
        (cpython-3.13-... -> cpython-3.13.14-...), and startup resolves through
        the alias path, so both spellings need to be mounted. Walk the chain
        and keep every install root it passes through.
        """
        roots: List[Path] = []

        def add(path: Path):
            if path.exists() and path not in roots:
                roots.append(path)

        cur, hops = py, 0
        while hops < 16:
            hops += 1
            add(cur.parent.parent)
            if cur.is_symlink():
                target = Path(os.readlink(cur))
                cur = target if target.is_absolute() else (cur.parent / target)
            else:
                break

        try:
            r = subprocess.run([str(py), "-c", "import sys; print(sys.base_prefix)"],
                               capture_output=True, text=True, timeout=60)
            if r.returncode == 0 and r.stdout.strip():
                add(Path(r.stdout.strip()))
        except Exception:
            pass

        return [p for p in roots if str(p) not in ("/", "/usr", "/usr/local")]

    def _isolate(self, argv: Sequence[str], workdir: Path,
                 ro_binds: Sequence[Path] = ()) -> List[str]:
        """Wrap a command in the strongest isolation tier this host supports."""
        if self.isolation == "bwrap":
            return _bwrap_argv(
                workdir, argv,
                ro_binds=[self.venv_dir, *(self._base_prefix or []), *ro_binds],
            )
        if self.isolation in UNSHARE_ARGV:
            return [*UNSHARE_ARGV[self.isolation], *argv]
        return list(argv)

    def _base_env(self, workdir: Path) -> Dict[str, str]:
        env = {
            "PATH": "/usr/bin:/bin",
            "HOME": str(workdir),
            "TMPDIR": str(workdir),
            "LC_ALL": "C.UTF-8",
            "PYTHONHASHSEED": "0",
            "PYTHONDONTWRITEBYTECODE": "1",
            "OPENBLAS_NUM_THREADS": "1",
            "OMP_NUM_THREADS": "1",
        }
        env.update(self._resolved_env())
        return env

    def run_argv(self, argv: Sequence[str], workdir: Path,
                 timeout: float = 15.0, env: Optional[Dict[str, str]] = None,
                 ro_binds: Sequence[Path] = ()) -> ExecResult:
        """Run an arbitrary command in `workdir`, which stays writable.

        Unlike run(), the caller owns the directory and it survives the call.
        That is what a multi-step agent needs: its edits have to persist across
        many tool invocations before anything is graded. Taking an argv rather
        than a source string is what lets a suite drive a linter or a test
        runner for a language other than Python.
        """
        run_env = self._base_env(workdir)
        if env:
            run_env.update(env)

        started = time.monotonic()
        try:
            proc = subprocess.Popen(
                self._isolate(argv, workdir, ro_binds),
                cwd=str(workdir),
                stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                errors="replace",
                env=run_env,
                preexec_fn=_apply_rlimits(timeout,
                                          with_nproc=self.isolation == "rlimit",
                                          as_bytes=self.address_space_limit),
            )
        except Exception as e:
            return ExecResult(False, "error", "", f"{type(e).__name__}: {e}")

        try:
            out, err = proc.communicate(timeout=timeout)
        except subprocess.TimeoutExpired:
            _kill_tree(proc)
            out, err = proc.communicate()
            return ExecResult(False, "timeout", out or "", err or "",
                              None, time.monotonic() - started)

        duration = time.monotonic() - started
        if proc.returncode == 0:
            return ExecResult(True, "ok", out, err, 0, duration)
        return ExecResult(False, "failed", out, err, proc.returncode, duration)

    def run(self, code: str, timeout: float = 15.0,
            env: Optional[Dict[str, str]] = None) -> ExecResult:
        """Run `code` to completion in the sandbox. Exit 0 == passed."""
        workdir = Path(tempfile.mkdtemp(prefix=f"{self.name}-", dir="/tmp"))
        try:
            prog = workdir / "candidate.py"
            prog.write_text(code)
            return self.run_argv([str(self.python), "-I", "-B", str(prog)],
                                 workdir, timeout=timeout, env=env)
        except Exception as e:
            return ExecResult(False, "error", "", f"{type(e).__name__}: {e}")
        finally:
            shutil.rmtree(workdir, ignore_errors=True)


def _kill_tree(proc: subprocess.Popen):
    """Kill the candidate and everything it spawned."""
    for sig in (signal.SIGKILL,):
        try:
            os.killpg(os.getpgid(proc.pid), sig)
        except (ProcessLookupError, PermissionError):
            try:
                proc.kill()
            except ProcessLookupError:
                pass