"""
Code execution engine for coding problems.

Supports local execution (python3 subprocess) and future Piston API execution.
The harness wraps user code with a JSON-emitting test runner so results can
be parsed deterministically regardless of what the user's code prints.

Design principles:
- Temp files are always cleaned up (finally block)
- Each test case is isolated: an exception in one case does not abort others
- Hidden test cases are tracked so the UI can suppress their expected values
- All public functions return plain dicts — no exceptions propagate to callers
"""

from __future__ import annotations

import base64
import json
import os
import shutil
import subprocess
import sys
import tempfile
from typing import Any

# resource is POSIX-only (not available on Windows).
try:
    import resource as _resource
    _HAS_RESOURCE = True
except ImportError:
    _HAS_RESOURCE = False


# ---------------------------------------------------------------------------
# Sandbox preexec (POSIX only)
# ---------------------------------------------------------------------------


def _sandbox_preexec() -> None:
    """Set resource limits on the child process before exec.

    Called by subprocess as preexec_fn — runs in the forked child before the
    executable is loaded.  Silently skips any limit that the OS does not
    support so the function is safe to call on all POSIX platforms.

    Limits applied:
    - RLIMIT_AS:    512 MB virtual address space (catches runaway allocations;
                    must exceed Python's baseline virtual mapping which can be
                    >256 MB on macOS due to shared libraries and heap arenas)
    - RLIMIT_CPU:   15 s CPU time (belt-and-suspenders alongside timeout=)
    - RLIMIT_NPROC: 50 child processes (prevents fork bombs)
    """
    if not _HAS_RESOURCE:
        return

    _512_MB = 512 * 1024 * 1024

    limits = [
        (_resource.RLIMIT_AS, (_512_MB, _512_MB)),
        (_resource.RLIMIT_CPU, (15, 15)),
    ]

    # RLIMIT_NPROC is not available on all POSIX systems (e.g. missing on some
    # BSDs and older Linux kernels), so guard individually.
    if hasattr(_resource, "RLIMIT_NPROC"):
        limits.append((_resource.RLIMIT_NPROC, (50, 50)))

    for resource_id, limit_pair in limits:
        try:
            _resource.setrlimit(resource_id, limit_pair)
        except (ValueError, _resource.error):
            # Silently skip unsupported or unprivileged limit changes.
            pass


# ---------------------------------------------------------------------------
# Runtime detection
# ---------------------------------------------------------------------------

# Map language slug → local executable name.
# For Java we check for the compiler (javac) to confirm the JDK is installed;
# the runtime (java) is implicitly available when the JDK is present.
_LOCAL_EXECUTABLES: dict[str, str] = {
    "python": "python3",
    "java": "javac",
}


def detect_runtime(language: str, piston_url: str | None = None) -> str:
    """Determine the best available execution runtime for *language*.

    Resolution order:
    1. Local executable — for Python: sys.executable is always present;
       for Java: javac must be on PATH (confirms JDK installed).
    2. Piston API (if *piston_url* is provided and reachable — checked lazily)
    3. 'unavailable'

    Returns one of: 'local', 'piston', 'unavailable'.
    """
    if language == "python":
        # sys.executable is the interpreter running this process — always valid.
        return "local"

    executable = _LOCAL_EXECUTABLES.get(language)
    if executable and shutil.which(executable):
        return "local"

    # Piston is accepted as available if a URL was supplied; actual connectivity
    # is validated only at execution time to keep detection fast.
    if piston_url:
        return "piston"

    return "unavailable"


# ---------------------------------------------------------------------------
# Python harness builder
# ---------------------------------------------------------------------------

# The harness template uses underscore-prefixed variable names throughout so
# they do not collide with names defined in user-submitted code.
#
# Test cases are embedded as a base64-encoded JSON blob.  This eliminates all
# quoting and escaping hazards: backslashes, embedded quotes, triple-quotes,
# and Unicode sequences in test inputs cannot break the harness regardless of
# their content.
#
# Output format (printed as JSON to stdout):
# {
#   "results": [
#     {"case": N, "passed": bool, "expected": str, "actual": str,
#      "time_ms": float, "hidden": bool}
#   ],
#   "passed_count": N,
#   "total_count": N,
#   "all_passed": bool
# }

# The {tc_b64} placeholder is a pure ASCII base64 string — safe to embed in
# any string literal context without any escaping.
_HARNESS_TEMPLATE = """\
import json as _json, time as _time, base64 as _b64

# ---- user code ----
{user_code}
# ---- end user code ----

_test_cases = _json.loads(_b64.b64decode("{tc_b64}").decode())
_results = []

for _i, _tc in enumerate(_test_cases):
    _expected = _tc.get("expected_output") or _tc["expected"]
    _hidden = bool(_tc.get("hidden", False))
    _t0 = _time.perf_counter()
    try:
        _actual = {function_name}(**_tc["input"])
        _passed = _actual == _expected
    except Exception as _exc:
        _actual = f"Exception: {{_exc}}"
        _passed = False
    _elapsed_ms = (_time.perf_counter() - _t0) * 1000
    _results.append({{
        "case": _i,
        "passed": _passed,
        "expected": repr(_expected),
        "actual": repr(_actual),
        "time_ms": round(_elapsed_ms, 3),
        "hidden": _hidden,
    }})

_passed_count = sum(1 for _r in _results if _r["passed"])
_total_count = len(_results)
print(_json.dumps({{
    "results": _results,
    "passed_count": _passed_count,
    "total_count": _total_count,
    "all_passed": _passed_count == _total_count,
}}))
"""


def build_python_harness(
    user_code: str,
    function_name: str,
    test_cases: list[dict[str, Any]],
) -> str:
    """Build a complete Python script that runs *user_code* against *test_cases*.

    The generated script prints a single JSON object to stdout with per-case
    results and aggregate pass/fail counts.  Each test case is run inside a
    try/except so a crashing solution produces per-case error entries rather
    than aborting the whole run.

    Test cases are embedded as a base64-encoded JSON blob so that special
    characters in inputs (quotes, backslashes, Unicode) cannot cause syntax
    errors or injection in the generated script.

    Comparison is done on native Python objects (not string representations),
    so list, int, bool, and None expected values all work correctly without
    any type coercion.

    Args:
        user_code:     The user's submitted Python source (may define helper
                       functions/classes above the main function).
        function_name: The name of the callable to invoke for each test case.
        test_cases:    List of dicts with keys: input (kwargs dict),
                       expected or expected_output (native Python value), hidden (bool).

    Returns:
        A complete Python script as a string suitable for exec or subprocess.
    """
    tc_b64 = base64.b64encode(json.dumps(test_cases).encode()).decode()

    return _HARNESS_TEMPLATE.format(
        user_code=user_code,
        function_name=function_name,
        tc_b64=tc_b64,
    )


# ---------------------------------------------------------------------------
# Local execution
# ---------------------------------------------------------------------------

def run_code_local(language: str, source: str, timeout: int = 10) -> dict:
    """Execute *source* locally via subprocess and return parsed results.

    Writes *source* to a temp file with the appropriate extension, runs it,
    and parses the JSON output.  The temp file is always deleted even if
    execution fails or times out.

    Args:
        language: Language slug ('python', 'java', …).  Only 'python' is
                  fully implemented; other languages return an error dict.
        source:   Complete source code to execute (for Python this is the
                  harness produced by build_python_harness).
        timeout:  Maximum wall-clock seconds before the process is killed.

    Returns:
        A result dict from parse_execution_result.
    """
    ext_map = {"python": ".py", "java": ".java"}
    ext = ext_map.get(language, ".txt")

    if language == "python":
        # Use the same interpreter running this process — avoids "python3 not
        # found" errors in venvs and conda environments where only the venv
        # interpreter is on PATH.
        executable = sys.executable
    elif language == "java":
        executable = "java"
    else:
        return parse_execution_result("", f"Unsupported language: {language}", 1)

    # Minimal environment for the child process: strip virtualenv paths,
    # credentials, and anything that could leak configuration to arbitrary code.
    clean_env = {
        k: os.environ[k]
        for k in ("PATH", "HOME", "LANG")
        if k in os.environ
    }

    tmp_path: str | None = None
    try:
        fd, tmp_path = tempfile.mkstemp(suffix=ext)
        try:
            os.write(fd, source.encode())
        finally:
            os.close(fd)

        subprocess_kwargs: dict = dict(
            capture_output=True,
            text=True,
            timeout=timeout,
            env=clean_env,
        )
        if _HAS_RESOURCE:
            subprocess_kwargs["preexec_fn"] = _sandbox_preexec

        proc = subprocess.run([executable, tmp_path], **subprocess_kwargs)
        return parse_execution_result(proc.stdout, proc.stderr, proc.returncode)

    except subprocess.TimeoutExpired:
        return parse_execution_result("", "", 0, timed_out=True)

    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


# ---------------------------------------------------------------------------
# Result parsing
# ---------------------------------------------------------------------------

def parse_execution_result(
    stdout: str,
    stderr: str,
    return_code: int,
    *,
    timed_out: bool = False,
) -> dict:
    """Parse raw subprocess output into a structured result dict.

    Handles four scenarios:
    1. Timeout: returns an error dict with "timed out" message.
    2. Valid JSON in stdout: returns the parsed dict.
    3. No JSON but stderr present: returns error dict with stderr message.
    4. No JSON and no stderr: returns generic error dict.

    Args:
        stdout:      The process's standard output string.
        stderr:      The process's standard error string.
        return_code: The process exit code.
        timed_out:   True when subprocess.TimeoutExpired was caught.

    Returns:
        A dict with at minimum: all_passed (bool), error (str or None).
        On success also includes: results, passed_count, total_count.
    """
    if timed_out:
        return {
            "all_passed": False,
            "passed_count": 0,
            "total_count": 0,
            "results": [],
            "error": "execution timed out",
        }

    # Try to extract JSON from stdout.  The harness always prints exactly one
    # JSON object on the last non-empty line; fall back to a full parse attempt.
    if stdout.strip():
        try:
            data = json.loads(stdout.strip())
            # Successful parse — enrich with a None error sentinel so callers
            # can always do result.get("error") without a KeyError.
            data.setdefault("error", None)
            return data
        except json.JSONDecodeError:
            pass

    # No valid JSON — execution failed before harness output (syntax/compile error).
    error_msg = stderr.strip() if stderr.strip() else "execution produced no output"
    return {
        "all_passed": False,
        "passed_count": 0,
        "total_count": 0,
        "results": [],
        "error": error_msg,
    }
