"""Isolated Python execution with hard limits.

Isolation layers (best available on the host):
  * macOS: `sandbox-exec` seatbelt profile — no network, writes only inside the run dir.
  * Linux (container): expected to run inside a hardened pod (NetworkPolicy deny-all,
    read-only rootfs); we still apply rlimits.
  * Always: `python -I -B`, rlimits (CPU, RAM, files, processes), wall-clock timeout,
    output truncation, auto-save of open matplotlib figures.
"""

from __future__ import annotations

import base64
import os
import platform
import resource
import shutil
import subprocess
import sys
import tempfile
import time
import uuid
from pathlib import Path

MAX_OUTPUT = 50 * 1024
MAX_FILES = 10
MAX_FILES_BYTES = 20 * 1024 * 1024
MEM_LIMIT = 512 * 1024 * 1024
PYTHON = os.environ.get("SANDBOX_PYTHON", sys.executable)

PRELUDE = '''
import os as _os, sys as _sys
_os.chdir(_os.environ["SANDBOX_WORKDIR"])
_sys.path.insert(0, _os.getcwd())
try:
    import matplotlib as _mpl
    _mpl.use("Agg")
    import matplotlib.pyplot as _plt
    _plt.rcParams["figure.dpi"] = 130
    _plt.rcParams["axes.grid"] = True
    _plt.rcParams["grid.alpha"] = 0.3
    _plt.rcParams["axes.spines.top"] = False
    _plt.rcParams["axes.spines.right"] = False
except Exception:
    _plt = None
import atexit as _atexit
def _save_figs():
    if _plt is None:
        return
    for i, num in enumerate(_plt.get_fignums(), 1):
        try:
            _plt.figure(num).savefig(f"outputs/figure_{i}.png", bbox_inches="tight")
        except Exception as e:
            print("figure save failed:", e, file=_sys.stderr)
_atexit.register(_save_figs)
'''


def _seatbelt_profile(workdir: str) -> str:
    return f"""(version 1)
(deny default)
(allow process-exec)
(allow process-fork)
(allow sysctl-read)
(allow mach-lookup)
(allow file-read*)
(deny file-read* (subpath "/Users") (with no-log))
(allow file-read* (subpath "{workdir}"))
(allow file-read* (subpath "{os.path.dirname(os.path.realpath(PYTHON))}"))
(allow file-read* (subpath "{sys.prefix}"))
(allow file-read* (subpath "{sys.base_prefix}"))
(allow file-read* (subpath "{os.path.expanduser('~')}/.local/share/uv"))
(allow file-read* (subpath "/opt/homebrew"))
(allow file-read* (subpath "/private/tmp"))
(allow file-write* (subpath "{workdir}"))
(allow file-write* (subpath "/private/var/folders"))
(allow file-write* (subpath "/private/tmp"))
(deny network*)
"""


def _limits() -> None:
    # Runs in the child before exec.
    try:
        resource.setrlimit(resource.RLIMIT_CPU, (25, 30))
        resource.setrlimit(resource.RLIMIT_FSIZE, (MAX_FILES_BYTES, MAX_FILES_BYTES))
        resource.setrlimit(resource.RLIMIT_NOFILE, (256, 256))
        if platform.system() != "Darwin":  # RLIMIT_AS breaks numpy on macOS
            resource.setrlimit(resource.RLIMIT_AS, (MEM_LIMIT, MEM_LIMIT))
        resource.setrlimit(resource.RLIMIT_NPROC, (32, 32))
    except (ValueError, OSError):
        pass


def run(code: str, files_in: list[dict], timeout_s: int = 30) -> dict:
    run_id = uuid.uuid4().hex
    base = Path(tempfile.gettempdir()) / "uqo-sandbox" / run_id
    (base / "inputs").mkdir(parents=True)
    (base / "outputs").mkdir()
    for f in files_in[:MAX_FILES]:
        name = os.path.basename(str(f.get("name", "input")))
        data = base64.b64decode(f.get("content_b64", "") or "")
        (base / "inputs" / name).write_bytes(data[:MAX_FILES_BYTES])
    main = base / "main.py"
    main.write_text(PRELUDE + "\n" + code, encoding="utf-8")

    env = {
        "PATH": "/usr/bin:/bin",
        "HOME": str(base),
        "TMPDIR": str(base),
        "SANDBOX_WORKDIR": str(base),
        "MPLCONFIGDIR": str(base / ".mpl"),
        "PYTHONHASHSEED": "0",
        "LANG": "fr_CA.UTF-8",
        "LC_ALL": "fr_CA.UTF-8",
        "OMP_NUM_THREADS": "2",
        "OPENBLAS_NUM_THREADS": "2",
    }
    cmd = [PYTHON, "-I", "-B", str(main)]
    if platform.system() == "Darwin" and shutil.which("sandbox-exec"):
        profile = base / "profile.sb"
        profile.write_text(_seatbelt_profile(str(base)))
        cmd = ["/usr/bin/sandbox-exec", "-f", str(profile)] + cmd

    t0 = time.perf_counter()
    truncated = False
    try:
        proc = subprocess.run(
            cmd, cwd=base, env=env, capture_output=True, timeout=timeout_s,
            preexec_fn=_limits, check=False,
        )
        stdout, stderr, code_ = proc.stdout, proc.stderr, proc.returncode
    except subprocess.TimeoutExpired as exc:
        stdout = exc.stdout or b""
        stderr = (exc.stderr or b"") + f"\nTemps d'exécution dépassé ({timeout_s} s).".encode()
        code_ = 124
    duration = int((time.perf_counter() - t0) * 1000)

    def clip(b: bytes) -> str:
        nonlocal truncated
        s = b.decode("utf-8", errors="replace")
        if len(s) > MAX_OUTPUT:
            truncated = True
            return s[:MAX_OUTPUT] + "\n… (tronqué)"
        return s

    files_out: list[dict] = []
    total = 0
    candidates = sorted(list((base / "outputs").iterdir()) +
                        [p for p in base.iterdir() if p.is_file() and p.name not in
                         {"main.py", "profile.sb"}])
    for p in candidates:
        if not p.is_file() or p.name.startswith("."):
            continue
        size = p.stat().st_size
        if size == 0 or total + size > MAX_FILES_BYTES or len(files_out) >= MAX_FILES:
            continue
        total += size
        files_out.append({"name": p.name, "size": size,
                          "content_b64": base64.b64encode(p.read_bytes()).decode()})
    shutil.rmtree(base, ignore_errors=True)
    return {"stdout": clip(stdout), "stderr": clip(stderr), "exit_code": code_,
            "duration_ms": duration, "files_out": files_out, "truncated": truncated}