"""Execute Python code in a sandboxed environment.

Supports two sandbox modes (set via AGENT_STDLIB_CODE_SANDBOX_MODE):
- "local": Process isolation with resource limits (default, fast)
- "docker": Full container isolation (more secure, requires Docker)

Security features:
- Timeout enforcement
- Memory limits
- CPU limits
- No network (in docker mode)
- Isolated filesystem

Namespace persistence:
- When the calling context has a ``_code_namespace``, variables defined in
  one call persist to the next call within the same turn.
"""

import logging
import os
from typing import Any, Literal

from core.pathways.context import Context
from core.pathways.models.tool import ToolOutput

logger = logging.getLogger(__name__)


def _get_sandbox_mode() -> Literal["local", "docker"]:
    """Get sandbox mode from environment."""
    mode = os.getenv("AGENT_STDLIB_CODE_SANDBOX_MODE", "local").lower()
    if mode == "docker":
        return "docker"
    return "local"


def _get_docker_image(profile: str | None = None) -> str:
    """Get Docker image for a profile."""
    if profile == "datascience":
        return os.getenv(
            "ALBUS_CODE_SANDBOX_DOCKER_IMAGE_DATASCIENCE",
            "albus-sandbox-datascience:latest",
        )
    if profile == "viz":
        return os.getenv("ALBUS_CODE_SANDBOX_DOCKER_IMAGE_VIZ", "albus-sandbox-viz:latest")
    if profile == "web":
        return os.getenv("ALBUS_CODE_SANDBOX_DOCKER_IMAGE_WEB", "albus-sandbox-web:latest")
    return os.getenv("AGENT_STDLIB_CODE_SANDBOX_DOCKER_IMAGE", "python:3.11-slim")


def _load_namespace(context: Context | None) -> dict[str, Any]:
    """Load the persisted code namespace from the execution context."""
    if context is None:
        return {}
    ns = context.run.get_code_namespace()
    return dict(ns) if ns else {}


def _save_namespace(
    context: Context | None,
    returned_ns: dict[str, Any],
) -> None:
    """Merge the returned namespace back into the execution context."""
    if context is None or not returned_ns:
        return
    context.run.update_code_namespace(returned_ns)


async def run(
    code: str,
    input: Any = None,
    timeout: int = 60,
    memory_mb: int = 512,
    profile: str | None = None,
    allow_packages: bool = True,
    context: Context | None = None,
) -> ToolOutput:
    """Execute Python code in a sandboxed environment.

    Args:
        code: Python code to execute.
        input: Input data accessible as ``input`` variable in the code.
        timeout: Maximum execution time in seconds (default: 60).
        memory_mb: Memory limit in MB (default: 512).
        profile: Docker image profile (``"datascience"``, ``"viz"``,
            ``"web"``, or ``None`` for basic).
        allow_packages: Allow site-packages imports (default: True).
        context: Execution context (injected by framework).

    The code should either:
    - Define ``def main(input): ...`` which returns the result
    - Set ``result = ...`` as a global variable
    - Print output to stdout

    Variables defined in the code persist across subsequent
    ``code.execute`` calls within the same conversation turn (only
    JSON-serializable values survive).

    Returns
    -------
        {
            "success": True/False,
            "result": <value from main() or result variable>,
            "stdout": <captured stdout>,
            "stderr": <captured stderr>,
            "error": <error message if failed>,
            "traceback": <Python traceback if failed>,
            "sandbox": "local" or "docker",
            "duration_ms": <execution time>,
            "persisted_vars": <list of variable names that survived>
        }
    """
    if not code:
        return ToolOutput(success=False, error="code is required")

    # Load persisted namespace from prior code.execute calls in this turn.
    ns = _load_namespace(context)

    try:
        from infrastructure.sandbox.python_runner import run_python_sandboxed

        sandbox_mode = _get_sandbox_mode()
        docker_image = (
            _get_docker_image(profile) if sandbox_mode == "docker" else "python:3.11-slim"
        )

        result = run_python_sandboxed(
            code=code,
            input_obj=input,
            namespace=ns if ns else None,
            timeout_ms=timeout * 1000,
            memory_mb=memory_mb,
            cpu_seconds=timeout,
            max_output_kb=1024,
            allow_site_packages=allow_packages,
            sandbox=sandbox_mode,
            docker_image=docker_image,
        )

        # Persist namespace back into the context for subsequent calls.
        returned_ns = result.get("namespace") or {}
        _save_namespace(context, returned_ns)

        # Normalize response
        ok = bool(result.get("ok", False))

        data: dict[str, Any] = {
            "result": result.get("result"),
            "stdout": result.get("stdout", ""),
            "stderr": result.get("stderr", ""),
            "sandbox": result.get("sandbox", sandbox_mode),
            "duration_ms": result.get("duration_ms", 0),
        }

        # Include traceback for failed executions so the LLM can self-debug.
        if not ok and result.get("traceback"):
            data["traceback"] = result["traceback"]

        # Report which variables are now persisted.
        if returned_ns:
            data["persisted_vars"] = sorted(returned_ns.keys())

        return ToolOutput(
            success=ok,
            error=(None if ok else str(result.get("error") or "Execution failed")),
            data=data,
        )

    except ImportError:
        # Fallback to simple exec if sandbox not available
        import io
        import sys

        old_stdout = sys.stdout
        old_stderr = sys.stderr
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()

        try:
            sys.stdout = stdout_capture
            sys.stderr = stderr_capture

            exec_namespace: dict[str, Any] = {
                "__builtins__": __builtins__,
                "__name__": "__main__",
                "input": input,
            }
            # Inject persisted namespace.
            exec_namespace.update(ns)

            exec(code, exec_namespace)  # noqa: S102

            # Get result
            result_val = None
            if "main" in exec_namespace and callable(exec_namespace["main"]):
                result_val = exec_namespace["main"](input)
            elif "result" in exec_namespace:
                result_val = exec_namespace["result"]

            return ToolOutput(
                success=True,
                data={
                    "result": result_val,
                    "stdout": stdout_capture.getvalue(),
                    "stderr": stderr_capture.getvalue(),
                    "sandbox": "none",
                    "duration_ms": 0,
                },
            )

        except Exception as e:
            return ToolOutput(
                success=False,
                error=str(e),
                data={
                    "stdout": stdout_capture.getvalue(),
                    "stderr": stderr_capture.getvalue(),
                    "sandbox": "none",
                },
            )
        finally:
            sys.stdout = old_stdout
            sys.stderr = old_stderr

    except Exception as e:
        return ToolOutput(success=False, error=str(e))
