use alloy_primitives::{Address, Bytes};
use revm::{
    Inspector,
    bytecode::opcode::ADDRESS,
    interpreter::{
        InstructionResult, Interpreter, InterpreterAction,
        interpreter_types::{Jumps, LoopControl},
    },
};

/// An inspector that enforces certain rules during script execution.
///
/// Currently, it only warns if the `ADDRESS` opcode is used within the script's main contract.
#[derive(Clone, Debug, Default)]
pub struct ScriptExecutionInspector {
    /// The address of the script contract being executed.
    pub script_address: Address,
}

impl<CTX> Inspector<CTX> for ScriptExecutionInspector {
    fn step(&mut self, interpreter: &mut Interpreter, _ecx: &mut CTX) {
        // Check if both target and bytecode address are the same as script contract address
        // (allow calling external libraries when bytecode address is different).
        if interpreter.bytecode.opcode() == ADDRESS
            && interpreter.input.target_address == self.script_address
            && interpreter.input.bytecode_address == Some(self.script_address)
        {
            interpreter.bytecode.set_action(InterpreterAction::new_return(
                InstructionResult::Revert,
                Bytes::from("Usage of `address(this)` detected in script contract. Script contracts are ephemeral and their addresses should not be relied upon."),
                interpreter.gas,
            ));
        }
    }
}
