"""Isolated execution of untrusted Python and Rust.

The app compiles and runs code that arbitrary visitors typed into a web page.
Running that in-process would hand every visitor the host's environment - which
on a deployed instance includes the API keys. Everything here exists to keep
those two worlds apart.

Guarantees, strongest first:

  1. Untrusted code runs in a *separate process* with an explicitly rebuilt
     environment. No API key, token, or secret is inherited - the child gets a
     hand-written allowlist of variables (PATH and the few the OS needs to
     start a process at all), never a copy of os.environ.
  2. Every run is wall-clock bounded and the whole process tree is killed on
     timeout, so an infinite loop cannot pin a worker forever.
  3. Work happens in a throwaway temp directory that is deleted afterwards, so
     nothing can accumulate on disk or read the repository.
  4. On Linux (how this deploys) POSIX rlimits additionally cap CPU seconds,
     address space, file size, and forked processes.

Item 4 is unavailable on Windows, so local dev is protected by 1-3 only. That
is fine for running your own code; the deployed container is where hostile
input actually arrives, and there the full set applies.

What this is NOT: a security boundary strong enough to run genuinely malicious
code with no further thought. A determined attacker inside the container is
still inside the container. The container is therefore the real boundary - the
process-level work here reduces blast radius, it does not replace isolation at
the platform layer.
"""

from __future__ import annotations

import os
import platform
import shutil
import subprocess
import sys
import tempfile
import time
from dataclasses import dataclass

IS_WINDOWS = platform.system() == "Windows"

# Wall-clock ceilings, seconds.
PYTHON_TIMEOUT = int(os.getenv("PYTHON_TIMEOUT", "60"))
COMPILE_TIMEOUT = int(os.getenv("COMPILE_TIMEOUT", "90"))
RUN_TIMEOUT = int(os.getenv("RUN_TIMEOUT", "60"))

# POSIX-only resource ceilings.
MAX_CPU_SECONDS = int(os.getenv("MAX_CPU_SECONDS", "60"))
# rustc with -C lto=fat is memory-hungry; 2 GB was tight enough to risk killing
# legitimate compiles. The free Spaces container has 16 GB, so this is generous
# without being unbounded.
MAX_MEMORY_MB = int(os.getenv("MAX_MEMORY_MB", "4096"))
MAX_OUTPUT_CHARS = 20_000


@dataclass
class Result:
    """Outcome of one sandboxed run."""

    ok: bool
    stdout: str
    stderr: str
    seconds: float
    stage: str = "run"  # "compile" or "run" - which step produced the error

    @property
    def display(self) -> str:
        """Human-readable text for the UI output box."""
        if self.ok:
            return self.stdout or "(the program produced no output)"
        label = "Compilation failed" if self.stage == "compile" else "Execution failed"
        detail = (self.stderr or self.stdout or "").strip() or "no diagnostics available"
        return f"{label}\n\n{detail}"


def _safe_env() -> dict[str, str]:
    """Build a minimal environment from scratch.

    Deliberately constructed by allowlist rather than by copying os.environ and
    deleting keys - a denylist silently fails the moment a new provider adds an
    env var, and that failure mode leaks a credential.
    """
    allowed: dict[str, str] = {}

    # The bare minimum for a process to start and find a toolchain.
    passthrough = ["PATH"]
    if IS_WINDOWS:
        # Windows cannot create a process without these.
        passthrough += ["SYSTEMROOT", "SYSTEMDRIVE", "COMSPEC", "PATHEXT",
                        "NUMBER_OF_PROCESSORS", "PROCESSOR_ARCHITECTURE",
                        "USERPROFILE", "WINDIR"]
    else:
        passthrough += ["LANG", "LC_ALL", "HOME"]

    for name in passthrough:
        value = os.environ.get(name)
        if value:
            allowed[name] = value

    # Keep the child's Python quiet and deterministic.
    allowed["PYTHONUNBUFFERED"] = "1"
    allowed["PYTHONDONTWRITEBYTECODE"] = "1"

    return allowed


def assert_safe_to_deploy() -> None:
    """Refuse to serve the public if secrets are sitting on disk.

    Scrubbing the child's environment stops submitted code from *inheriting* a
    key, but it cannot stop it from calling open('.env'). Verified by test: a
    sandboxed program read the dotenv file in plaintext.

    On a public host that path has to be closed by not having the file at all -
    secrets belong in the platform's secret store, injected as environment
    variables into the parent only. This turns that requirement into a
    startup check instead of a thing to remember.
    """
    if os.getenv("PUBLIC_DEPLOYMENT", "0") != "1":
        return

    here = os.path.dirname(os.path.abspath(__file__))
    offenders = [
        name for name in (".env", ".env.local", ".env.production")
        if os.path.exists(os.path.join(here, name))
    ]
    if offenders:
        raise RuntimeError(
            "Refusing to start in public mode with secret files on disk: "
            f"{', '.join(offenders)}.\n"
            "Submitted code runs on this filesystem and can read them. "
            "Delete the file and supply secrets as environment variables "
            "(on Hugging Face Spaces: Settings -> Variables and secrets)."
        )


def _preexec():
    """Apply POSIX resource limits in the child before exec. Linux/macOS only."""
    import resource

    resource.setrlimit(resource.RLIMIT_CPU, (MAX_CPU_SECONDS, MAX_CPU_SECONDS))
    mem = MAX_MEMORY_MB * 1024 * 1024
    resource.setrlimit(resource.RLIMIT_AS, (mem, mem))
    resource.setrlimit(resource.RLIMIT_FSIZE, (64 * 1024 * 1024, 64 * 1024 * 1024))

    # RLIMIT_NPROC is deliberately NOT set.
    #
    # It counts processes per *user*, not per process tree, so any cap here also
    # counts every thread the host container already runs under this uid. In
    # practice that starved the compilers of threads:
    #   rustc: Os { code: 11, kind: WouldBlock, ... }
    #   g++:   cannot execute cc1plus: vfork: Resource temporarily unavailable
    # Both are EAGAIN from the fork itself. Because the limit is shared with
    # processes we did not start and cannot count, any value is either too low
    # to compile or too high to matter - it is not a usable control here.
    #
    # A fork bomb is still bounded by RLIMIT_CPU, RLIMIT_AS, and the wall-clock
    # kill, which terminates the entire process group.
    # New process group so a timeout kill takes any children with it.
    os.setsid()


def _popen_kwargs() -> dict:
    kwargs: dict = {}
    if IS_WINDOWS:
        # Lets us terminate the whole tree on timeout.
        kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
    else:
        kwargs["preexec_fn"] = _preexec
    return kwargs


def _truncate(text: str) -> str:
    if text is None:
        return ""
    if len(text) <= MAX_OUTPUT_CHARS:
        return text
    return text[:MAX_OUTPUT_CHARS] + f"\n\n... output truncated at {MAX_OUTPUT_CHARS:,} characters ..."


def _kill_tree(proc: subprocess.Popen) -> None:
    """Kill a process and everything it spawned."""
    try:
        if IS_WINDOWS:
            subprocess.run(["taskkill", "/F", "/T", "/PID", str(proc.pid)],
                           capture_output=True, timeout=10)
        else:
            import signal
            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
    except Exception:
        pass
    try:
        proc.wait(timeout=5)
    except Exception:
        pass


def _run(cmd: list[str], cwd: str, timeout: int, stage: str) -> Result:
    """Run one command under the full set of restrictions."""
    started = time.perf_counter()
    try:
        proc = subprocess.Popen(
            cmd,
            cwd=cwd,
            env=_safe_env(),
            stdin=subprocess.DEVNULL,   # never block waiting on input
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            **_popen_kwargs(),
        )
    except FileNotFoundError as e:
        return Result(False, "", f"Toolchain not found: {e}", 0.0, stage)

    try:
        stdout, stderr = proc.communicate(timeout=timeout)
    except subprocess.TimeoutExpired:
        _kill_tree(proc)
        elapsed = time.perf_counter() - started
        return Result(
            False, "",
            f"Timed out after {timeout}s and was terminated.\n"
            f"The sandbox caps how long submitted code may run.",
            elapsed, stage,
        )

    elapsed = time.perf_counter() - started
    return Result(proc.returncode == 0, _truncate(stdout), _truncate(stderr), elapsed, stage)


# --------------------------------------------------------------------------
# public API
# --------------------------------------------------------------------------


def run_python(code: str) -> Result:
    """Execute untrusted Python in a subprocess.

    The notebook original used exec() in the web server's own process with full
    builtins, which let any visitor read os.environ - and therefore every API
    key on the host. This runs it out-of-process against a scrubbed environment
    instead.
    """
    if not code.strip():
        return Result(False, "", "No Python code to run.", 0.0)

    with tempfile.TemporaryDirectory(prefix="pyrun_") as tmp:
        path = os.path.join(tmp, "program.py")
        with open(path, "w", encoding="utf-8") as f:
            f.write(code)
        # -I: isolated mode. Ignores PYTHON* env vars and keeps the script's
        # directory off sys.path, so it cannot import anything from this repo.
        return _run([sys.executable, "-I", "-S", path], tmp, PYTHON_TIMEOUT, "run")


@dataclass(frozen=True)
class Language:
    """A compiled target the app can translate Python into."""

    key: str            # internal id
    display: str        # shown in the UI
    extension: str      # source file suffix
    highlight: str      # gradio Code component language
    compilers: tuple[str, ...]   # executables to look for, in order of preference
    install_hint: str

    def find_compiler(self) -> str | None:
        for name in self.compilers:
            found = shutil.which(name)
            if found:
                return found
        return None


LANGUAGES: dict[str, Language] = {
    "rust": Language(
        key="rust",
        display="Rust",
        extension="rs",
        highlight="cpp",     # gradio has no rust mode; cpp reads acceptably
        compilers=("rustc",),
        install_hint="https://rustup.rs",
    ),
    "cpp": Language(
        key="cpp",
        display="C++",
        extension="cpp",
        highlight="cpp",
        compilers=("g++", "clang++", "c++"),
        install_hint="https://gcc.gnu.org or https://winlibs.com on Windows",
    ),
}

DEFAULT_LANGUAGE = "rust"


def build_compile_command(language: Language, compiler: str, source: str,
                          binary: str) -> list[str]:
    """Optimisation flags per toolchain.

    Both sets deliberately trade compile time for runtime - whole-program
    optimisation, native CPU targeting, no unwinding machinery where it can be
    dropped. Compile time is excluded from the measurement, so this is free.
    """
    if language.key == "rust":
        return [
            compiler, source,
            "-C", "opt-level=3",
            "-C", "target-cpu=native",
            "-C", "codegen-units=1",
            "-C", "lto=fat",
            "-C", "panic=abort",
            "-C", "strip=symbols",
            "-o", binary,
        ]
    return [
        compiler, source,
        "-std=c++17",
        "-O3",
        "-march=native",
        "-funroll-loops",
        "-fno-plt",
        "-flto",
        "-o", binary,
    ]


def run_compiled(language_key: str, code: str) -> tuple[Result, str]:
    """Compile and execute untrusted code. Returns (result, compile command)."""
    language = LANGUAGES.get(language_key)
    if language is None:
        return Result(False, "", f"Unknown language: {language_key}", 0.0, "compile"), ""

    if not code.strip():
        return Result(False, "", f"No {language.display} code to compile.",
                      0.0, "compile"), ""

    compiler = language.find_compiler()
    if not compiler:
        return Result(
            False, "",
            f"No {language.display} compiler found on PATH "
            f"(looked for: {', '.join(language.compilers)}).\n"
            f"Install one from {language.install_hint} to run generated code.",
            0.0, "compile",
        ), ""

    binary = "program.exe" if IS_WINDOWS else "program"
    source_name = f"main.{language.extension}"

    with tempfile.TemporaryDirectory(prefix=f"{language.key}run_") as tmp:
        with open(os.path.join(tmp, source_name), "w", encoding="utf-8") as f:
            f.write(code)

        cmd = build_compile_command(language, compiler, source_name, binary)
        pretty = " ".join([os.path.basename(cmd[0])] + cmd[1:])

        compiled = _run(cmd, tmp, COMPILE_TIMEOUT, "compile")
        if not compiled.ok:
            return compiled, pretty

        # Time execution only - trading compile time for runtime is the entire
        # premise of these flags, so charging the result for it measures wrong.
        return _run([os.path.join(tmp, binary)], tmp, RUN_TIMEOUT, "run"), pretty


def available_languages() -> list[str]:
    """Language keys with a working compiler on this machine."""
    return [key for key, lang in LANGUAGES.items() if lang.find_compiler()]


_probe_cache: dict[str, bool] = {}


def cpp_has_int128() -> bool:
    """Whether this C++ toolchain supports the __int128 extension.

    Worth probing rather than assuming. It is a GCC/Clang extension present on
    64-bit Linux (so the deployed container has it) but missing from MinGW on
    Windows. Models reach for it unprompted when a benchmark needs wide
    accumulators, and telling them "avoid non-portable extensions" is too weak
    to stop it - stating the capability as a fact is not.
    """
    if "int128" in _probe_cache:
        return _probe_cache["int128"]

    lang = LANGUAGES["cpp"]
    compiler = lang.find_compiler()
    if not compiler:
        _probe_cache["int128"] = False
        return False

    snippet = "unsigned __int128 x = 1; int main(){ return (int)(x & 1) - 1; }"
    with tempfile.TemporaryDirectory(prefix="probe_") as tmp:
        path = os.path.join(tmp, "probe.cpp")
        with open(path, "w", encoding="utf-8") as f:
            f.write(snippet)
        out = os.path.join(tmp, "probe.exe" if IS_WINDOWS else "probe")
        try:
            proc = subprocess.run(
                [compiler, "-std=c++17", "-fsyntax-only", path],
                capture_output=True, text=True, timeout=25, env=_safe_env(),
            )
            ok = proc.returncode == 0
        except Exception:
            ok = False

    _probe_cache["int128"] = ok
    return ok


def _version_of(path: str) -> str:
    try:
        out = subprocess.run([path, "--version"], capture_output=True,
                             text=True, timeout=5)
        return (out.stdout or out.stderr).strip().splitlines()[0]
    except Exception:
        return os.path.basename(path)


def toolchain_status() -> dict[str, str]:
    """Report what the sandbox can actually run, for the UI to display."""
    status: dict[str, str] = {}
    for key, lang in LANGUAGES.items():
        found = lang.find_compiler()
        status[key] = _version_of(found) if found else ""
    status["python"] = f"Python {platform.python_version()}"
    status["isolation"] = (
        "full (rlimits + scrubbed env)" if not IS_WINDOWS
        else "partial (scrubbed env, no rlimits on Windows)"
    )
    return status