"""Security sandbox for user-code execution and file deserialization.

Two layers of defence for ``exec()``-based user code:

1. **AST validation** (``validate_user_code``) — parses the code string
   and walks the tree *before* execution, rejecting dangerous patterns:
   dunder attribute access (``__class__``, ``__subclasses__``), reflection
   helpers (``getattr``, ``type``, ``vars``), import statements, class
   definitions, and scope-escaping keywords (``global``, ``nonlocal``).
   This closes known CPython sandbox-escape vectors at the structural
   level.

2. **Restricted builtins** (``safe_globals``) — runtime defence-in-depth
   that removes ``__import__``, ``open``, ``eval``, ``exec``, ``compile``,
   ``breakpoint``, ``globals``, ``locals``, and ``input`` from the
   namespace passed to ``exec()``.

Also provides:
- ``safe_unpickle(path)`` — a ``RestrictedUnpickler`` that only allows
  known-safe classes (numpy, sklearn, catboost, etc.).
- ``validate_project_path(path)`` — ensures a path resolves inside the
  project root directory, preventing directory-traversal attacks.
"""

from __future__ import annotations

import ast
import builtins
import pickle
from pathlib import Path
from typing import Any

from haute._logging import get_logger
from haute._types import HauteError

logger = get_logger(component="sandbox")

# ---------------------------------------------------------------------------
# Project-root path validation
# ---------------------------------------------------------------------------

_PROJECT_ROOT: Path | None = None


def _get_project_root() -> Path:
    """Return the cached project root (cwd at import time)."""
    global _PROJECT_ROOT
    if _PROJECT_ROOT is None:
        _PROJECT_ROOT = Path.cwd().resolve()
    return _PROJECT_ROOT


def set_project_root(root: Path) -> None:
    """Override the project root (used by tests and CLI)."""
    global _PROJECT_ROOT
    _PROJECT_ROOT = root.resolve()


def validate_project_path(path: str | Path) -> Path:
    """Resolve *path* and verify it is inside the project root.

    Raises:
        ValueError: If the path escapes the project directory.
    """
    resolved = Path(path).resolve()
    root = _get_project_root()
    if not resolved.is_relative_to(root):
        raise ValueError(
            f"Path '{path}' resolves to '{resolved}' which is outside "
            f"the project root '{root}'"
        )
    return resolved


# ---------------------------------------------------------------------------
# Restricted builtins for exec()
# ---------------------------------------------------------------------------

# Builtins that allow arbitrary code execution or system access.
_BLOCKED_BUILTINS = frozenset({
    "__import__",
    "breakpoint",
    "compile",
    "eval",
    "exec",
    "globals",
    "locals",
    "open",
    "input",
    "memoryview",
})

_SAFE_BUILTINS: dict[str, Any] = {
    name: getattr(builtins, name)
    for name in dir(builtins)
    if not name.startswith("_") and name not in _BLOCKED_BUILTINS
}
# Keep __builtins__ itself pointing to the restricted set so nested
# lookups (e.g. list comprehensions) work correctly.
_SAFE_BUILTINS["__builtins__"] = _SAFE_BUILTINS


def safe_globals(*, allow_imports: bool = False, **extra: Any) -> dict[str, Any]:
    """Build a restricted global namespace for ``exec()``.

    Includes safe builtins + any extra bindings (e.g. ``pl=polars``).
    Blocks ``__import__``, ``open``, ``eval``, ``exec``, ``compile``,
    ``breakpoint``, ``globals``, ``locals``, and ``input``.

    *allow_imports* restores ``__import__`` — used for preamble code
    that legitimately imports from project utilities.
    """
    ns = dict(_SAFE_BUILTINS)
    if allow_imports:
        ns["__import__"] = builtins.__import__
        ns["__builtins__"] = {**_SAFE_BUILTINS, "__import__": builtins.__import__}
    ns.update(extra)
    return ns


# ---------------------------------------------------------------------------
# AST-level code validation — runs BEFORE exec()
# ---------------------------------------------------------------------------

# Attribute names that enable sandbox escapes via the Python type system.
_BLOCKED_ATTRS = frozenset({
    "__subclasses__",
    "__bases__",
    "__mro__",
    "__class__",
    "__globals__",
    "__code__",
    "__func__",
    "__self__",
    "__module__",
    "__dict__",
    "__init_subclass__",
    "__set_name__",
    "__reduce__",
    "__reduce_ex__",
    "__getattr__",
    "__setattr__",
    "__delattr__",
    "__import__",
    "__builtins__",
    "__loader__",
    "__spec__",
})

# Built-in function names that can be used to bypass attribute restrictions.
_BLOCKED_CALLS = frozenset({
    "getattr",
    "setattr",
    "delattr",
    "type",
    "vars",
    "dir",
    "hasattr",
    "classmethod",
    "staticmethod",
    "super",
    "__import__",
    "eval",
    "exec",
    "compile",
    "open",
    "breakpoint",
    "globals",
    "locals",
    "input",
    "exit",
    "quit",
    "help",
})


class UnsafeCodeError(HauteError):
    """Raised when AST validation detects a dangerous pattern."""


class _ASTValidator(ast.NodeVisitor):
    """Walk an AST and raise ``UnsafeCodeError`` on dangerous patterns.

    Blocks:
    - Dunder attribute access (``obj.__class__``, ``obj.__subclasses__()``)
    - Calls to reflection helpers (``getattr``, ``type``, ``vars``, etc.)
    - Import statements (unless ``allow_imports=True``)
    - Class, async, and lambda definitions
    - Star expressions in assignments (``a, *b = ...``)
    """

    def __init__(self, *, allow_imports: bool = False) -> None:
        super().__init__()
        self.allow_imports = allow_imports

    def visit_Attribute(self, node: ast.Attribute) -> None:
        if node.attr.startswith("__") and node.attr.endswith("__"):
            if node.attr in _BLOCKED_ATTRS:
                raise UnsafeCodeError(
                    f"Access to '{node.attr}' is blocked in pipeline code"
                )
        self.generic_visit(node)

    def visit_Call(self, node: ast.Call) -> None:
        # Block calls to dangerous built-in names
        if isinstance(node.func, ast.Name) and node.func.id in _BLOCKED_CALLS:
            raise UnsafeCodeError(
                f"Call to '{node.func.id}()' is blocked in pipeline code"
            )
        self.generic_visit(node)

    def visit_Import(self, node: ast.Import) -> None:
        if not self.allow_imports:
            raise UnsafeCodeError("import statements are blocked in pipeline code")

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        if not self.allow_imports:
            raise UnsafeCodeError("import statements are blocked in pipeline code")

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        raise UnsafeCodeError("class definitions are blocked in pipeline code")

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
        raise UnsafeCodeError(
            "async function definitions are blocked in pipeline code"
        )

    def visit_Global(self, node: ast.Global) -> None:
        raise UnsafeCodeError("global statements are blocked in pipeline code")

    def visit_Nonlocal(self, node: ast.Nonlocal) -> None:
        raise UnsafeCodeError(
            "nonlocal statements are blocked in pipeline code"
        )


_validator = _ASTValidator()
_preamble_validator = _ASTValidator(allow_imports=True)


def validate_user_code(code: str, *, allow_imports: bool = False) -> None:
    """Parse *code* and check for dangerous AST patterns.

    Raises ``UnsafeCodeError`` if the code contains blocked constructs
    (dunder access, imports, getattr, class defs, etc.).

    Called by ``_exec_user_code`` before ``exec()`` so dangerous code
    is rejected at the structural level — not just at runtime via
    restricted builtins.

    *allow_imports* permits ``import`` / ``from … import`` statements,
    used for preamble code which legitimately imports from utility modules.

    Results for safe code are cached by code string so repeated
    executions of the same node (preview, trace) skip the AST parse.
    """
    _validate_user_code_cached(code, allow_imports=allow_imports)


def _validate_user_code_cached(
    code: str,
    *,
    allow_imports: bool = False,
    _cache: dict[tuple[str, bool], bool] = {},  # noqa: B006
) -> None:
    """Inner validation with per-code-string caching.

    Uses a mutable default dict as a simple cache.  Safe-code results
    (``True``) are cached; unsafe code always raises before caching.
    """
    cache_key = (code, allow_imports)
    if cache_key in _cache:
        return

    try:
        tree = ast.parse(code)
    except SyntaxError:
        # Let exec() produce the proper SyntaxError with line numbers
        _cache[cache_key] = True
        return
    v = _preamble_validator if allow_imports else _validator
    v.visit(tree)
    _cache[cache_key] = True


# ---------------------------------------------------------------------------
# Restricted unpickler
# ---------------------------------------------------------------------------

# Allowlisted (module, qualname) prefixes for pickle deserialization.
# These cover the common model/data types used in pricing pipelines.
_ALLOWED_PICKLE_PREFIXES: list[tuple[str, ...]] = [
    ("numpy",),
    ("sklearn",),
    ("scipy",),
    ("catboost",),
    ("xgboost",),
    ("lightgbm",),
    ("pandas",),
    ("polars",),
    ("joblib",),
    ("collections",),
    ("builtins",),
    ("_codecs",),
    ("copyreg",),
    ("datetime",),
]


class _RestrictedUnpickler(pickle.Unpickler):
    """Unpickler that only allows known-safe classes.

    Prevents arbitrary code execution via crafted pickle payloads
    while still supporting common ML model and data formats.
    """

    def find_class(self, module: str, name: str) -> Any:
        for prefix in _ALLOWED_PICKLE_PREFIXES:
            if module.startswith(prefix[0]):
                return super().find_class(module, name)
        raise pickle.UnpicklingError(
            f"Blocked unpickling of {module}.{name} — "
            f"class not in the allowlist. If this is a legitimate model "
            f"class, add its module to _ALLOWED_PICKLE_PREFIXES in "
            f"src/haute/_sandbox.py"
        )


def safe_unpickle(path: str | Path) -> Any:
    """Deserialize a pickle file using the restricted unpickler.

    Also validates the path is within the project root.
    """
    validated = validate_project_path(path)
    with open(validated, "rb") as f:
        return _RestrictedUnpickler(f).load()


def safe_joblib_load(path: str | Path) -> Any:
    """Deserialize a joblib file using a restricted unpickler.

    ``joblib.load()`` uses pickle internally but provides no class
    restriction hook.  This function patches joblib's ``NumpyUnpickler``
    with the same ``find_class`` allowlist used by ``safe_unpickle``,
    then restores the original after loading.

    Also validates the path is within the project root.
    """
    validated = validate_project_path(path)

    try:
        from joblib.numpy_pickle import NumpyUnpickler
    except ImportError:
        # joblib not installed — fall back to restricted pickle
        logger.warning("joblib_missing", msg="falling back to safe_unpickle")
        return safe_unpickle(validated)

    original_find_class = NumpyUnpickler.find_class

    def _restricted_joblib_find_class(self: Any, module: str, name: str) -> Any:
        """find_class with allowlist, delegating to the original on match."""
        for prefix in _ALLOWED_PICKLE_PREFIXES:
            if module.startswith(prefix[0]):
                return original_find_class(self, module, name)
        raise pickle.UnpicklingError(
            f"Blocked unpickling of {module}.{name} — "
            f"class not in the allowlist. If this is a legitimate model "
            f"class, add its module to _ALLOWED_PICKLE_PREFIXES in "
            f"src/haute/_sandbox.py"
        )

    NumpyUnpickler.find_class = _restricted_joblib_find_class
    try:
        import joblib

        return joblib.load(validated)
    finally:
        NumpyUnpickler.find_class = original_find_class
