use crate::bytes;
use core::fmt;
use object::{Bytes, LittleEndian, U32};

/// Information about trap.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct TrapInformation {
    /// The offset of the trapping instruction in native code.
    ///
    /// This is relative to the beginning of the function.
    pub code_offset: u32,

    /// Code of the trap.
    pub trap_code: CompiledTrap,
}

/// Possible traps that can be compiled into WebAssembly modules.
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum CompiledTrap {
    /// A normal trap, expected to possibly be hit at runtime.
    Normal(Trap),
    /// An internal assertion in the compiled code itself.
    InternalAssert,
    /// The GC heap was detected as being corrupt.
    GcHeapCorrupt,
}

impl CompiledTrap {
    /// Encodes this as a byte.
    pub fn as_u8(&self) -> u8 {
        match self {
            CompiledTrap::Normal(trap) => {
                let ret = *trap as u8;
                debug_assert_ne!(ret, CompiledTrap::InternalAssert.as_u8());
                debug_assert_ne!(ret, CompiledTrap::GcHeapCorrupt.as_u8());
                ret
            }
            CompiledTrap::InternalAssert => 0xFF,
            CompiledTrap::GcHeapCorrupt => 0xFE,
        }
    }

    /// Decodes a byte as a trap.
    pub fn from_u8(byte: u8) -> Option<CompiledTrap> {
        if byte == 0xFF {
            Some(CompiledTrap::InternalAssert)
        } else if byte == 0xFE {
            Some(CompiledTrap::GcHeapCorrupt)
        } else {
            Trap::from_u8(byte).map(CompiledTrap::Normal)
        }
    }
}

impl From<Trap> for CompiledTrap {
    fn from(trap: Trap) -> CompiledTrap {
        CompiledTrap::Normal(trap)
    }
}

macro_rules! generate_trap_type {
    (pub enum Trap {
        $(
            $(#[$doc:meta])*
            $name:ident = $msg:tt,
        )*
    }) => {
        #[non_exhaustive]
        #[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
        #[expect(missing_docs, reason = "self-describing variants")]
        pub enum Trap {
            $(
                $(#[$doc])*
                $name,
            )*
        }

        impl Trap {
            /// Converts a byte back into a `Trap` if its in-bounds
            pub fn from_u8(byte: u8) -> Option<Trap> {
                $(
                    if byte == Trap::$name as u8 {
                        return Some(Trap::$name);
                    }
                )*
                None
            }
        }

        impl fmt::Display for Trap {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                let desc = match self {
                    $(Self::$name => $msg,)*
                };
                write!(f, "wasm trap: {desc}")
            }
        }
    }
}

// The code can be accessed from the c-api, where the possible values are
// translated into enum values defined there:
//
// *  the const assertions in c-api/src/trap.rs, and
// * `wasmtime_trap_code_enum` in c-api/include/wasmtime/trap.h.
//
// These need to be kept in sync.
generate_trap_type! {
    pub enum Trap {
        /// The current stack space was exhausted.
        StackOverflow = "call stack exhausted",

        /// An out-of-bounds memory access.
        MemoryOutOfBounds = "out of bounds memory access",

        /// A wasm atomic operation was presented with a not-naturally-aligned linear-memory address.
        HeapMisaligned = "unaligned atomic",

        /// An out-of-bounds access to a table.
        TableOutOfBounds = "undefined element: out of bounds table access",

        /// Indirect call to a null table entry.
        IndirectCallToNull = "uninitialized element",

        /// Signature mismatch on indirect call.
        BadSignature = "indirect call type mismatch",

        /// An integer arithmetic operation caused an overflow.
        IntegerOverflow = "integer overflow",

        /// An integer division by zero.
        IntegerDivisionByZero = "integer divide by zero",

        /// Failed float-to-int conversion.
        BadConversionToInteger = "invalid conversion to integer",

        /// Code that was supposed to have been unreachable was reached.
        UnreachableCodeReached = "wasm `unreachable` instruction executed",

        /// Execution has potentially run too long and may be interrupted.
        Interrupt = "interrupt",

        /// When wasm code is configured to consume fuel and it runs out of fuel
        /// then this trap will be raised.
        OutOfFuel = "all fuel consumed by WebAssembly",

        /// Used to indicate that a trap was raised by atomic wait operations on non shared memory.
        AtomicWaitNonSharedMemory = "atomic wait on non-shared memory",

        /// Call to a null reference.
        NullReference = "null reference",

        /// Attempt to access beyond the bounds of an array.
        ArrayOutOfBounds = "out of bounds array access",

        /// Attempted an allocation that was too large to succeed.
        AllocationTooLarge = "allocation size too large",

        /// Attempted to cast a reference to a type that it is not an instance of.
        CastFailure = "cast failure",

        /// When the `component-model` feature is enabled this trap represents a
        /// scenario where one component tried to call another component but it
        /// would have violated the reentrance rules of the component model,
        /// triggering a trap instead.
        CannotEnterComponent = "cannot enter component instance",

        /// Async-lifted export failed to produce a result by calling `task.return`
        /// before returning `STATUS_DONE` and/or after all host tasks completed.
        NoAsyncResult = "async-lifted export failed to produce a result",

        /// We are suspending to a tag for which there is no active handler.
        UnhandledTag = "unhandled tag",

        /// Attempt to resume a continuation twice.
        ContinuationAlreadyConsumed = "continuation already consumed",

        /// A Pulley opcode was executed at runtime when the opcode was disabled at
        /// compile time.
        DisabledOpcode = "pulley opcode disabled at compile time was executed",

        /// Async event loop deadlocked; i.e. it cannot make further progress given
        /// that all host tasks have completed and any/all host-owned stream/future
        /// handles have been dropped.
        AsyncDeadlock = "deadlock detected: event loop cannot make further progress",

        /// When the `component-model` feature is enabled this trap represents a
        /// scenario where a component instance tried to call an import or intrinsic
        /// when it wasn't allowed to, e.g. from a post-return function.
        CannotLeaveComponent = "cannot leave component instance",

        /// A synchronous task attempted to make a potentially blocking call prior
        /// to returning.
        CannotBlockSyncTask = "cannot block a synchronous task before returning",

        /// A component tried to lift a `char` with an invalid bit pattern.
        InvalidChar = "invalid `char` bit pattern",

        /// Debug assertion generated for a fused adapter regarding the expected
        /// completion of a string encoding operation.
        DebugAssertStringEncodingFinished = "should have finished string encoding",

        /// Debug assertion generated for a fused adapter regarding a string
        /// encoding operation.
        DebugAssertEqualCodeUnits = "code units should be equal",

        /// Debug assertion generated for a fused adapter regarding the alignment of
        /// a pointer.
        DebugAssertPointerAligned = "pointer should be aligned",

        /// Debug assertion generated for a fused adapter regarding the upper bits
        /// of a 64-bit value.
        DebugAssertUpperBitsUnset = "upper bits should be unset",

        /// A component tried to lift or lower a string past the end of its memory.
        StringOutOfBounds = "string content out-of-bounds",

        /// A component tried to lift or lower a list past the end of its memory.
        ListOutOfBounds = "list content out-of-bounds",

        /// A component used an invalid discriminant when lowering a variant value.
        InvalidDiscriminant = "invalid variant discriminant",

        /// A component passed an unaligned pointer when lifting or lowering a
        /// value.
        UnalignedPointer = "unaligned pointer",

        /// `task.cancel` was called by a task which has not been cancelled.
        TaskCancelNotCancelled = "`task.cancel` called by task which has not been cancelled",

        /// `task.return` or `task.cancel` was called more than once for the
        /// current task.
        TaskCancelOrReturnTwice = "`task.return` or `task.cancel` called more than once for current task",

        /// `subtask.cancel` was called after terminal status was already
        /// delivered.
        SubtaskCancelAfterTerminal = "`subtask.cancel` called after terminal status delivered",

        /// Invalid `task.return` signature and/or options for the current task.
        TaskReturnInvalid = "invalid `task.return` signature and/or options for current task",

        /// Cannot drop waitable set with waiters in it.
        WaitableSetDropHasWaiters = "cannot drop waitable set with waiters",

        /// Cannot drop a subtask which has not yet resolved.
        SubtaskDropNotResolved = "cannot drop a subtask which has not yet resolved",

        /// Start function does not match the expected type.
        ThreadNewIndirectInvalidType = "start function does not match expected type (currently only `(i32) -> ()` is supported)",

        /// The start function index points to an uninitialized function.
        ThreadNewIndirectUninitialized = "the start function index points to an uninitialized function",

        /// Backpressure-related intrinsics overflowed the built-in 16-bit
        /// counter.
        BackpressureOverflow = "backpressure counter overflow",

        /// Invalid code returned from `callback` of `async`-lifted function.
        UnsupportedCallbackCode = "unsupported callback code",

        /// Cannot resume a thread which is not suspended.
        CannotResumeThread = "cannot resume thread which is not suspended",

        /// Cannot issue a read/write on a future/stream while there is a
        /// pending operation already.
        ConcurrentFutureStreamOp = "cannot have concurrent operations active on a future/stream",

        /// A reference count (for e.g. an `error-context`) overflowed.
        ReferenceCountOverflow = "reference count overflow",

        /// A read/write on a stream must be <2**28 items.
        StreamOpTooBig = "stream read/write count too large",

        /// The guest either attempted to add a waitable to a waitable set while
        /// it was being used in a synchronous operation or tried to use it in a
        /// synchronous operation while it was added to a waitable set.
        WaitableSyncAndAsync = "waitable cannot be used synchronously while added to a waitable set",

        /// An exception propagated out of a component without being caught.
        UncaughtException = "uncaught exception propagated out of component",

        // if adding a variant here be sure to update `trap.rs` and `trap.h` as
        // mentioned above
    }
}

impl core::error::Error for Trap {}

/// Number of trap entries packed into one block of the trap section.
///
/// See `TrapEncodingBuilder` in `crate::compile` for the full section format.
/// Chosen as a balance between the fixed-width index overhead per block (8
/// bytes, amortized across entries) and the amount of linear decoding required
/// to look up a single pc within a block.
pub(crate) const TRAP_BLOCK_SIZE: usize = 128;

/// Decodes the provided trap information section and attempts to find the trap
/// code corresponding to the `offset` specified.
///
/// The `section` provided is expected to have been built by
/// `TrapEncodingBuilder` in `crate::compile`, whose documentation describes
/// the format decoded here. Additionally the `offset` should be a relative
/// offset within the text section of the compilation image.
pub fn lookup_trap_code(section: &[u8], offset: usize) -> Option<CompiledTrap> {
    let section = parse(section)?;
    let offset = u32::try_from(offset).ok()?;

    // Find the last block whose first pc is `<= offset`; only that block can
    // contain `offset`. Note that this is a precise search because trap pcs
    // should always be precise as well as our metadata about them, which means
    // we expect an exact match to correspond to a trap opcode.
    let block = section
        .block_index
        .partition_point(|[first_offset, _]| first_offset.get(LittleEndian) <= offset)
        .checked_sub(1)?;

    for (pc, byte) in section.block_entries(block)? {
        if pc == offset {
            let trap = CompiledTrap::from_u8(byte);
            debug_assert!(trap.is_some(), "missing mapping for {byte}");
            return trap;
        }
        if pc > offset {
            break;
        }
    }
    None
}

/// A parsed view of the trap section.
///
/// The fields here correspond to the pieces of the section layout described
/// on `TrapEncodingBuilder` in `crate::compile`.
#[derive(Clone, Copy)]
struct TrapSection<'a> {
    /// Total number of trap entries in this section.
    entries: usize,
    /// One `(first_offset, block_pos)` pair per block.
    block_index: &'a [[U32<LittleEndian>; 2]],
    /// Variable-length block bodies, index by `block_pos` in the `block_index`
    /// table above.
    block_bodies: &'a [u8],
}

impl<'a> TrapSection<'a> {
    /// Returns an iterator of `(text_offset, trap_code_byte)` for all entries
    /// in `block`, or `None` if the section is malformed.
    fn block_entries(&self, block_index: usize) -> Option<BlockEntries<'a>> {
        let [first_offset, block_pos] = self.block_index.get(block_index)?;
        let first_offset = first_offset.get(LittleEndian);
        let block_pos = block_pos.get(LittleEndian);
        let mut block = self.block_bodies.get(usize::try_from(block_pos).ok()?..)?;
        let default_code = bytes::pop(&mut block)?;
        let remaining = core::cmp::min(
            TRAP_BLOCK_SIZE,
            self.entries.checked_sub(block_index * TRAP_BLOCK_SIZE)?,
        );
        Some(BlockEntries {
            block,
            prev_offset: first_offset,
            default_code,
            remaining,
        })
    }
}

/// Iterator over the entries of a single block, decoding the
/// delta-and-code-flag varints described in the "block body" portion of the
/// section format on `TrapEncodingBuilder` in `crate::compile`.
struct BlockEntries<'a> {
    block: &'a [u8],
    prev_offset: u32,
    default_code: u8,
    remaining: usize,
}

impl Iterator for BlockEntries<'_> {
    type Item = (u32, u8);

    fn next(&mut self) -> Option<(u32, u8)> {
        self.remaining = self.remaining.checked_sub(1)?;
        let token = bytes::read_uleb(&mut self.block)?;
        let delta = u32::try_from(token >> 1).ok()?;
        let cur_offset = self.prev_offset.checked_add(delta)?;
        self.prev_offset = cur_offset;
        let code = if token & 1 != 0 {
            bytes::pop(&mut self.block)?
        } else {
            self.default_code
        };
        Some((cur_offset, code))
    }
}

fn parse(section: &[u8]) -> Option<TrapSection<'_>> {
    let mut section = Bytes(section);
    // NB: this matches the encoding written by `TrapEncodingBuilder`.
    let entries = section.read::<U32<LittleEndian>>().ok()?;
    let entries = usize::try_from(entries.get(LittleEndian)).ok()?;
    let num_blocks = section.read::<U32<LittleEndian>>().ok()?;
    let num_blocks = usize::try_from(num_blocks.get(LittleEndian)).ok()?;
    let (block_index, block_bodies) =
        object::slice_from_bytes::<[U32<LittleEndian>; 2]>(section.0, num_blocks).ok()?;
    Some(TrapSection {
        entries,
        block_index,
        block_bodies,
    })
}

/// Returns an iterator over all of the traps encoded in `section`, which should
/// have been produced by `TrapEncodingBuilder`.
pub fn iterate_traps(section: &[u8]) -> Option<impl Iterator<Item = (u32, CompiledTrap)> + '_> {
    let section = parse(section)?;
    Some(
        (0..section.block_index.len())
            .flat_map(move |block| section.block_entries(block).into_iter().flatten())
            .map(|(pc, byte)| (pc, CompiledTrap::from_u8(byte).unwrap())),
    )
}
