"""Pickleable wrappers for sandbox-defined functions and classes."""

import ast
import copy
import inspect
from typing import Any

from .builtins import SAFE_EXCEPTIONS, SAFE_FN_NAMES, make_safe_builtins

_SAFE_BUILTIN_NAMES = (
    set(SAFE_FN_NAMES)
    | set(SAFE_EXCEPTIONS)
    | {
        "True",
        "False",
        "None",
        "Ellipsis",
        "NotImplemented",
        "print",
        "getattr",
        "hasattr",
        "locals",
    }
)


def _collect_global_names(code: Any) -> set[str]:
    """Collect co_names from a code object and all nested code objects.

    Generator expressions, comprehensions, and nested functions create
    separate code objects.  Global names used inside them appear in the
    nested code's co_names, not the parent's.
    """
    names: set[str] = set(code.co_names)
    for const in code.co_consts:
        if hasattr(const, "co_names"):
            names.update(_collect_global_names(const))
    return names


class StFunction:
    """Callable, pickleable wrapper for a sandbox-defined function.

    Stores the rewritten AST so the function can be serialized and
    recompiled after deserialization.
    """

    def __init__(
        self,
        name: str,
        compiled_fn: Any,
        func_ast: ast.FunctionDef | ast.AsyncFunctionDef,
    ) -> None:
        self._name = name
        self._compiled = compiled_fn
        self._func_ast = func_ast
        self._sandbox: Any = None  # set by activate() or auto-activation
        self._gates: dict[str, Any] | None = None

        # Copy plain-data metadata from compiled function
        self.__name__ = name
        if compiled_fn is not None:
            self.__qualname__ = getattr(compiled_fn, "__qualname__", name)
            self.__doc__ = getattr(compiled_fn, "__doc__", None)
            self.__annotations__ = getattr(compiled_fn, "__annotations__", {})
            self.__defaults__ = getattr(compiled_fn, "__defaults__", None)
            self.__kwdefaults__ = getattr(compiled_fn, "__kwdefaults__", None)
            try:
                self.__signature__ = inspect.signature(compiled_fn)
            except (ValueError, TypeError):
                self.__signature__ = None
        else:
            self.__qualname__ = name
            self.__doc__ = None
            self.__annotations__ = {}
            self.__defaults__ = None
            self.__kwdefaults__ = None
            self.__signature__ = None

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if self._compiled is None:
            raise RuntimeError(
                f"StFunction '{self._name}' is not active "
                f"-- call activate() with gate functions first"
            )
        if self._sandbox is not None and self._gates is not None:
            return self._sandbox._call_in_context(
                self._compiled, self._gates, args, kwargs
            )
        return self._compiled(*args, **kwargs)

    def __repr__(self) -> str:
        status = "active" if self._compiled is not None else "inactive"
        return f"<StFunction '{self._name}' ({status})>"

    @property
    def global_refs(self) -> set[str]:
        """Names this function references as globals (excluding builtins/gates)."""
        stored = getattr(self, "_global_ref_names", None)
        if stored is not None:
            return set(stored)
        if self._compiled is not None and hasattr(self._compiled, "__code__"):
            return {
                n
                for n in _collect_global_names(self._compiled.__code__)
                if not n.startswith("__st_")
                and n not in _SAFE_BUILTIN_NAMES
                and n != self._name
            }
        return set()

    def __getstate__(self) -> dict[str, Any]:
        state = self.__dict__.copy()
        compiled = state.pop("_compiled", None)
        state.pop("_sandbox", None)
        state.pop("_gates", None)

        # Freeze closure variables from the compiled function
        # (skip if _compiled is a wrapper like StFunction, not a raw function)
        if (
            compiled is not None
            and hasattr(compiled, "__closure__")
            and compiled.__closure__
        ):
            freevars = compiled.__code__.co_freevars
            frozen: dict[str, Any] = {}
            for name, cell in zip(freevars, compiled.__closure__):
                if name.startswith("__st_"):
                    continue  # Gates are re-injected on activate
                try:
                    frozen[name] = cell.cell_contents
                except ValueError:
                    continue  # Empty cell
            if frozen:
                state["_frozen_closure"] = frozen

        # Freeze global references (StFunction/StClass only) and store
        # all global ref names for introspection via global_refs property
        # (skip if _compiled is a wrapper like StFunction, not a raw function)
        if compiled is not None and hasattr(compiled, "__globals__"):
            globs = compiled.__globals__
            frozen_globals: dict[str, Any] = {}
            global_ref_names: set[str] = set()
            for name in _collect_global_names(compiled.__code__):
                if name.startswith("__st_"):
                    continue
                if name in _SAFE_BUILTIN_NAMES:
                    continue
                if name == self._name:
                    continue
                global_ref_names.add(name)
                if name in globs and isinstance(globs[name], (StFunction, StClass)):
                    frozen_globals[name] = globs[name]
            if frozen_globals:
                state["_frozen_globals"] = frozen_globals
            if global_ref_names:
                state["_global_ref_names"] = global_ref_names

        return state

    def __setstate__(self, state: dict[str, Any]) -> None:
        self.__dict__.update(state)
        self._compiled = None
        self._sandbox = None
        self._gates = None

    def activate(
        self,
        gates: dict[str, Any],
        *,
        sandbox: Any = None,
        namespace: dict[str, Any] | None = None,
    ) -> None:
        """Recompile from stored AST with the given gate functions.

        Args:
            gates: Gate function dict (from make_gates).
            sandbox: Optional Sandbox reference for direct-call context.
            namespace: Optional namespace for globals the function references.

        Namespace priority (lowest to highest):
            1. Frozen globals (fallback from pickle)
            2. Caller namespace (late-binding override)
            3. Frozen closure (value captures, sacred)
            4. Gates + builtins (always on top)
        """
        if getattr(self, "_activating", False):
            return  # Guard against circular activation (e.g. mutual refs)
        self._activating = True
        try:
            self._activate_inner(gates, sandbox=sandbox, namespace=namespace)
        finally:
            self._activating = False

    def _activate_inner(
        self,
        gates: dict[str, Any],
        *,
        sandbox: Any = None,
        namespace: dict[str, Any] | None = None,
    ) -> None:
        ns: dict[str, Any] = {}

        frozen_globals = getattr(self, "_frozen_globals", None)
        if frozen_globals:
            ns.update(frozen_globals)

        if namespace:
            ns.update(namespace)

        frozen_closure = getattr(self, "_frozen_closure", None)
        if frozen_closure:
            ns.update(frozen_closure)

        ns.update(gates)
        ns["__builtins__"] = make_safe_builtins(gates["__st_getattr__"])
        ns["__name__"] = "__sandtrap__"

        # Auto-activate frozen globals and frozen closure values
        for source in (frozen_globals, frozen_closure):
            if source:
                for val in source.values():
                    activate_value(val, gates, sandbox=sandbox, namespace=ns)

        # Resolve ModuleRef objects in the namespace via the import gate
        import_gate = gates.get("__st_import__")
        if import_gate is not None:
            for k, v in list(ns.items()):
                if isinstance(v, ModuleRef):
                    try:
                        top = v.name.split(".")[0]
                        if k == top:
                            ns[k] = import_gate(v.name)
                        else:
                            ns[k] = import_gate(v.name, alias=k)
                    except Exception:
                        pass

        # Wrap function AST in a module for compilation
        func_copy = copy.deepcopy(self._func_ast)
        module = ast.Module(body=[func_copy], type_ignores=[])
        ast.fix_missing_locations(module)

        code = compile(module, f"<sandtrap:fn:{self._name}>", "exec")
        exec(code, ns)  # noqa: S102

        self._compiled = ns[self._name]
        self._sandbox = sandbox
        self._gates = gates


class StClass:
    """Pickleable wrapper for a sandbox-defined class.

    Stores the rewritten AST and frozen references to decorators/bases
    so the class can be serialized and recompiled after deserialization.
    """

    def __init__(
        self,
        name: str,
        compiled_cls: type,
        class_ast: ast.ClassDef,
        frozen_refs: dict[str, Any] | None = None,
    ) -> None:
        self._name = name
        self._compiled_cls = compiled_cls
        self._class_ast = class_ast
        self._frozen_refs = frozen_refs or {}
        self._st_getattr_gate: Any = None
        self._sandbox: Any = None
        self._gates: dict[str, Any] | None = None

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if self._compiled_cls is None:
            raise RuntimeError(
                f"StClass '{self._name}' is not active -- call activate() first"
            )
        if self._sandbox is not None and self._gates is not None:
            instance = self._sandbox._call_in_context(
                self._compiled_cls, self._gates, args, kwargs
            )
        else:
            instance = self._compiled_cls(*args, **kwargs)
        return StInstance(self, instance, self._st_getattr_gate)

    def __mro_entries__(self, bases: tuple) -> tuple:
        """Allow StClass to be used as a base class in class statements."""
        return (self._compiled_cls,)

    def __getattr__(self, name: str) -> Any:
        if name.startswith("_"):
            raise AttributeError(name)
        compiled = self.__dict__.get("_compiled_cls")
        if compiled is None:
            raise RuntimeError(
                f"StClass '{self._name}' is not active -- call activate() first"
            )
        return getattr(compiled, name)

    def __repr__(self) -> str:
        status = "active" if self._compiled_cls is not None else "inactive"
        return f"<StClass '{self._name}' ({status})>"

    def __getstate__(self) -> dict[str, Any]:
        state = self.__dict__.copy()
        state.pop("_compiled_cls", None)
        state.pop("_st_getattr_gate", None)
        state.pop("_sandbox", None)
        state.pop("_gates", None)
        return state

    def __setstate__(self, state: dict[str, Any]) -> None:
        self.__dict__.update(state)
        self._compiled_cls = None
        self._st_getattr_gate = None
        self._sandbox = None
        self._gates = None

    def activate(
        self,
        gates: dict[str, Any],
        *,
        sandbox: Any = None,
        namespace: dict[str, Any] | None = None,
    ) -> None:
        """Recompile class from stored AST with the given gate functions."""
        if getattr(self, "_activating", False):
            return  # Guard against circular activation (e.g. mutual base refs)
        self._activating = True
        try:
            self._activate_inner(gates, sandbox=sandbox, namespace=namespace)
        finally:
            self._activating = False

    def _activate_inner(
        self,
        gates: dict[str, Any],
        *,
        sandbox: Any = None,
        namespace: dict[str, Any] | None = None,
    ) -> None:
        # Frozen refs first (lowest priority), then caller namespace overrides
        ns: dict[str, Any] = {}
        if self._frozen_refs:
            ns.update(self._frozen_refs)
        if namespace:
            ns.update(namespace)
        ns.update(gates)
        ns["__builtins__"] = make_safe_builtins(gates["__st_getattr__"])
        ns["__name__"] = "__sandtrap__"

        # Auto-activate frozen ref deps — use the values from ns (which
        # may be namespace overrides of frozen refs) so the compiled class
        # sees the activated version.
        if self._frozen_refs:
            for name in self._frozen_refs:
                val = ns.get(name)
                activate_value(val, gates, sandbox=sandbox, namespace=ns)

        class_copy = copy.deepcopy(self._class_ast)
        module = ast.Module(body=[class_copy], type_ignores=[])
        ast.fix_missing_locations(module)

        code = compile(module, f"<sandtrap:cls:{self._name}>", "exec")
        exec(code, ns)  # noqa: S102

        self._compiled_cls = ns[self._name]
        self._st_getattr_gate = gates.get("__st_getattr__")
        self._sandbox = sandbox
        self._gates = gates


class StInstance:
    """Pickleable wrapper for an instance of a sandbox-defined class.

    Proxies attribute access to the underlying real instance and stores
    the StClass reference + instance __dict__ for serialization.
    """

    __slots__ = ("_st_class", "_st_instance", "_frozen_attrs", "_st_getattr_gate")

    def __init__(
        self, sb_class: StClass, instance: Any, getattr_gate: Any = None
    ) -> None:
        object.__setattr__(self, "_st_class", sb_class)
        object.__setattr__(self, "_st_instance", instance)
        object.__setattr__(self, "_st_getattr_gate", getattr_gate)

    def __getattr__(self, name: str) -> Any:
        instance = object.__getattribute__(self, "_st_instance")
        if instance is None:
            raise RuntimeError("StInstance is not active -- call activate() first")
        gate = object.__getattribute__(self, "_st_getattr_gate")
        if gate is not None:
            return gate(instance, name)
        return getattr(instance, name)

    def __setattr__(self, name: str, value: Any) -> None:
        instance = object.__getattribute__(self, "_st_instance")
        if instance is None:
            raise RuntimeError("StInstance is not active -- call activate() first")
        setattr(instance, name, value)

    def __delattr__(self, name: str) -> None:
        instance = object.__getattribute__(self, "_st_instance")
        if instance is None:
            raise RuntimeError("StInstance is not active -- call activate() first")
        delattr(instance, name)

    def __repr__(self) -> str:
        try:
            instance = object.__getattribute__(self, "_st_instance")
            if instance is not None:
                return repr(instance)
        except AttributeError:
            pass
        sb_class = object.__getattribute__(self, "_st_class")
        return f"<StInstance of '{sb_class._name}' (inactive)>"

    def __getstate__(self) -> dict[str, Any]:
        instance = object.__getattribute__(self, "_st_instance")
        attrs = instance.__dict__.copy() if instance is not None else {}
        return {
            "st_class": object.__getattribute__(self, "_st_class"),
            "attrs": attrs,
        }

    def __setstate__(self, state: dict[str, Any]) -> None:
        object.__setattr__(self, "_st_class", state["st_class"])
        object.__setattr__(self, "_st_instance", None)
        object.__setattr__(self, "_frozen_attrs", state["attrs"])
        object.__setattr__(self, "_st_getattr_gate", None)

    def activate(
        self,
        *,
        gates: dict[str, Any] | None = None,
        sandbox: Any = None,
        namespace: dict[str, Any] | None = None,
    ) -> None:
        """Restore the real instance from frozen attrs.

        The StClass must be activated before calling this.
        When gates/sandbox/namespace are provided, nested StFunction/StClass/
        StInstance values in frozen attrs are auto-activated.
        """
        sb_class = object.__getattribute__(self, "_st_class")
        if sb_class._compiled_cls is None:
            raise RuntimeError("StClass must be activated before its instances")
        cls = sb_class._compiled_cls
        instance = cls.__new__(cls)
        frozen = object.__getattribute__(self, "_frozen_attrs")

        # Auto-activate nested wrappers in frozen attrs
        if gates is not None:
            for val in frozen.values():
                activate_value(val, gates, sandbox=sandbox, namespace=namespace)

        instance.__dict__.update(frozen)
        object.__setattr__(self, "_st_instance", instance)
        object.__setattr__(self, "_st_getattr_gate", sb_class._st_getattr_gate)


def activate_value(
    val: Any,
    gates: dict[str, Any],
    *,
    sandbox: Any = None,
    namespace: dict[str, Any] | None = None,
) -> None:
    """Activate an inactive StFunction, StClass, or StInstance."""
    if isinstance(val, StFunction) and val._compiled is None:
        val.activate(gates, sandbox=sandbox, namespace=namespace)
    elif isinstance(val, StClass) and val._compiled_cls is None:
        val.activate(gates, sandbox=sandbox, namespace=namespace)
    elif isinstance(val, StInstance):
        sb_class = object.__getattribute__(val, "_st_class")
        if sb_class._compiled_cls is None:
            sb_class.activate(gates, sandbox=sandbox, namespace=namespace)
        if object.__getattribute__(val, "_st_instance") is None:
            val.activate(gates=gates, sandbox=sandbox, namespace=namespace)


def _make_dunder_forwarder(name: str):
    """Create a forwarding method for a dunder on StInstance.

    These go directly to the underlying instance, bypassing the attr gate.
    Protocol dunders (__len__, __iter__, __add__, etc.) are safe — they only
    operate on the object's own data and don't expose interpreter internals.
    """

    def forwarder(self, *args, **kwargs):
        instance = object.__getattribute__(self, "_st_instance")
        return getattr(instance, name)(*args, **kwargs)

    forwarder.__name__ = name
    forwarder.__qualname__ = f"StInstance.{name}"
    return forwarder


# Dunders that need forwarding for implicit protocol dispatch (str(), len(), etc.).
# __repr__ is handled explicitly above; __init__/__getattr__/__setattr__/__delattr__
# are part of StInstance's own proxy machinery.
_FORWARDED_DUNDERS = [
    "__str__",
    "__len__",
    "__bool__",
    "__hash__",
    "__iter__",
    "__next__",
    "__contains__",
    "__getitem__",
    "__setitem__",
    "__delitem__",
    "__eq__",
    "__ne__",
    "__lt__",
    "__le__",
    "__gt__",
    "__ge__",
    "__add__",
    "__radd__",
    "__sub__",
    "__rsub__",
    "__mul__",
    "__rmul__",
    "__truediv__",
    "__rtruediv__",
    "__floordiv__",
    "__rfloordiv__",
    "__mod__",
    "__rmod__",
    "__pow__",
    "__rpow__",
    "__neg__",
    "__pos__",
    "__abs__",
    "__int__",
    "__float__",
    "__index__",
    "__enter__",
    "__exit__",
    "__aenter__",
    "__aexit__",
    "__aiter__",
    "__anext__",
    "__call__",
]

for _dname in _FORWARDED_DUNDERS:
    setattr(StInstance, _dname, _make_dunder_forwarder(_dname))


class ModuleRef:
    """Pickleable reference to a module.

    Python modules can't survive pickle.  This ref stores just the module
    name so ``Sandbox._auto_activate`` can re-import it via the
    ``__st_import__`` gate on the next turn.  Works for both VFS modules
    and policy-registered modules (including aliased imports like
    ``import math as m``).
    """

    def __init__(self, name: str, file: str | None = None) -> None:
        self.name = name
        self.file = file

    def __repr__(self) -> str:
        return f"<ModuleRef '{self.name}'>"
