"""V2 restricted quantum executor; v1 remains frozen in quantum_sandbox.py.

AST screening is combined with nobody UID, cleared environment, process/resource
limits and seccomp denial of filesystem opens, networking and process execution.
Libraries are imported/warmed in the parent before any untrusted code runs.
"""

import ast
import contextlib
import ctypes
import ctypes.util
import errno
import importlib
import io
import multiprocessing as mp
import os
import resource

for variable in ["OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"]:
    os.environ[variable] = "1"

import numpy as np

MODULES = {
    "numpy",
    "numpy.linalg",
    "math",
    "cmath",
    "qiskit",
    "qiskit.quantum_info",
    "pennylane",
    "cirq",
}
ATTRIBUTES = set(
    "array asarray zeros ones eye pi sqrt exp sin cos conj conjugate transpose reshape flatten ravel squeeze astype complex64 complex128 float32 float64 norm kron real imag T shape size data QuantumCircuit Statevector from_instruction h x y z ry rx rz cx cz swap s sdg t tdg sx barrier append device qnode state RY RX RZ Rot PhaseShift Identity PauliX PauliY PauliZ Hadamard CNOT CZ SWAP LineQubit range Circuit Simulator simulate final_state_vector H X Y Z CXPowGate ZPowGate XPowGate HPowGate controlled on inverse quantum_info".split()
)
# Documented Cirq gate aliases, independent of any v2 model outputs. They were
# absent from v1's conservative supported API. No isolation rule is relaxed.
ATTRIBUTES.update({"CX", "Rx", "Ry", "Rz"})
BUILTINS = {
    "range": lambda *a: bounded_range(*a),
    "enumerate": enumerate,
    "zip": zip,
    "len": len,
    "list": list,
    "tuple": tuple,
    "float": float,
    "int": int,
    "complex": complex,
    "print": print,
    "abs": abs,
    "sum": sum,
    "min": min,
    "max": max,
    "reversed": reversed,
}
FORBIDDEN = (
    ast.ClassDef,
    ast.AsyncFunctionDef,
    ast.Await,
    ast.While,
    ast.Try,
    ast.With,
    ast.AsyncWith,
    ast.Raise,
    ast.Delete,
    ast.Global,
    ast.Nonlocal,
    ast.Lambda,
    ast.Yield,
    ast.YieldFrom,
    ast.Match,
)
LIB = None


def bounded_range(*args):
    value = range(*args)
    if len(value) > 1024:
        raise ValueError("range limit")
    return value


def screen(program):
    if len(program) > 25000:
        raise ValueError("program length limit")
    tree = ast.parse(program)
    if sum(1 for _ in ast.walk(tree)) > 6000:
        raise ValueError("AST node limit")
    for node in ast.walk(tree):
        if isinstance(node, FORBIDDEN):
            raise ValueError("unsupported syntax: " + type(node).__name__)
        if isinstance(node, ast.Name) and (
            node.id.startswith("_")
            or node.id
            in {
                "eval",
                "exec",
                "compile",
                "open",
                "getattr",
                "setattr",
                "globals",
                "locals",
                "vars",
                "type",
                "object",
                "input",
                "help",
                "breakpoint",
            }
        ):
            raise ValueError("forbidden name")
        if isinstance(node, ast.Attribute) and node.attr not in ATTRIBUTES:
            raise ValueError("unsupported attribute: " + node.attr)
        if isinstance(node, ast.Import):
            if any(a.name not in MODULES or (a.asname or "").startswith("_") for a in node.names):
                raise ValueError("unsupported import")
        if isinstance(node, ast.ImportFrom):
            if (
                node.level
                or node.module not in MODULES
                or any(
                    a.name not in ATTRIBUTES or (a.asname or "").startswith("_") for a in node.names
                )
            ):
                raise ValueError("unsupported from import")
        if isinstance(node, ast.FunctionDef) and node.name.startswith("_"):
            raise ValueError("forbidden function name")
    return tree


def initialize():
    global LIB
    assert os.name == "posix" and os.getuid() == 0, "Linux root parent required for UID isolation"
    # Resolve/load seccomp before restricting filesystem access.
    library = ctypes.util.find_library("seccomp")
    assert library, "libseccomp is required; never fall back to AST-only execution"
    LIB = ctypes.CDLL(library, use_errno=True)
    LIB.seccomp_init.argtypes = [ctypes.c_uint32]
    LIB.seccomp_init.restype = ctypes.c_void_p
    LIB.seccomp_syscall_resolve_name.argtypes = [ctypes.c_char_p]
    LIB.seccomp_syscall_resolve_name.restype = ctypes.c_int
    LIB.seccomp_rule_add.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_int, ctypes.c_uint]
    LIB.seccomp_load.argtypes = [ctypes.c_void_p]
    LIB.seccomp_release.argtypes = [ctypes.c_void_p]
    from qiskit import QuantumCircuit
    from qiskit.quantum_info import Statevector
    import pennylane as qml
    import cirq

    qc = QuantumCircuit(2)
    qc.h(0)
    qc.ry(0.3, 1)
    qc.rz(0.17, 0)
    qc.x(0)
    qc.z(1)
    qc.cx(0, 1)
    Statevector.from_instruction(qc)
    device = qml.device("default.qubit", wires=[1, 0])

    @qml.qnode(device)
    def warm():
        qml.Hadamard(0)
        qml.RY(0.3, wires=1)
        qml.RZ(0.17, wires=0)
        qml.PauliX(0)
        qml.PauliZ(1)
        qml.CNOT(wires=[0, 1])
        return qml.state()

    warm()
    qubits = cirq.LineQubit.range(2)
    cirq.Simulator(dtype=np.complex128).simulate(
        cirq.Circuit(
            cirq.H(qubits[0]),
            cirq.ry(0.3)(qubits[1]),
            cirq.rz(0.17)(qubits[0]),
            cirq.X(qubits[0]),
            cirq.Z(qubits[1]),
            cirq.CNOT(*qubits),
        ),
        qubit_order=list(reversed(qubits)),
    )
    # Warm Python import machinery and serialization before fork/lockdown.
    for module in MODULES:
        importlib.import_module(module)
    importlib.import_module("socket")
    return dict(
        qiskit=importlib.import_module("qiskit").__version__,
        pennylane=qml.__version__,
        cirq=cirq.__version__,
    )


def lockdown():
    os.environ.clear()
    os.setgroups([])
    os.setgid(65534)
    os.setuid(65534)
    resource.setrlimit(resource.RLIMIT_CPU, (15, 15))
    # Native numeric runtimes reserve large virtual address spaces. Physical
    # memory remains bounded by the HF CPU job container's memory allocation.
    resource.setrlimit(resource.RLIMIT_AS, (16 * 1024**3, 16 * 1024**3))
    resource.setrlimit(resource.RLIMIT_FSIZE, (0, 0))
    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
    resource.setrlimit(resource.RLIMIT_NPROC, (32, 32))
    libc = ctypes.CDLL(None, use_errno=True)
    assert libc.prctl(38, 1, 0, 0, 0) == 0, "PR_SET_NO_NEW_PRIVS failed"
    policy = LIB.seccomp_init(0x7FFF0000)  # SCMP_ACT_ALLOW
    assert policy
    deny = 0x00050000 | errno.EPERM  # SCMP_ACT_ERRNO(EPERM)
    for name in [
        "open",
        "openat",
        "openat2",
        "creat",
        "socket",
        "socketpair",
        "connect",
        "accept",
        "accept4",
        "bind",
        "listen",
        "sendto",
        "sendmsg",
        "execve",
        "execveat",
        "ptrace",
        "process_vm_readv",
        "process_vm_writev",
        "mount",
        "umount2",
        "bpf",
        "userfaultfd",
    ]:
        syscall = LIB.seccomp_syscall_resolve_name(name.encode())
        if syscall >= 0:
            assert LIB.seccomp_rule_add(policy, deny, syscall, 0) == 0, name
    assert LIB.seccomp_load(policy) == 0, "seccomp load failed"
    LIB.seccomp_release(policy)


def safe_import(name, globals=None, locals=None, fromlist=(), level=0):
    if name not in MODULES or level:
        raise ImportError("module not allowed")
    module = importlib.import_module(name)
    return module if fromlist else importlib.import_module(name.split(".")[0])


def worker(program, qubits, connection, mode):
    try:
        lockdown()
        if mode == "isolation_check":
            import socket

            file_blocked = network_blocked = False
            try:
                open("/etc/hostname").read()
            except PermissionError:
                file_blocked = True
            try:
                socket.socket()
            except PermissionError:
                network_blocked = True
            connection.send(
                dict(
                    file_opens_blocked=file_blocked,
                    network_blocked=network_blocked,
                    uid=os.getuid(),
                    environment_empty=not os.environ,
                )
            )
            return
        tree = screen(program)
        safe = dict(BUILTINS, __import__=safe_import)
        namespace = {"__builtins__": safe}
        with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
            exec(compile(tree, "<restricted-quantum>", "exec"), namespace, namespace)
        value = np.asarray(namespace["state"], dtype=np.complex128).reshape(-1)
        assert value.shape == (2**qubits,) and np.isfinite(value).all()
        connection.send(
            dict(
                status="executed",
                real=value.real.tolist(),
                imag=value.imag.tolist(),
                norm=float(np.linalg.norm(value)),
            )
        )
    except BaseException as error:
        connection.send(
            dict(status="rejected_or_error", error=type(error).__name__, detail=str(error)[:300])
        )
    finally:
        connection.close()


def run(program, qubits, mode="program"):
    assert 1 <= qubits <= 8
    ctx = mp.get_context("fork")
    receive, send = ctx.Pipe(duplex=False)
    process = ctx.Process(target=worker, args=(program, qubits, send, mode))
    process.start()
    send.close()
    if receive.poll(20):
        try:
            result = receive.recv()
        except EOFError:
            result = dict(status="worker_exit", exitcode=process.exitcode)
    else:
        result = dict(status="timeout")
    process.join(1)
    if process.is_alive():
        process.kill()
        process.join(2)
    receive.close()
    return result