luau 0.732.0

use core::marker::PhantomData;
use core::ptr::NonNull;
use std::rc::Rc;
use std::{mem, ptr};

use luau_common::ByteSlice;
use luau_vm::Thread as VmThread;
use luau_vm::native::{NativeCallContext, NativeCallResult};
use luau_vm::thread::upvalue_index;
use luau_vm::{VmError, VmExit};

use crate::error::Error;
use crate::lua::LuaRef;
use crate::lua::runtime::RuntimeData;
use crate::value::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti, MultiValue, Value};

/// Arguments supplied to an active Rust callback.
///
/// Values are converted directly from the callback stack. Consuming this value
/// with [`Arguments::finish`] completes the callback and prevents it from being
/// finalized more than once.
pub struct Arguments<'call> {
    lua: LuaRef<'call>,
    base_top: i32,
    count: i32,
    cursor: i32,
    position: usize,
}

/// The completed result of an active Rust callback.
///
/// Values of this type can only be produced by [`Arguments::finish`].
pub struct CallbackReturn<'call> {
    count: usize,
    _marker: PhantomData<LuaRef<'call>>,
}

/// A lazy view over values remaining on an active callback stack.
///
/// Creating this view does not convert or root any values. Values are converted
/// individually when accessed or iterated.
pub struct Varargs<'call> {
    lua: LuaRef<'call>,
    base_top: i32,
    count: i32,
}

impl<'call> Arguments<'call> {
    pub(crate) fn new(lua: LuaRef<'call>, base_top: i32, count: i32) -> Self {
        Self {
            lua,
            base_top,
            count: count.max(0),
            cursor: 0,
            position: 0,
        }
    }

    /// Returns the number of arguments supplied to the callback.
    pub const fn len(&self) -> usize {
        self.count as usize
    }

    /// Returns whether the callback was invoked without arguments.
    pub const fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// Returns the number of arguments not yet consumed by this cursor.
    pub const fn remaining(&self) -> usize {
        (self.count - self.cursor) as usize
    }

    /// Converts the next argument to `T`, or converts `nil` when no argument remains.
    ///
    /// This is intentionally generic per call, so successive arguments can be
    /// converted to different types and the cursor cannot implement `Iterator`.
    #[allow(clippy::should_implement_trait)]
    pub fn next<T>(&mut self) -> Result<T, Error>
    where
        T: FromLua<'call>,
    {
        let position = self.position + 1;
        let value = if self.cursor < self.count {
            // Arguments owns this initialized callback stack range.
            unsafe { T::from_stack(&self.lua.current_thread(), self.base_top + self.cursor + 1) }
        } else {
            T::from_lua(Value::Nil, self.lua)
        }
        .map_err(|error| Error::bad_argument(position, error))?;
        self.cursor = self.cursor.saturating_add(1).min(self.count);
        self.position = position;
        Ok(value)
    }

    /// Converts all unconsumed arguments to `T` and advances to the end.
    pub fn take_remaining<T>(&mut self) -> Result<T, Error>
    where
        T: FromLuaMulti<'call>,
    {
        let remaining = self.count - self.cursor;
        let position = self.position + 1;
        // Arguments owns this initialized callback stack range.
        let value = unsafe {
            T::from_stack_multi(
                &self.lua.current_thread(),
                self.base_top + self.cursor,
                remaining,
            )
        }
        .map_err(|error| Error::bad_argument(position, error))?;
        self.cursor = self.count;
        self.position = self.position.saturating_add(remaining as usize);
        Ok(value)
    }

    /// Takes a lazy view over all unconsumed arguments and advances to the end.
    pub fn take_varargs(&mut self) -> Varargs<'call> {
        let remaining = self.count - self.cursor;
        let varargs = Varargs {
            lua: self.lua,
            base_top: self.base_top + self.cursor,
            count: remaining,
        };
        self.cursor = self.count;
        self.position = self.position.saturating_add(remaining as usize);
        varargs
    }

    /// Pushes the callback results and consumes this invocation.
    pub fn finish<R>(self, values: R) -> Result<CallbackReturn<'call>, Error>
    where
        R: IntoLuaMulti<'call>,
    {
        let count = unsafe { values.push_into_stack_multi(&self.lua.current_thread())? };
        Ok(CallbackReturn {
            count,
            _marker: PhantomData,
        })
    }
}

impl<'call> Varargs<'call> {
    /// Returns the number of values in this view.
    pub const fn len(&self) -> usize {
        self.count as usize
    }

    /// Returns whether this view contains no values.
    pub const fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// Converts one value by zero-based index without converting its siblings.
    pub fn get<T>(&self, index: usize) -> Result<Option<T>, Error>
    where
        T: FromLua<'call>,
    {
        if index >= self.len() {
            return Ok(None);
        }
        let index = i32::try_from(index).map_err(|_| Error::index_out_of_bounds())?;
        // Varargs owns this initialized callback stack range.
        unsafe { T::from_stack(&self.lua.current_thread(), self.base_top + index + 1) }.map(Some)
    }

    /// Iterates over the values as heterogeneous [`Value`] handles.
    pub fn iter(&self) -> impl Iterator<Item = Result<Value<'call>, Error>> + '_ {
        (0..self.count)
            .map(|offset| Value::from_stack(&self.lua.current_thread(), self.base_top + offset + 1))
    }

    /// Iterates over the values, converting every element to `T` on demand.
    pub fn iter_as<T>(&self) -> impl Iterator<Item = Result<T, Error>> + '_
    where
        T: FromLua<'call>,
    {
        (0..self.count).map(|offset| {
            // Varargs owns this initialized callback stack range.
            unsafe { T::from_stack(&self.lua.current_thread(), self.base_top + offset + 1) }
        })
    }

    /// Converts and roots every value into a materialized [`MultiValue`].
    pub fn materialize(&self) -> Result<MultiValue<'call>, Error> {
        MultiValue::from_stack(&self.lua.current_thread(), self.base_top, self.count)
    }
}

pub(crate) trait Callback {
    fn call<'call>(
        &self,
        context: NativeCallContext<'call>,
        runtime: &'call RuntimeData,
    ) -> NativeCallResult;
}

pub(crate) struct CallbackEntry {
    callback: Option<Box<dyn Callback>>,
}

pub(crate) fn finish_callback(
    thread: &VmThread,
    runtime: &RuntimeData,
    result: Result<CallbackReturn<'_>, Error>,
    expected_results: Option<usize>,
) -> NativeCallResult {
    let result = match result {
        Ok(result) => result,
        Err(error) => return raise_callback_error(thread, runtime, error),
    };
    if let Some(expected) = expected_results
        && result.count != expected
    {
        return raise_callback_error(
            thread,
            runtime,
            Error::runtime(format_args!(
                "callback returned {} values, expected {expected}",
                result.count
            )),
        );
    }
    Ok(result.count)
}

pub(crate) fn raise_callback_error<T>(
    thread: &VmThread,
    runtime: &RuntimeData,
    error: Error,
) -> Result<T, VmExit> {
    raise_callback_vm_error(thread, runtime, error).map_err(Into::into)
}

pub(crate) fn raise_callback_vm_error<T>(
    thread: &VmThread,
    runtime: &RuntimeData,
    error: Error,
) -> Result<T, VmError> {
    let traceback = unsafe { thread.debug_trace() }
        .map(|traceback| String::from_utf8_lossy(traceback.as_bytes()).into_owned())
        .unwrap_or_default();
    let safe_thread = crate::Thread::new(thread, runtime);
    let pushed = unsafe {
        Error::CallbackError {
            traceback,
            cause: Rc::new(error),
        }
        .push_into_stack(&safe_thread)
    };
    match pushed {
        Ok(()) => Err(VmError::Runtime),
        Err(error) => error.raise_error(thread),
    }
}

pub(crate) fn push_callback(
    thread: &VmThread,
    callback: Box<dyn Callback>,
) -> Result<NonNull<CallbackEntry>, Error> {
    unsafe {
        let data = thread
            .new_userdata_dtor(mem::size_of::<CallbackEntry>(), drop_callback)
            .map_err(|exit| Error::from_thread_exit(thread, exit))?;
        let mut entry = NonNull::new_unchecked(data.cast::<CallbackEntry>());
        entry.as_ptr().write(CallbackEntry {
            callback: Some(callback),
        });

        match thread.push_native_closure(callback_trampoline, None, 1) {
            Ok(()) => Ok(entry),
            Err(exit) => {
                entry.as_mut().callback = None;
                thread.pop(1);
                Err(Error::from_thread_exit(thread, exit))
            }
        }
    }
}

pub(crate) unsafe fn invalidate_callback(mut entry: NonNull<CallbackEntry>) {
    unsafe {
        entry.as_mut().callback = None;
    }
}

fn callback_trampoline(context: NativeCallContext<'_>) -> NativeCallResult {
    unsafe {
        let thread = context.raw_thread();
        let data = thread.to_userdata(upvalue_index(1));
        if data.is_null() {
            return luau_vm::error!(thread, "callback upvalue must be userdata")
                .map_err(Into::into);
        }

        let entry = &*data.cast::<CallbackEntry>();
        let runtime = RuntimeData::from_thread(thread);
        let Some(callback) = &entry.callback else {
            return raise_callback_error(thread, runtime, Error::CallbackDestructed);
        };
        runtime.with_thread(thread, || callback.call(context, runtime))
    }
}

fn drop_callback(data: *mut ()) {
    unsafe {
        ptr::drop_in_place(data.cast::<CallbackEntry>());
    }
}