"""Sandbox helpers for Excel execute_code.

This is not process isolation. It narrows what Python code running inside
Excel can do by:

- exposing a restricted builtin set
- allowing imports only from a small stdlib allowlist
- patching reachable module globals so exposed helpers cannot recover the
  unrestricted builtin table
- reactively blocking dangerous audit events while user code is executing
"""

from __future__ import annotations

import builtins
import os
import sys
import threading
from contextlib import contextmanager

# User code is serialized onto the main thread via `_run_on_main`, so a single
# process-global flag is safe. Making it global (rather than thread-local)
# means the audit hook blocks on *any* thread while a sandbox is active —
# closes escapes where user code reaches a threading primitive via
# `sys.modules` and spawns a worker that wouldn't inherit a thread-local flag.
_SANDBOX_ACTIVE_DEPTH = 0
_SANDBOX_LOCAL = threading.local()
_AUDIT_HOOK_INSTALLED = False

_BLOCKED_IMPORT_ROOTS = frozenset(
    {
        '_thread',
        'asyncio',
        'builtins',
        'concurrent',
        'ctypes',
        'ensurepip',
        'importlib',
        'multiprocessing',
        'os',
        'pip',
        'runpy',
        'site',
        'socket',
        'subprocess',
        'sys',
        'threading',
        'venv',
        'webbrowser',
    }
)

_BLOCKED_BUILTINS = frozenset(
    {
        '__import__',
        'breakpoint',
        'compile',
        'eval',
        'exec',
        'exit',
        'input',
        'memoryview',
        'open',
        'quit',
        '__loader__',
        '__spec__',
    }
)

_BLOCKED_AUDIT_EVENTS = frozenset(
    {
        'os.mkdir',
        'os.remove',
        'os.rename',
        'os.replace',
        'os.rmdir',
        'ctypes.dlopen',
        'os.exec',
        'os.popen',
        'os.spawn',
        'os.spawnv',
        'os.spawnve',
        'os.startfile',
        'os.system',
        'os.unlink',
        'socket.bind',
        'socket.connect',
        'socket.listen',
        'socket.sendto',
        'subprocess.Popen',
        'webbrowser.open',
    }
)

_WRITE_OPEN_FLAGS = os.O_WRONLY | os.O_RDWR | os.O_APPEND | os.O_CREAT | os.O_TRUNC
_ORIGINAL_BUILTINS = builtins.__dict__
_ORIGINAL_IMPORT = builtins.__import__


def _is_sandbox_active() -> bool:
    return _SANDBOX_ACTIVE_DEPTH > 0


def _current_safe_builtins() -> dict | None:
    return getattr(_SANDBOX_LOCAL, 'safe_builtins', None)


def _patched_modules() -> dict[int, tuple[dict, object]]:
    patched = getattr(_SANDBOX_LOCAL, 'patched_modules', None)
    if patched is None:
        patched = {}
        _SANDBOX_LOCAL.patched_modules = patched
    return patched


def _is_blocked_import(name: str, level: int) -> bool:
    if level != 0:
        return True
    root = name.partition('.')[0]
    return root in _BLOCKED_IMPORT_ROOTS


def _safe_import(name: str, globals=None, locals=None, fromlist=(), level: int = 0):
    if _is_blocked_import(name, level):
        raise PermissionError(f"Import '{name}' is not allowed in Excel execute_code.")
    module = _ORIGINAL_IMPORT(name, globals, locals, fromlist, level)
    safe_builtins = _current_safe_builtins()
    if safe_builtins is not None:
        _patch_loaded_module_builtins(safe_builtins, name)
    return module


def build_safe_builtins() -> dict:
    safe = {k: v for k, v in _ORIGINAL_BUILTINS.items() if k not in _BLOCKED_BUILTINS}
    safe['__import__'] = _safe_import
    return safe


def _should_patch_module(name: str, import_root: str | None = None) -> bool:
    if name == 'shortcut_xl' or name.startswith('shortcut_xl.'):
        return True
    if import_root is None:
        return False
    root = import_root.partition('.')[0]
    return name == root or name.startswith(f'{root}.')


def _patch_loaded_module_builtins(safe_builtins: dict, import_root: str | None = None) -> None:
    patched = _patched_modules()
    for name, module in list(sys.modules.items()):
        if module is None or not _should_patch_module(name, import_root):
            continue
        module_dict = getattr(module, '__dict__', None)
        if not isinstance(module_dict, dict):
            continue
        key = id(module_dict)
        if key in patched:
            continue
        patched[key] = (module_dict, module_dict.get('__builtins__', _ORIGINAL_BUILTINS))
        module_dict['__builtins__'] = safe_builtins


def _restore_patched_module_builtins() -> None:
    patched = getattr(_SANDBOX_LOCAL, 'patched_modules', None)
    if not patched:
        return
    for module_dict, original_builtins in reversed(list(patched.values())):
        module_dict['__builtins__'] = original_builtins
    patched.clear()


def _is_write_open(args: tuple) -> bool:
    mode = args[1] if len(args) > 1 else None
    flags = args[2] if len(args) > 2 else 0
    if isinstance(mode, str) and any(flag in mode for flag in ('w', 'a', 'x', '+')):
        return True
    if isinstance(flags, int) and flags & _WRITE_OPEN_FLAGS:
        return True
    return False


def sandbox_audit_hook(event: str, args: tuple) -> None:
    # Defense in depth only: audit hooks observe Python runtime events after
    # execution has already entered the process. They are useful for blocking
    # dangerous operations, but they are not equivalent to OS/process isolation.
    if not _is_sandbox_active():
        return
    if event in _BLOCKED_AUDIT_EVENTS:
        raise PermissionError(f"Operation '{event}' is not allowed in Excel execute_code.")
    if event == 'open' and _is_write_open(args):
        raise PermissionError('File writes are not allowed in Excel execute_code.')


def install_audit_hook() -> None:
    global _AUDIT_HOOK_INSTALLED
    if _AUDIT_HOOK_INSTALLED:
        return
    sys.addaudithook(sandbox_audit_hook)
    _AUDIT_HOOK_INSTALLED = True


@contextmanager
def active_exec_sandbox(safe_builtins: dict):
    global _SANDBOX_ACTIVE_DEPTH
    prev_safe_builtins = getattr(_SANDBOX_LOCAL, 'safe_builtins', None)
    prev_patched_modules = getattr(_SANDBOX_LOCAL, 'patched_modules', None)

    _SANDBOX_ACTIVE_DEPTH += 1
    _SANDBOX_LOCAL.safe_builtins = safe_builtins
    _SANDBOX_LOCAL.patched_modules = {}
    _patch_loaded_module_builtins(safe_builtins)

    try:
        yield
    finally:
        _restore_patched_module_builtins()
        _SANDBOX_ACTIVE_DEPTH -= 1
        _SANDBOX_LOCAL.safe_builtins = prev_safe_builtins
        _SANDBOX_LOCAL.patched_modules = prev_patched_modules


install_audit_hook()
