rustpython-vm 0.6.0

use crate::{
    AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
    builtins::{PyCode, PyDictRef, PyNamespace, PyUtf8StrRef, code::CoMonitoringData},
    function::FuncArgs,
};
use core::sync::atomic::Ordering;
use crossbeam_utils::atomic::AtomicCell;
use std::collections::{HashMap, HashSet};

pub(crate) const TOOL_LIMIT: usize = 6;
const EVENTS_COUNT: usize = 19;
const LOCAL_EVENTS_COUNT: usize = 11;
const UNGROUPED_EVENTS_COUNT: usize = 18;

/// Event identifier (`PY_MONITORING_EVENT_*`), stored in `tstate->what_event`.
/// `None` on the VM field is the `< 0` sentinel (not in a monitoring callback).
#[repr(i32)]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub(crate) enum MonitoringEvent {
    PyStart = 0,
    PyResume = 1,
    PyReturn = 2,
    PyYield = 3,
    Call = 4,
    Line = 5,
    Instruction = 6,
    Jump = 7,
    BranchLeft = 8,
    BranchRight = 9,
    StopIteration = 10,
    Raise = 11,
    ExceptionHandled = 12,
    PyUnwind = 13,
    PyThrow = 14,
    Reraise = 15,
    CReturn = 16,
    CRaise = 17,
    Branch = 18,
}

impl MonitoringEvent {
    pub(crate) const fn mask(self) -> u32 {
        1 << (self as u32)
    }

    pub(crate) const fn from_id(id: usize) -> Option<Self> {
        match id {
            0 => Some(Self::PyStart),
            1 => Some(Self::PyResume),
            2 => Some(Self::PyReturn),
            3 => Some(Self::PyYield),
            4 => Some(Self::Call),
            5 => Some(Self::Line),
            6 => Some(Self::Instruction),
            7 => Some(Self::Jump),
            8 => Some(Self::BranchLeft),
            9 => Some(Self::BranchRight),
            10 => Some(Self::StopIteration),
            11 => Some(Self::Raise),
            12 => Some(Self::ExceptionHandled),
            13 => Some(Self::PyUnwind),
            14 => Some(Self::PyThrow),
            15 => Some(Self::Reraise),
            16 => Some(Self::CReturn),
            17 => Some(Self::CRaise),
            18 => Some(Self::Branch),
            _ => None,
        }
    }
}

const C_RETURN_MASK: u32 = MonitoringEvent::CReturn.mask() | MonitoringEvent::CRaise.mask();

const EVENT_NAMES: [&str; EVENTS_COUNT] = [
    "PY_START",
    "PY_RESUME",
    "PY_RETURN",
    "PY_YIELD",
    "CALL",
    "LINE",
    "INSTRUCTION",
    "JUMP",
    "BRANCH_LEFT",
    "BRANCH_RIGHT",
    "STOP_ITERATION",
    "RAISE",
    "EXCEPTION_HANDLED",
    "PY_UNWIND",
    "PY_THROW",
    "RERAISE",
    "C_RETURN",
    "C_RAISE",
    "BRANCH",
];

/// Interpreter-level monitoring state, shared by all threads.
pub struct MonitoringState {
    pub tool_names: [Option<String>; TOOL_LIMIT],
    pub global_events: [u32; TOOL_LIMIT],
    pub local_events: HashMap<(usize, usize), u32>,
    pub callbacks: HashMap<(usize, usize), PyObjectRef>,
    /// Per-instruction disabled tools: (code_id, offset, tool)
    pub disabled: HashSet<(usize, usize, usize)>,
    /// Cached MISSING sentinel singleton
    pub missing: Option<PyObjectRef>,
    /// Cached DISABLE sentinel singleton
    pub disable: Option<PyObjectRef>,
}

impl Default for MonitoringState {
    fn default() -> Self {
        Self {
            tool_names: Default::default(),
            global_events: [0; TOOL_LIMIT],
            local_events: HashMap::new(),
            callbacks: HashMap::new(),
            disabled: HashSet::new(),
            missing: None,
            disable: None,
        }
    }
}

impl MonitoringState {
    /// Compute the OR of all tools' global_events + local_events.
    /// This is used for the fast-path atomic mask to skip monitoring
    /// when no events are registered at all.
    pub fn combined_events(&self) -> u32 {
        let global = self.global_events.iter().fold(0, |acc, &e| acc | e);
        let local = self.local_events.values().fold(0, |acc, &e| acc | e);
        global | local
    }

    /// Compute the events that apply to a specific code object:
    /// global events OR'd with local events registered for that code.
    /// This prevents events like INSTRUCTION that are local to one code
    /// from being applied to unrelated code objects.
    pub fn events_for_code(&self, code_id: usize) -> u32 {
        let global = self.global_events.iter().fold(0, |acc, &e| acc | e);
        let local = self
            .local_events
            .iter()
            .filter(|((_, cid), _)| *cid == code_id)
            .fold(0, |acc, (_, &e)| acc | e);
        global | local
    }
}

/// Global atomic mask: OR of all tools' events. Checked in the hot path
/// to skip monitoring overhead when no events are registered.
/// Lives in PyGlobalState alongside the PyMutex<MonitoringState>.
pub(crate) type MonitoringEventsMask = AtomicCell<u32>;

/// Get the MISSING sentinel, creating it if necessary.
pub(crate) fn get_missing(vm: &VirtualMachine) -> PyObjectRef {
    let mut state = vm.state.monitoring.lock();
    if let Some(ref m) = state.missing {
        m.clone()
    } else {
        let m: PyObjectRef = sys_monitoring::MonitoringSentinel.into_ref(&vm.ctx).into();
        state.missing = Some(m.clone());
        m
    }
}

/// Get the DISABLE sentinel, creating it if necessary.
pub(crate) fn get_disable(vm: &VirtualMachine) -> PyObjectRef {
    let mut state = vm.state.monitoring.lock();
    if let Some(ref d) = state.disable {
        d.clone()
    } else {
        let d: PyObjectRef = sys_monitoring::MonitoringSentinel.into_ref(&vm.ctx).into();
        state.disable = Some(d.clone());
        d
    }
}

fn check_valid_tool(tool_id: i32, vm: &VirtualMachine) -> PyResult<usize> {
    if !(0..TOOL_LIMIT as i32).contains(&tool_id) {
        return Err(vm.new_value_error(format!("invalid tool {tool_id} (must be between 0 and 5)")));
    }
    Ok(tool_id as usize)
}

fn check_tool_in_use(tool: usize, vm: &VirtualMachine) -> PyResult<()> {
    let state = vm.state.monitoring.lock();
    if state.tool_names[tool].is_some() {
        Ok(())
    } else {
        Err(vm.new_value_error(format!("tool {tool} is not in use")))
    }
}

fn parse_single_event(event: i32, vm: &VirtualMachine) -> PyResult<usize> {
    let event = u32::try_from(event)
        .map_err(|_| vm.new_value_error("The callback can only be set for one event at a time"))?;
    if event.count_ones() != 1 {
        return Err(vm.new_value_error("The callback can only be set for one event at a time"));
    }
    let event_id = event.trailing_zeros() as usize;
    if MonitoringEvent::from_id(event_id).is_none() {
        return Err(vm.new_value_error(format!("invalid event {event}")));
    }
    Ok(event_id)
}

fn normalize_event_set(event_set: i32, local: bool, vm: &VirtualMachine) -> PyResult<u32> {
    let kind = if local {
        "local event set"
    } else {
        "event set"
    };
    if event_set < 0 {
        return Err(vm.new_value_error(format!("invalid {kind} 0x{event_set:x}")));
    }

    let mut event_set = event_set as u32;
    if event_set >= (1 << EVENTS_COUNT) {
        return Err(vm.new_value_error(format!("invalid {kind} 0x{event_set:x}")));
    }

    if (event_set & C_RETURN_MASK) != 0
        && (event_set & MonitoringEvent::Call.mask()) != MonitoringEvent::Call.mask()
    {
        return Err(vm.new_value_error("cannot set C_RETURN or C_RAISE events independently"));
    }

    event_set &= !C_RETURN_MASK;

    if (event_set & MonitoringEvent::Branch.mask()) != 0 {
        event_set &= !MonitoringEvent::Branch.mask();
        event_set |= MonitoringEvent::BranchLeft.mask() | MonitoringEvent::BranchRight.mask();
    }

    if local && event_set >= (1 << LOCAL_EVENTS_COUNT) {
        return Err(vm.new_value_error(format!("invalid local event set 0x{event_set:x}")));
    }

    Ok(event_set)
}

/// Event that causes this opcode to be rewritten to INSTRUMENTED_*.
/// RESUME uses `oparg`: 0 → PY_START, nonzero → PY_RESUME.
fn event_for_opcode(
    op: rustpython_compiler_core::bytecode::Instruction,
    oparg: u8,
) -> Option<MonitoringEvent> {
    use rustpython_compiler_core::bytecode::Opcode;
    match op.deoptimize().as_opcode() {
        Opcode::ReturnValue => Some(MonitoringEvent::PyReturn),
        Opcode::Call | Opcode::CallKw | Opcode::CallFunctionEx | Opcode::LoadSuperAttr => {
            Some(MonitoringEvent::Call)
        }
        Opcode::YieldValue => Some(MonitoringEvent::PyYield),
        Opcode::JumpForward | Opcode::JumpBackward => Some(MonitoringEvent::Jump),
        Opcode::PopJumpIfFalse
        | Opcode::PopJumpIfTrue
        | Opcode::PopJumpIfNone
        | Opcode::PopJumpIfNotNone => Some(MonitoringEvent::BranchRight),
        Opcode::ForIter => Some(MonitoringEvent::BranchLeft),
        Opcode::PopIter => Some(MonitoringEvent::BranchRight),
        Opcode::EndFor | Opcode::EndSend => Some(MonitoringEvent::StopIteration),
        Opcode::NotTaken => Some(MonitoringEvent::BranchLeft),
        Opcode::EndAsyncFor => Some(MonitoringEvent::BranchRight),
        Opcode::Resume => {
            if oparg != 0 {
                Some(MonitoringEvent::PyResume)
            } else {
                Some(MonitoringEvent::PyStart)
            }
        }
        _ => None,
    }
}

fn opcode_event_is_active(
    op: rustpython_compiler_core::bytecode::Instruction,
    oparg: u8,
    events: u32,
) -> bool {
    event_for_opcode(op, oparg).is_some_and(|ev| events & ev.mask() != 0)
}

/// Walk real instructions, skipping specialization CACHE payloads.
/// Cache slots may contain pointer bits that are not valid opcodes.
fn for_each_instruction(
    code: &Py<PyCode>,
    mut f: impl FnMut(usize, rustpython_compiler_core::bytecode::Instruction, u8),
) {
    let len = code.code.instructions.len();
    let mut i = 0;
    while i < len {
        let op = code.code.instructions.read_op(i);
        let oparg = code.code.instructions.read_arg(i).as_u8();
        let caches = op.deoptimize().cache_entries();
        f(i, op, oparg);
        i += 1 + caches;
    }
}

/// Rewrite a code object's bytecode in-place with layered instrumentation.
///
/// Three layers (outermost first):
/// 1. INSTRUMENTED_LINE — wraps line-start instructions (stores original in side-table)
/// 2. INSTRUMENTED_INSTRUCTION — wraps all traceable instructions (stores original in side-table)
/// 3. Regular INSTRUMENTED_* — direct 1:1 opcode swap (no side-table needed)
///
/// De-instrumentation peels layers in reverse order.
/// Specialized opcodes are restored to base only when their EVENT_FOR_OPCODE is active.
pub(crate) fn instrument_code(code: &Py<PyCode>, events: u32) {
    use rustpython_compiler_core::bytecode::{self, Instruction};

    let len = code.code.instructions.len();
    let mut monitoring_data = code.monitoring_data.lock();

    // === Phase 1-3: De-instrument all layers (outermost first) ===

    // Phase 1: Remove INSTRUMENTED_LINE → restore from side-table
    if let Some(data) = monitoring_data.as_mut() {
        for (i, opcode) in data.line_opcodes.iter_mut().enumerate().take(len) {
            if *opcode != 0 {
                let original =
                    Instruction::try_from(*opcode).expect("invalid opcode in line side-table");
                unsafe {
                    code.code.instructions.replace_op(i, original);
                }
                *opcode = 0;
            }
        }
    }

    // Phase 2: Remove INSTRUMENTED_INSTRUCTION → restore from side-table
    if let Some(data) = monitoring_data.as_mut() {
        for (i, opcode) in data
            .per_instruction_opcodes
            .iter_mut()
            .enumerate()
            .take(len)
        {
            if *opcode != 0 {
                let original = Instruction::try_from(*opcode)
                    .expect("invalid opcode in instruction side-table");
                unsafe {
                    code.code.instructions.replace_op(i, original);
                }
                *opcode = 0;
            }
        }
    }

    // Phase 3: Restore INSTRUMENTED_* whose event is no longer active.
    // Restore specialized opcodes to base only when that family will be instrumented.
    {
        let mut i = 0;
        while i < len {
            let op = code.code.instructions.read_op(i);
            let oparg = code.code.instructions.read_arg(i).as_u8();
            let base_op = op.deoptimize();
            let caches = base_op.cache_entries();
            // LINE/INSTRUCTION wrap every instruction, so specialized
            // CACHE payloads must be cleared first.
            let want_instrumented =
                events & (MonitoringEvent::Line.mask() | MonitoringEvent::Instruction.mask()) != 0
                    || opcode_event_is_active(base_op, oparg, events);
            let restore_base = if want_instrumented {
                !op.is_instrumented() && u8::from(base_op) != u8::from(op)
            } else {
                op.is_instrumented()
            };
            if restore_base {
                unsafe {
                    code.code.instructions.replace_op(i, base_op);
                }
                for c in 1..=caches {
                    if i + c < len {
                        unsafe {
                            code.code.instructions.write_cache_u16(i + c, 0);
                        }
                    }
                }
            }
            i += 1 + caches;
        }
    }

    if events == 0 {
        *monitoring_data = None;
        return;
    }

    // === Phase 4-6: Re-instrument (innermost first) ===

    // Ensure monitoring data exists
    if monitoring_data.is_none() {
        *monitoring_data = Some(CoMonitoringData {
            line_opcodes: vec![0u8; len],
            per_instruction_opcodes: vec![0u8; len],
        });
    }
    let data = monitoring_data.as_mut().unwrap();
    // Resize if code length changed (shouldn't happen, but be safe)
    data.line_opcodes.resize(len, 0);
    data.per_instruction_opcodes.resize(len, 0);

    // Find _co_firsttraceable: index of first RESUME instruction
    let mut first_traceable = None;
    for_each_instruction(code, |i, op, _| {
        if first_traceable.is_none()
            && matches!(
                op,
                Instruction::Resume { .. }
                    | Instruction::ResumeCheck
                    | Instruction::InstrumentedResume
            )
        {
            first_traceable = Some(i);
        }
    });
    let first_traceable = first_traceable.unwrap_or(0);

    // Phase 4: Place regular INSTRUMENTED_* opcodes whose event is active.
    // Walk by cache_entries so specialized CACHE payloads are not decoded as opcodes.
    {
        let mut i = 0;
        while i < len {
            let op = code.code.instructions.read_op(i);
            let oparg = code.code.instructions.read_arg(i).as_u8();
            let caches = op.deoptimize().cache_entries();
            if (events & (MonitoringEvent::Line.mask() | MonitoringEvent::Instruction.mask()) != 0
                || opcode_event_is_active(op, oparg, events))
                && let Some(instrumented) = op.to_instrumented()
            {
                unsafe {
                    code.code.instructions.replace_op(i, instrumented);
                }
            }
            i += 1 + caches;
        }
    }

    // Phase 5: Place INSTRUMENTED_INSTRUCTION when Instruction is active
    // LINE/INSTRUCTION restore specialized opcodes first, so CACHE slots are valid.
    if events & MonitoringEvent::Instruction.mask() != 0 {
        for i in first_traceable..len {
            let op = code.code.instructions[i].op;
            // Skip ExtendedArg
            if matches!(op, Instruction::ExtendedArg) {
                continue;
            }
            // Excluded: RESUME, END_FOR, CACHE (and their instrumented variants)
            let base = op.to_base().unwrap_or(op);
            if matches!(
                base,
                Instruction::Resume { .. } | Instruction::EndFor | Instruction::Cache
            ) {
                continue;
            }
            // Store current opcode (may already be INSTRUMENTED_*) and replace
            data.per_instruction_opcodes[i] = u8::from(op);
            unsafe {
                code.code
                    .instructions
                    .replace_op(i, Instruction::InstrumentedInstruction);
            }
        }
    }

    // Phase 6: Place INSTRUMENTED_LINE when Line is active
    // Mirrors CPython's initialize_lines: first determine which positions
    // are line starts, then mark branch/jump targets, then place opcodes.
    if events & MonitoringEvent::Line.mask() != 0 {
        // is_line_start[i] = true if position i should have INSTRUMENTED_LINE
        let mut is_line_start = vec![false; len];
        let line_locations = rustpython_compiler_core::marshal::linetable_to_locations(
            &code.code.linetable,
            code.code.first_line_number.map_or(-1, |x| x.get() as i32),
            len,
        );

        // Build NO_LOCATION mask from linetable
        let no_loc_mask = bytecode::build_no_location_mask(&code.code.linetable, len);

        // First pass: mark positions where the source line changes
        let mut prev_line: Option<u32> = None;
        for (i, unit) in code
            .code
            .instructions
            .iter()
            .enumerate()
            .take(len)
            .skip(first_traceable)
        {
            let op = unit.op;
            let base = op.to_base().unwrap_or(op);
            if matches!(base, Instruction::ExtendedArg) {
                continue;
            }
            // Excluded opcodes
            if matches!(
                base,
                Instruction::Resume { .. }
                    | Instruction::EndFor
                    | Instruction::EndSend
                    | Instruction::PopIter
                    | Instruction::EndAsyncFor
                    | Instruction::Cache
            ) {
                continue;
            }
            // Skip NO_LOCATION instructions
            if no_loc_mask.get(i).copied().unwrap_or(false) {
                continue;
            }
            if let Some((loc, _)) = line_locations.get(i) {
                let line = loc.line.get() as u32;
                let is_new = prev_line != Some(line);
                prev_line = Some(line);
                if is_new && line > 0 {
                    is_line_start[i] = true;
                }
            }
        }

        // Second pass: mark branch/jump targets as line starts.
        // Every jump/branch target must be a line start, even if on the
        // same source line as the preceding instruction. Critical for loops
        // (JUMP_BACKWARD → FOR_ITER).
        let mut arg_state = bytecode::OpArgState::default();
        let mut instr_idx = first_traceable;
        for unit in code.code.instructions[first_traceable..len].iter().copied() {
            let (op, arg) = arg_state.get(unit);
            let base = op.to_base().unwrap_or(op);

            if matches!(base, Instruction::ExtendedArg) || matches!(base, Instruction::Cache) {
                instr_idx += 1;
                continue;
            }

            let caches = base.cache_entries();
            let after_caches = instr_idx + 1 + caches;
            let delta = u32::from(arg) as usize;

            let target: Option<usize> = match base {
                // Forward relative jumps
                Instruction::PopJumpIfFalse { .. }
                | Instruction::PopJumpIfTrue { .. }
                | Instruction::PopJumpIfNone { .. }
                | Instruction::PopJumpIfNotNone { .. }
                | Instruction::JumpForward { .. } => Some(after_caches + delta),
                // Backward relative jumps
                Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } => {
                    Some(after_caches.wrapping_sub(delta))
                }
                Instruction::ForIter { .. } | Instruction::Send { .. } => {
                    // Skip over END_FOR/END_SEND
                    Some(after_caches + delta + 1)
                }
                _ => None,
            };

            if let Some(target_idx) = target
                && target_idx < len
                && !is_line_start[target_idx]
                && !no_loc_mask.get(target_idx).copied().unwrap_or(false)
            {
                let target_op = code.code.instructions[target_idx].op;
                let target_base = target_op.to_base().unwrap_or(target_op);
                // Skip synthetic cleanup targets.
                if matches!(target_base, Instruction::PopIter) {
                    instr_idx += 1;
                    continue;
                }
                if let Some((loc, _)) = line_locations.get(target_idx)
                    && loc.line.get() > 0
                {
                    is_line_start[target_idx] = true;
                }
            }
            instr_idx += 1;
        }

        // Third pass: mark exception handler targets as line starts.
        for entry in bytecode::decode_exception_table(&code.code.exceptiontable) {
            let target_idx = entry.target as usize;
            if target_idx < len
                && !is_line_start[target_idx]
                && !no_loc_mask.get(target_idx).copied().unwrap_or(false)
            {
                let target_op = code.code.instructions.read_op(target_idx);
                let target_base = target_op.to_base().unwrap_or(target_op);
                if !matches!(target_base, Instruction::PopIter)
                    && let Some((loc, _)) = line_locations.get(target_idx)
                    && loc.line.get() > 0
                {
                    is_line_start[target_idx] = true;
                }
            }
        }

        // Fourth pass: actually place INSTRUMENTED_LINE at all marked positions
        for (i, marked) in is_line_start
            .iter()
            .copied()
            .enumerate()
            .take(len)
            .skip(first_traceable)
        {
            if marked {
                let op = code.code.instructions.read_op(i);
                data.line_opcodes[i] = u8::from(op);
                unsafe {
                    code.code
                        .instructions
                        .replace_op(i, Instruction::InstrumentedLine);
                }
            }
        }
    }
}

/// Update the global monitoring_events atomic mask from current state.
fn update_events_mask(vm: &VirtualMachine, state: &MonitoringState) {
    let events = state.combined_events();
    vm.state.monitoring_events.store(events);
    let new_ver = vm
        .state
        .instrumentation_version
        .fetch_add(1, Ordering::Release)
        + 1;
    // Eagerly re-instrument all frames on the current thread's stack so that
    // code objects already past their RESUME pick up the new event set.
    // Each code object gets only the events that apply to it (global + its
    // own local events), preventing e.g. INSTRUCTION from being applied to
    // unrelated code objects.
    // Re-instrument all frames on the current thread's stack, including
    // data stack frames that have no FrameObject.
    {
        let mut cur = crate::vm::thread::get_current_frame();
        while !cur.is_null() {
            let iframe_ref = unsafe { &*cur };
            let code = iframe_ref.code();
            let code_ver = code.instrumentation_version.load(Ordering::Acquire);
            if code_ver != new_ver {
                let code_events = state.events_for_code(code.get_id());
                instrument_code(code, code_events);
                code.instrumentation_version
                    .store(new_ver, Ordering::Release);
            }
            cur = iframe_ref.previous();
        }
    }
}

fn use_tool_id(tool_id: i32, name: &str, vm: &VirtualMachine) -> PyResult<()> {
    let tool = check_valid_tool(tool_id, vm)?;
    let mut state = vm.state.monitoring.lock();
    if state.tool_names[tool].is_some() {
        return Err(vm.new_value_error(format!("tool {tool_id} is already in use")));
    }
    state.tool_names[tool] = Some(name.to_owned());
    Ok(())
}

fn clear_tool_id(tool_id: i32, vm: &VirtualMachine) -> PyResult<()> {
    let tool = check_valid_tool(tool_id, vm)?;
    let mut state = vm.state.monitoring.lock();
    if state.tool_names[tool].is_some() {
        state.global_events[tool] = 0;
        state
            .local_events
            .retain(|(local_tool, _), _| *local_tool != tool);
        state.callbacks.retain(|(cb_tool, _), _| *cb_tool != tool);
        state.disabled.retain(|&(_, _, t)| t != tool);
    }
    update_events_mask(vm, &state);
    Ok(())
}

fn free_tool_id(tool_id: i32, vm: &VirtualMachine) -> PyResult<()> {
    let tool = check_valid_tool(tool_id, vm)?;
    let mut state = vm.state.monitoring.lock();
    if state.tool_names[tool].is_some() {
        state.global_events[tool] = 0;
        state
            .local_events
            .retain(|(local_tool, _), _| *local_tool != tool);
        state.callbacks.retain(|(cb_tool, _), _| *cb_tool != tool);
        state.disabled.retain(|&(_, _, t)| t != tool);
        state.tool_names[tool] = None;
    }
    update_events_mask(vm, &state);
    Ok(())
}

fn get_tool(tool_id: i32, vm: &VirtualMachine) -> PyResult<Option<String>> {
    let tool = check_valid_tool(tool_id, vm)?;
    let state = vm.state.monitoring.lock();
    Ok(state.tool_names[tool].clone())
}

fn register_callback(
    tool_id: i32,
    event: i32,
    func: PyObjectRef,
    vm: &VirtualMachine,
) -> PyResult<PyObjectRef> {
    let tool = check_valid_tool(tool_id, vm)?;
    let event_id = parse_single_event(event, vm)?;

    vm.audit("sys.monitoring.register_callback", || (func.clone(),))?;

    let mut state = vm.state.monitoring.lock();
    let prev = state
        .callbacks
        .remove(&(tool, event_id))
        .unwrap_or_else(|| vm.ctx.none());
    let branch_id = MonitoringEvent::Branch as usize;
    let branch_left_id = MonitoringEvent::BranchLeft as usize;
    let branch_right_id = MonitoringEvent::BranchRight as usize;
    if !vm.is_none(&func) {
        state.callbacks.insert((tool, event_id), func.clone());
        // BRANCH is a composite event: also register for BRANCH_LEFT/RIGHT
        if event_id == branch_id {
            state.callbacks.insert((tool, branch_left_id), func.clone());
            state.callbacks.insert((tool, branch_right_id), func);
        }
    } else {
        // Also clear BRANCH_LEFT/RIGHT when clearing BRANCH
        if event_id == branch_id {
            state.callbacks.remove(&(tool, branch_left_id));
            state.callbacks.remove(&(tool, branch_right_id));
        }
    }
    Ok(prev)
}

fn get_events(tool_id: i32, vm: &VirtualMachine) -> PyResult<u32> {
    let tool = check_valid_tool(tool_id, vm)?;
    let state = vm.state.monitoring.lock();
    Ok(state.global_events[tool])
}

fn set_events(tool_id: i32, event_set: i32, vm: &VirtualMachine) -> PyResult<()> {
    let tool = check_valid_tool(tool_id, vm)?;
    check_tool_in_use(tool, vm)?;
    let normalized = normalize_event_set(event_set, false, vm)?;
    let mut state = vm.state.monitoring.lock();
    state.global_events[tool] = normalized;
    update_events_mask(vm, &state);
    Ok(())
}

fn get_local_events(tool_id: i32, code: &PyObject, vm: &VirtualMachine) -> PyResult<u32> {
    if code.downcast_ref::<PyCode>().is_none() {
        return Err(vm.new_type_error("code must be a code object"));
    }
    let tool = check_valid_tool(tool_id, vm)?;
    let code_id = code.get_id();
    let state = vm.state.monitoring.lock();
    Ok(state
        .local_events
        .get(&(tool, code_id))
        .copied()
        .unwrap_or(0))
}

fn set_local_events(
    tool_id: i32,
    code: &PyObject,
    event_set: i32,
    vm: &VirtualMachine,
) -> PyResult<()> {
    if code.downcast_ref::<PyCode>().is_none() {
        return Err(vm.new_type_error("code must be a code object"));
    }
    let tool = check_valid_tool(tool_id, vm)?;
    check_tool_in_use(tool, vm)?;
    let normalized = normalize_event_set(event_set, true, vm)?;
    let code_id = code.get_id();
    let mut state = vm.state.monitoring.lock();
    if normalized == 0 {
        state.local_events.remove(&(tool, code_id));
    } else {
        state.local_events.insert((tool, code_id), normalized);
    }
    update_events_mask(vm, &state);
    Ok(())
}

fn restart_events(vm: &VirtualMachine) {
    let mut state = vm.state.monitoring.lock();
    state.disabled.clear();
}

fn all_events(vm: &VirtualMachine) -> PyResult<PyDictRef> {
    // Collect data under the lock, then release before calling into Python VM.
    let masks: Vec<(&str, u8)> = {
        let state = vm.state.monitoring.lock();
        EVENT_NAMES
            .iter()
            .take(UNGROUPED_EVENTS_COUNT)
            .enumerate()
            .filter_map(|(event_id, event_name)| {
                let event_bit = 1u32 << event_id;
                let mut tools_mask = 0u8;
                for (i, tool) in state.global_events.iter().enumerate().take(TOOL_LIMIT) {
                    if (tool & event_bit) != 0 {
                        tools_mask |= 1 << i;
                    }
                }
                if tools_mask != 0 {
                    Some((*event_name, tools_mask))
                } else {
                    None
                }
            })
            .collect()
    };
    let all_events = vm.ctx.new_dict();
    for (name, mask) in masks {
        all_events.set_item(name, vm.ctx.new_int(mask).into(), vm)?;
    }
    Ok(all_events)
}

// Event dispatch

use core::cell::Cell;

thread_local! {
    /// Re-entrancy guard: prevents monitoring callbacks from triggering
    /// additional monitoring events (which would cause infinite recursion).
    static FIRING: Cell<bool> = const { Cell::new(false) };

    /// Tracks whether a RERAISE event has been fired since the last
    /// EXCEPTION_HANDLED. Used to suppress duplicate RERAISE from
    /// cleanup handlers that chain through multiple exception table entries.
    static RERAISE_PENDING: Cell<bool> = const { Cell::new(false) };
}

/// Fire an event for all tools that have the event bit set.
/// `cb_extra` contains the callback arguments after the code object.
fn fire(
    vm: &VirtualMachine,
    event: MonitoringEvent,
    code: &Py<PyCode>,
    offset: u32,
    cb_extra: &[PyObjectRef],
) -> PyResult<()> {
    // Prevent recursive event firing
    if vm.tracing_is_suppressed() || FIRING.with(|f| f.get()) {
        return Ok(());
    }

    let event_id = event as usize;
    let event_mask = event.mask();
    let code_id = code.get_id();

    // C_RETURN and C_RAISE are implicitly enabled when CALL is set.
    let check_bit = if matches!(event, MonitoringEvent::CReturn | MonitoringEvent::CRaise) {
        event_mask | MonitoringEvent::Call.mask()
    } else {
        event_mask
    };

    // Collect callbacks and snapshot the DISABLE sentinel under a single lock.
    let (callbacks, disable_sentinel): (Vec<(usize, PyObjectRef)>, Option<PyObjectRef>) = {
        let state = vm.state.monitoring.lock();
        let mut cbs = Vec::new();
        for tool in 0..TOOL_LIMIT {
            let global = state.global_events[tool];
            let local = state
                .local_events
                .get(&(tool, code_id))
                .copied()
                .unwrap_or(0);
            if ((global | local) & check_bit) == 0 {
                continue;
            }
            if state.disabled.contains(&(code_id, offset as usize, tool)) {
                continue;
            }
            if let Some(cb) = state.callbacks.get(&(tool, event_id)) {
                cbs.push((tool, cb.clone()));
            }
        }
        (cbs, state.disable.clone())
    };

    if callbacks.is_empty() {
        return Ok(());
    }

    let mut args_vec = Vec::with_capacity(1 + cb_extra.len());
    args_vec.push(code.to_owned().into());
    args_vec.extend_from_slice(cb_extra);
    let args = FuncArgs::from(args_vec);

    FIRING.with(|f| f.set(true));
    vm.enter_tracing();
    let old_what = vm.what_event.replace(Some(event));
    let result = (|| {
        for (tool, cb) in callbacks {
            let result = cb.call(args.clone(), vm)?;
            if disable_sentinel.as_ref().is_some_and(|d| result.is(d)) {
                // Only local events (event_id < LOCAL_EVENTS_COUNT) can be disabled.
                // Non-local events (RAISE, EXCEPTION_HANDLED, PY_UNWIND, etc.)
                // cannot be disabled per code object.
                if event_id >= LOCAL_EVENTS_COUNT {
                    // Remove the callback.
                    let mut state = vm.state.monitoring.lock();
                    state.callbacks.remove(&(tool, event_id));
                    return Err(vm.new_value_error(format!(
                        "Cannot disable {} events. Callback removed.",
                        EVENT_NAMES[event_id]
                    )));
                }
                let mut state = vm.state.monitoring.lock();
                state.disabled.insert((code_id, offset as usize, tool));
            }
        }
        Ok(())
    })();
    vm.what_event.set(old_what);
    vm.leave_tracing();
    FIRING.with(|f| f.set(false));
    result
}

// Public dispatch functions (called from frame.rs)

pub(crate) fn fire_py_start(vm: &VirtualMachine, code: &Py<PyCode>, offset: u32) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::PyStart,
        code,
        offset,
        &[vm.ctx.new_int(offset).into()],
    )
}

pub(crate) fn fire_py_resume(vm: &VirtualMachine, code: &Py<PyCode>, offset: u32) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::PyResume,
        code,
        offset,
        &[vm.ctx.new_int(offset).into()],
    )
}

pub(crate) fn fire_py_return(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    retval: &PyObject,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::PyReturn,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), retval.to_owned()],
    )
}

pub(crate) fn fire_py_yield(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    retval: &PyObject,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::PyYield,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), retval.to_owned()],
    )
}

pub(crate) fn fire_call(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    callable: &PyObject,
    arg0: PyObjectRef,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::Call,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), callable.to_owned(), arg0],
    )
}

pub(crate) fn fire_c_return(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    callable: &PyObject,
    arg0: PyObjectRef,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::CReturn,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), callable.to_owned(), arg0],
    )
}

pub(crate) fn fire_c_raise(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    callable: &PyObject,
    arg0: PyObjectRef,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::CRaise,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), callable.to_owned(), arg0],
    )
}

pub(crate) fn fire_line(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    line: u32,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::Line,
        code,
        offset,
        &[vm.ctx.new_int(line).into()],
    )
}

pub(crate) fn fire_instruction(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::Instruction,
        code,
        offset,
        &[vm.ctx.new_int(offset).into()],
    )
}

pub(crate) fn fire_raise(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    exception: &PyObject,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::Raise,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), exception.to_owned()],
    )
}

/// Only fires if no RERAISE has been fired since the last EXCEPTION_HANDLED,
/// preventing duplicate events from chained cleanup handlers.
pub(crate) fn fire_reraise(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    exception: &PyObject,
) -> PyResult<()> {
    if RERAISE_PENDING.with(|f| f.get()) {
        return Ok(());
    }
    RERAISE_PENDING.with(|f| f.set(true));
    let result = fire(
        vm,
        MonitoringEvent::Reraise,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), exception.to_owned()],
    );
    if result.is_err() {
        RERAISE_PENDING.with(|f| f.set(false));
    }
    result
}

pub(crate) fn fire_exception_handled(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    exception: &PyObject,
) -> PyResult<()> {
    RERAISE_PENDING.with(|f| f.set(false));
    fire(
        vm,
        MonitoringEvent::ExceptionHandled,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), exception.to_owned()],
    )
}

pub(crate) fn fire_py_unwind(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    exception: &PyObject,
) -> PyResult<()> {
    RERAISE_PENDING.with(|f| f.set(false));
    fire(
        vm,
        MonitoringEvent::PyUnwind,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), exception.to_owned()],
    )
}

pub(crate) fn fire_py_throw(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    exception: &PyObject,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::PyThrow,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), exception.to_owned()],
    )
}

/// If `value` is already a `StopIteration`, pass it directly; otherwise wrap
/// it in a new `StopIteration(value)` — matching `PyMonitoring_FireStopIterationEvent`.
pub(crate) fn fire_stop_iteration(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    value: &PyObject,
) -> PyResult<()> {
    let exc: PyObjectRef = if value.fast_isinstance(vm.ctx.exceptions.stop_iteration) {
        value.to_owned()
    } else {
        vm.new_stop_iteration(Some(value.to_owned())).into()
    };
    fire(
        vm,
        MonitoringEvent::StopIteration,
        code,
        offset,
        &[vm.ctx.new_int(offset).into(), exc],
    )
}

pub(crate) fn fire_jump(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    destination: u32,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::Jump,
        code,
        offset,
        &[
            vm.ctx.new_int(offset).into(),
            vm.ctx.new_int(destination).into(),
        ],
    )
}

pub(crate) fn fire_branch_left(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    destination: u32,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::BranchLeft,
        code,
        offset,
        &[
            vm.ctx.new_int(offset).into(),
            vm.ctx.new_int(destination).into(),
        ],
    )
}

pub(crate) fn fire_branch_right(
    vm: &VirtualMachine,
    code: &Py<PyCode>,
    offset: u32,
    destination: u32,
) -> PyResult<()> {
    fire(
        vm,
        MonitoringEvent::BranchRight,
        code,
        offset,
        &[
            vm.ctx.new_int(offset).into(),
            vm.ctx.new_int(destination).into(),
        ],
    )
}

#[pymodule(sub, name = "sys.monitoring")]
pub(super) mod sys_monitoring {
    use super::*;

    #[pyclass(no_attr, module = "sys.monitoring", name = "_Sentinel")]
    #[derive(Debug, PyPayload)]
    pub(super) struct MonitoringSentinel;

    #[pyclass]
    impl MonitoringSentinel {}

    #[pyattr(name = "DEBUGGER_ID")]
    const DEBUGGER_ID: u8 = 0;
    #[pyattr(name = "COVERAGE_ID")]
    const COVERAGE_ID: u8 = 1;
    #[pyattr(name = "PROFILER_ID")]
    const PROFILER_ID: u8 = 2;
    #[pyattr(name = "OPTIMIZER_ID")]
    const OPTIMIZER_ID: u8 = 5;

    #[pyattr(once, name = "DISABLE")]
    fn disable(vm: &VirtualMachine) -> PyObjectRef {
        super::get_disable(vm)
    }

    #[pyattr(once, name = "MISSING")]
    fn missing(vm: &VirtualMachine) -> PyObjectRef {
        super::get_missing(vm)
    }

    #[pyattr(once)]
    fn events(vm: &VirtualMachine) -> PyRef<PyNamespace> {
        let events = PyNamespace::default().into_ref(&vm.ctx);
        for (event_id, event_name) in EVENT_NAMES.iter().enumerate() {
            events
                .as_object()
                .set_attr(*event_name, vm.ctx.new_int(1u32 << event_id), vm)
                .expect("setting sys.monitoring.events attribute should not fail");
        }
        events
            .as_object()
            .set_attr("NO_EVENTS", vm.ctx.new_int(0), vm)
            .expect("setting sys.monitoring.events.NO_EVENTS should not fail");
        events
    }

    #[pyfunction]
    fn use_tool_id(tool_id: i32, name: PyUtf8StrRef, vm: &VirtualMachine) -> PyResult<()> {
        super::use_tool_id(tool_id, name.as_str(), vm)
    }

    #[pyfunction]
    fn clear_tool_id(tool_id: i32, vm: &VirtualMachine) -> PyResult<()> {
        super::clear_tool_id(tool_id, vm)
    }

    #[pyfunction]
    fn free_tool_id(tool_id: i32, vm: &VirtualMachine) -> PyResult<()> {
        super::free_tool_id(tool_id, vm)
    }

    #[pyfunction]
    fn get_tool(tool_id: i32, vm: &VirtualMachine) -> PyResult<Option<String>> {
        super::get_tool(tool_id, vm)
    }

    #[pyfunction]
    fn register_callback(
        tool_id: i32,
        event: i32,
        func: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyObjectRef> {
        super::register_callback(tool_id, event, func, vm)
    }

    #[pyfunction]
    fn get_events(tool_id: i32, vm: &VirtualMachine) -> PyResult<u32> {
        super::get_events(tool_id, vm)
    }

    #[pyfunction]
    fn set_events(tool_id: i32, event_set: i32, vm: &VirtualMachine) -> PyResult<()> {
        super::set_events(tool_id, event_set, vm)
    }

    #[pyfunction]
    fn get_local_events(tool_id: i32, code: PyObjectRef, vm: &VirtualMachine) -> PyResult<u32> {
        super::get_local_events(tool_id, &code, vm)
    }

    #[pyfunction]
    fn set_local_events(
        tool_id: i32,
        code: PyObjectRef,
        event_set: i32,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        super::set_local_events(tool_id, &code, event_set, vm)
    }

    #[pyfunction]
    fn restart_events(vm: &VirtualMachine) {
        super::restart_events(vm)
    }

    #[pyfunction]
    fn _all_events(vm: &VirtualMachine) -> PyResult<PyDictRef> {
        super::all_events(vm)
    }
}