from __future__ import annotations

import contextlib
import json
import os
import resource
import shutil
import signal
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from typing import Any, Literal


def _truthy_env(name: str) -> bool:
    raw = (os.getenv(name) or "").strip().lower()
    return raw in ("1", "true", "yes", "y", "on")


def _clamp_int(v: Any, *, default: int, lo: int, hi: int) -> int:
    try:
        x = int(v)
    except Exception:
        return default
    return max(lo, min(hi, x))


_RUNNER_SOURCE = r'''
from __future__ import annotations

import json
import os
import sys
import traceback
import types
from typing import Any

# Keys that must never leak into the persisted namespace.
_NAMESPACE_EXCLUDE = frozenset({
    "__name__", "__builtins__", "__doc__", "__loader__", "__spec__",
    "__package__", "__cached__", "__file__", "input",
})

# Maximum serialized namespace size (bytes) — prevents runaway memory.
_MAX_NAMESPACE_BYTES = 1_048_576  # 1 MB


def _jsonable(x: Any) -> Any:
    try:
        json.dumps(x)
        return x
    except Exception:
        return {"__repr__": repr(x)}


def _is_jsonable(x: Any) -> bool:
    """Return True if *x* can be round-tripped through JSON."""
    try:
        json.dumps(x)
        return True
    except Exception:
        return False


def _capture_namespace(g: dict[str, Any]) -> dict[str, Any]:
    """Extract user-defined, JSON-serializable globals for persistence."""
    ns: dict[str, Any] = {}
    for k, v in g.items():
        if k.startswith("_") or k in _NAMESPACE_EXCLUDE:
            continue
        # Skip modules, functions, classes, and other non-data objects.
        if isinstance(v, (types.ModuleType, types.FunctionType, type)):
            continue
        if _is_jsonable(v):
            ns[k] = v
    # Enforce size cap — silently drop if too large.
    try:
        blob = json.dumps(ns, ensure_ascii=False)
        if len(blob.encode("utf-8")) > _MAX_NAMESPACE_BYTES:
            return {}
    except Exception:
        return {}
    return ns


def _load_stdin_json() -> dict[str, Any]:
    raw = sys.stdin.read()
    if not raw.strip():
        return {}
    try:
        obj = json.loads(raw)
        return obj if isinstance(obj, dict) else {}
    except Exception:
        return {}


def main() -> int:
    payload = _load_stdin_json()
    code = payload.get("code")
    inp = payload.get("input")
    ns_in = payload.get("namespace") or {}

    out: dict[str, Any] = {"ok": False}
    try:
        if not isinstance(code, str) or not code.strip():
            raise ValueError("missing_code")

        g: dict[str, Any] = {"__name__": "__main__", "input": inp}
        # Inject persisted namespace from prior code.execute calls.
        if isinstance(ns_in, dict):
            g.update(ns_in)
        exec(compile(code, "<python.run>", "exec"), g, g)  # noqa: S102 - intentional tool boundary

        res: Any = None
        fn = g.get("main")
        if callable(fn):
            res = fn(inp)
        elif "result" in g:
            res = g.get("result")
        else:
            res = None

        out = {
            "ok": True,
            "result": _jsonable(res),
            "namespace": _capture_namespace(g),
        }
    except Exception as e:  # noqa: BLE001 - tool boundary
        out = {
            "ok": False,
            "error": str(e),
            "traceback": traceback.format_exc(),
            "namespace": _capture_namespace(g) if "g" in dir() else {},
        }

    try:
        with open("__albus_result.json", "w", encoding="utf-8") as f:
            json.dump(out, f, ensure_ascii=False)
    except Exception:
        # If we can't write results, there's nothing else we can do.
        pass
    return 0 if out.get("ok") is True else 1


if __name__ == "__main__":
    raise SystemExit(main())
'''


def _docker_available() -> tuple[bool, str]:
    """Return (available, note)."""
    if shutil.which("docker") is None:
        return False, "docker_not_found"
    try:
        # Quick, low-noise check that the daemon is reachable.
        p = subprocess.run(
            ["docker", "version", "--format", "{{.Server.Version}}"],
            capture_output=True,
            timeout=2.5,
            check=False,
        )
        if p.returncode != 0:
            err = (p.stderr or b"").decode("utf-8", errors="replace").strip()
            return False, f"docker_unavailable:{err[:200]}"
        return True, "ok"
    except Exception as e:  # noqa: BLE001 - sandbox boundary
        return False, f"docker_unavailable:{e}"


def _docker_image_present(image: str) -> bool:
    try:
        p = subprocess.run(
            ["docker", "image", "inspect", image],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            timeout=2.5,
            check=False,
        )
        return p.returncode == 0
    except Exception:
        return False


def _run_python_sandboxed_docker(
    *,
    code: str,
    input_obj: Any = None,
    namespace: dict[str, Any] | None = None,
    timeout_ms: int = 10_000,
    memory_mb: int = 256,
    cpu_seconds: int = 5,
    max_output_kb: int = 256,
    allow_site_packages: bool = False,
    docker_image: str = "python:3.11-slim",
) -> dict[str, Any]:
    """Execute code inside a hardened Docker sandbox.

    Security properties (best-effort, depends on host Docker daemon):
    - No network: --network=none
    - No host filesystem access: only mounts a temp workdir to /work
    - Read-only root filesystem: --read-only (workdir remains writable)
    - Drops Linux capabilities: --cap-drop=ALL + no-new-privileges
    """
    t0 = time.time()
    ok, note = _docker_available()
    if not ok:
        return {
            "ok": False,
            "error": f"docker_required:{note}",
            "stdout": "",
            "stderr": "",
            "duration_ms": (time.time() - t0) * 1000.0,
        }

    # Do not attempt to pull images automatically; in secure deployments, images
    # should be pre-pulled / pinned by ops.
    if not _docker_image_present(docker_image):
        return {
            "ok": False,
            "error": f"docker_image_missing:{docker_image}",
            "stdout": "",
            "stderr": "",
            "duration_ms": (time.time() - t0) * 1000.0,
        }

    timeout_ms_i = _clamp_int(timeout_ms, default=10_000, lo=1, hi=600_000)
    memory_mb_i = _clamp_int(memory_mb, default=256, lo=16, hi=16_384)
    cpu_seconds_i = _clamp_int(cpu_seconds, default=5, lo=1, hi=600)
    max_output_kb_i = _clamp_int(max_output_kb, default=256, lo=1, hi=8192)

    with tempfile.TemporaryDirectory(prefix="albus_py_docker_") as td:
        # Ensure an unprivileged container user can write results.
        with contextlib.suppress(Exception):
            os.chmod(td, 0o777)

        wd = Path(td)
        runner_path = wd / "_runner.py"
        runner_path.write_text(_RUNNER_SOURCE, encoding="utf-8")

        payload: dict[str, Any] = {"code": code, "input": input_obj}
        if namespace:
            payload["namespace"] = namespace
        stdin = (json.dumps(payload, ensure_ascii=False) + "\n").encode("utf-8")

        stdout_limit = max_output_kb_i * 1024
        stderr_limit = max_output_kb_i * 1024

        # Hardened container flags:
        # - No network
        # - Read-only root
        # - Drop caps + no-new-privileges
        # - Limit pids, memory, cpu
        # - Limit file sizes and open files via ulimit where supported
        docker_cmd = [
            "docker",
            "run",
            "--rm",
            # Keep stdin open so we can pass the JSON payload via communicate().
            # Without this, sys.stdin.read() inside the container is empty and user code never runs.
            "-i",
            "--network=none",
            "--read-only",
            "--security-opt",
            "no-new-privileges",
            "--cap-drop=ALL",
            "--pids-limit=64",
            "--memory",
            f"{int(memory_mb_i)}m",
            "--cpus",
            "1",
            "--ulimit",
            "nofile=64:64",
            "--ulimit",
            "fsize=5242880:5242880",
            "--user",
            "65534:65534",
            "-v",
            f"{wd!s}:/work:rw",
            "-w",
            "/work",
            docker_image,
            "python",
            "-u",
        ]

        if not allow_site_packages:
            docker_cmd += ["-I", "-S"]

        docker_cmd += ["/work/_runner.py"]

        proc = subprocess.Popen(
            docker_cmd,
            cwd=str(wd),
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        try:
            out_b, err_b = proc.communicate(stdin, timeout=timeout_ms_i / 1000.0)
        except subprocess.TimeoutExpired:
            with contextlib.suppress(Exception):
                proc.kill()
            return {
                "ok": False,
                "error": "timeout",
                "stdout": "",
                "stderr": "",
                "duration_ms": (time.time() - t0) * 1000.0,
            }

        stdout = (out_b or b"")[:stdout_limit].decode("utf-8", errors="replace")
        stderr = (err_b or b"")[:stderr_limit].decode("utf-8", errors="replace")

        result_path = wd / "__albus_result.json"
        result_obj: dict[str, Any] | None = None
        if result_path.exists():
            try:
                result_obj = json.loads(result_path.read_text(encoding="utf-8"))
            except Exception:
                result_obj = None

        if not isinstance(result_obj, dict):
            return {
                "ok": False,
                "error": f"docker_runner_failed:exit_code={proc.returncode}",
                "stdout": stdout,
                "stderr": stderr,
                "duration_ms": (time.time() - t0) * 1000.0,
            }

        result_obj["stdout"] = stdout
        result_obj["stderr"] = stderr
        result_obj["duration_ms"] = (time.time() - t0) * 1000.0
        # Helpful for debugging but not required.
        result_obj.setdefault("sandbox", "docker")
        result_obj.setdefault("cpu_seconds", cpu_seconds_i)
        return result_obj


def run_python_sandboxed(
    *,
    code: str,
    input_obj: Any = None,
    namespace: dict[str, Any] | None = None,
    timeout_ms: int = 10_000,
    memory_mb: int = 256,
    cpu_seconds: int = 5,
    max_output_kb: int = 256,
    allow_site_packages: bool = False,
    python_executable: str | None = None,
    sandbox: Literal["local", "docker"] = "local",
    docker_image: str = "python:3.11-slim",
    extra_env: dict[str, str] | None = None,
) -> dict[str, Any]:
    """Execute Python code in a separate process with best-effort local sandboxing.

    Contract:
    - Input: JSON object with ``code`` (required), ``input`` (optional),
      and ``namespace`` (optional dict of persisted variables).
    - User code may define:
        - ``def main(input): ...`` (preferred), or
        - set a global ``result = ...``
    - Output: JSON object::

        {"ok": bool, "result": any?, "namespace": dict?,
         "stdout": str, "stderr": str, "error": str?}

    Args:
        namespace: Dict of JSON-serializable variables to inject into the
            execution globals before running the code.  After execution
            the runner captures user-defined globals and returns them in
            ``result["namespace"]`` so they can be persisted across calls.
        extra_env: Additional environment variables to pass to the subprocess.
                   Used for builder mode to pass ALBUS_SRC_PATH.
    """
    if sandbox == "docker":
        return _run_python_sandboxed_docker(
            code=code,
            input_obj=input_obj,
            namespace=namespace,
            timeout_ms=timeout_ms,
            memory_mb=memory_mb,
            cpu_seconds=cpu_seconds,
            max_output_kb=max_output_kb,
            allow_site_packages=allow_site_packages,
            docker_image=docker_image,
        )

    t0 = time.time()
    timeout_ms_i = _clamp_int(timeout_ms, default=10_000, lo=1, hi=600_000)
    memory_mb_i = _clamp_int(memory_mb, default=256, lo=16, hi=16_384)
    cpu_seconds_i = _clamp_int(cpu_seconds, default=5, lo=1, hi=600)
    max_output_kb_i = _clamp_int(max_output_kb, default=256, lo=1, hi=8192)

    py = str(python_executable).strip() if python_executable else sys.executable
    if not py:
        py = sys.executable

    with tempfile.TemporaryDirectory(prefix="albus_py_") as td:
        wd = Path(td)
        runner_path = wd / "_runner.py"
        runner_path.write_text(_RUNNER_SOURCE, encoding="utf-8")

        env = {
            # Minimal env: keep it deterministic.
            "PYTHONDONTWRITEBYTECODE": "1",
            "PYTHONIOENCODING": "utf-8",
            "PYTHONHASHSEED": "0",
            # Explicitly strip user-site imports even if user tries.
            "PYTHONNOUSERSITE": "1",
        }

        # Allow opt-in passthrough (debug only).
        if _truthy_env("PATHWAY_PYTHON_TOOL_PASSTHROUGH_ENV"):
            for k, v in os.environ.items():
                if isinstance(k, str) and isinstance(v, str):
                    env.setdefault(k, v)

        # Merge extra_env (for builder mode etc.)
        if extra_env:
            for k, v in extra_env.items():
                if isinstance(k, str) and isinstance(v, str):
                    env[k] = v

        args = [py, "-u"]
        if not allow_site_packages:
            # -I: isolate environment, -S: don't import site
            args += ["-I", "-S"]
        args.append(str(runner_path))

        payload: dict[str, Any] = {"code": code, "input": input_obj}
        if namespace:
            payload["namespace"] = namespace
        stdin = (json.dumps(payload, ensure_ascii=False) + "\n").encode("utf-8")

        stdout_limit = max_output_kb_i * 1024
        stderr_limit = max_output_kb_i * 1024

        def _limit_child() -> None:
            # Best-effort POSIX resource limits (macOS/Linux).
            with contextlib.suppress(Exception):
                os.setsid()
            try:
                # CPU time (seconds).
                resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds_i, cpu_seconds_i + 1))
                # Address space (bytes).
                mem_bytes = int(memory_mb_i) * 1024 * 1024
                resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes))
                # File size (bytes).
                resource.setrlimit(resource.RLIMIT_FSIZE, (5 * 1024 * 1024, 5 * 1024 * 1024))
                # Open files.
                resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
            except Exception:
                # Non-fatal; we still have wall-clock timeout and process isolation.
                pass

        proc = subprocess.Popen(
            args,
            cwd=str(wd),
            env=env,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            preexec_fn=_limit_child if os.name == "posix" else None,
        )

        try:
            out_b, err_b = proc.communicate(stdin, timeout=timeout_ms_i / 1000.0)
        except subprocess.TimeoutExpired:
            try:
                # Kill process group if possible.
                if os.name == "posix":
                    os.killpg(proc.pid, signal.SIGKILL)
                else:
                    proc.kill()
            except Exception:
                pass
            return {
                "ok": False,
                "error": "timeout",
                "stdout": "",
                "stderr": "",
                "duration_ms": (time.time() - t0) * 1000.0,
            }

        stdout = (out_b or b"")[:stdout_limit].decode("utf-8", errors="replace")
        stderr = (err_b or b"")[:stderr_limit].decode("utf-8", errors="replace")

        result_path = wd / "__albus_result.json"
        result_obj: dict[str, Any] | None = None
        if result_path.exists():
            try:
                result_obj = json.loads(result_path.read_text(encoding="utf-8"))
            except Exception:
                result_obj = None

        if not isinstance(result_obj, dict):
            # Fall back to exit-code based failure.
            return {
                "ok": False,
                "error": f"python_runner_failed:exit_code={proc.returncode}",
                "stdout": stdout,
                "stderr": stderr,
                "duration_ms": (time.time() - t0) * 1000.0,
            }

        # Attach captured streams + timing.
        result_obj["stdout"] = stdout
        result_obj["stderr"] = stderr
        result_obj["duration_ms"] = (time.time() - t0) * 1000.0
        return result_obj


__all__ = ["run_python_sandboxed"]
