"""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
from typing import Any, Literal

from core.config import get_settings
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."""
    return get_settings().agent_stdlib_code_sandbox_mode


_DOCKER_IMAGE_PROFILES: dict[str, str] = {
    "datascience": "albus_code_sandbox_docker_image_datascience",
    "viz": "albus_code_sandbox_docker_image_viz",
    "web": "albus_code_sandbox_docker_image_web",
}


def _get_docker_image(profile: str | None = None) -> str:
    """Get Docker image for a profile."""
    settings = get_settings()
    if profile and profile in _DOCKER_IMAGE_PROFILES:
        return getattr(settings, _DOCKER_IMAGE_PROFILES[profile])
    return settings.agent_stdlib_code_sandbox_docker_image


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 or 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:
        return ToolOutput(
            success=False,
            error="Sandbox unavailable. Install the sandbox infrastructure or set "
            "AGENT_STDLIB_CODE_SANDBOX_MODE to configure a runner.",
        )

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