DebugLibrary.cs

using System.Runtime.CompilerServices;
using Lua.Internal;
using Lua.Internal.CompilerServices;
using Lua.Runtime;

namespace Lua.Standard;

public sealed class DebugLibrary
{
    public static readonly DebugLibrary Instance = new();
    static readonly ConditionalWeakTable<UpValue, object> upValueIds = new();
    static readonly ConditionalWeakTable<LuaValue[], object[]> csUpValueIds = new();

    public DebugLibrary()
    {
        var libraryName = "debug";
        Functions =
        [
            new(libraryName, "getlocal", GetLocal),
            new(libraryName, "setlocal", SetLocal),
            new(libraryName, "getupvalue", GetUpValue),
            new(libraryName, "setupvalue", SetUpValue),
            new(libraryName, "getmetatable", GetMetatable),
            new(libraryName, "setmetatable", SetMetatable),
            new(libraryName, "getuservalue", GetUserValue),
            new(libraryName, "setuservalue", SetUserValue),
            new(libraryName, "traceback", Traceback),
            new(libraryName, "getregistry", GetRegistry),
            new(libraryName, "upvalueid", UpValueId),
            new(libraryName, "upvaluejoin", UpValueJoin),
            new(libraryName, "gethook", GetHook),
            new(libraryName, "sethook", SetHook),
            new(libraryName, "getinfo", GetInfo),
        ];
    }

    public readonly LibraryFunction[] Functions;

    static LuaState GetLuaThread(in LuaFunctionExecutionContext context, out int argOffset)
    {
        if (context.ArgumentCount < 1)
        {
            argOffset = 0;
            return context.State;
        }

        if (context.GetArgument(0).TryRead<LuaState>(out var state))
        {
            argOffset = 1;
            return state;
        }

        argOffset = 0;
        return context.State;
    }

    static ref LuaValue FindLocal(LuaState state, int level, int index, out string? name)
    {
        if (index == 0)
        {
            name = null;
            return ref Unsafe.NullRef<LuaValue>();
        }

        var callStack = state.GetCallStackFrames();
        var frame = callStack[^(level + 1)];
        if (index < 0)
        {
            index = -index - 1;
            var frameVariableArgumentCount = frame.VariableArgumentCount;
            if (frameVariableArgumentCount > 0 && index < frameVariableArgumentCount)
            {
                name = "(*vararg)";
                return ref state.GetLocalReference(frame.Base - frameVariableArgumentCount + index);
            }

            name = null;
            return ref Unsafe.NullRef<LuaValue>();
        }

        index -= 1;

        var frameBase = frame.Base;

        if (frame.Function is LuaClosure closure)
        {
            var locals = closure.Proto.LocalVariables;
            var nextFrame = callStack[^level];
            var currentPc = nextFrame.CallerInstructionIndex;
            {
                var nextFrameBase = closure.Proto.Code[currentPc].OpCode
                    is OpCode.Call
                        or OpCode.TailCall
                    ? nextFrame.ReturnBase
                    : nextFrame.Base - nextFrame.VariableArgumentCount;
                if (index >= nextFrameBase - frameBase)
                {
                    name = null;
                    return ref Unsafe.NullRef<LuaValue>();
                }
            }

            var localId = index + 1;
            foreach (var l in locals)
            {
                if (currentPc < l.StartPc)
                {
                    break;
                }

                if (l.EndPc <= currentPc)
                {
                    continue;
                }

                localId--;
                if (localId == 0)
                {
                    name = l.Name;
                    return ref state.GetLocalReference(frameBase + index);
                }
            }
        }
        else
        {
            var nextFrameBase =
                level != 0
                    ? callStack[^level].Base - callStack[^level].VariableArgumentCount
                    : state.Stack.Count;

            if (index >= nextFrameBase - frameBase)
            {
                name = null;
                return ref Unsafe.NullRef<LuaValue>();
            }
        }

        name = "(*temporary)";
        return ref state.GetLocalReference(frameBase + index);
    }

    public ValueTask<int> GetLocal(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        static LuaValue GetParam(LuaFunction function, int index)
        {
            if (function is LuaClosure closure)
            {
                var locals = closure.Proto.LocalVariables;
                var paramCount = Math.Min(closure.Proto.ParameterCount, locals.Length);
                if (0 <= index && index < paramCount)
                {
                    return locals[index].Name;
                }
            }

            return LuaValue.Nil;
        }

        var state = GetLuaThread(context, out var argOffset);

        var index = context.GetIntegerArgument(argOffset + 1);
        if (context.GetArgument(argOffset).TryReadFunction(out var f))
        {
            return new(context.Return(GetParam(f, index - 1)));
        }

        var level = context.GetIntegerArgument(argOffset);

        var discarded = state.GetNativeTailCaller().Prototype != null;
        if (discarded && level == 1)
        {
            // The discarded caller's registers are gone, so it has no readable locals.
            return new(context.Return(LuaValue.Nil));
        }

        if (discarded && level > 1)
            level--;

        if (level < 0 || level >= state.GetCallStackFrames().Length)
        {
            context.ThrowBadArgument(1, "level out of range");
        }

        ref var local = ref FindLocal(state, level, index, out var name);
        if (name is null)
        {
            return new(context.Return(LuaValue.Nil));
        }

        return new(context.Return(name, local));
    }

    public ValueTask<int> SetLocal(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var state = GetLuaThread(context, out var argOffset);

        var value = context.GetArgument(argOffset + 2);
        var index = context.GetIntegerArgument(argOffset + 1);
        var level = context.GetIntegerArgument(argOffset);

        var discarded = state.GetNativeTailCaller().Prototype != null;
        if (discarded && level == 1)
        {
            // The discarded caller's registers are gone, so it has no writable locals.
            return new(context.Return(LuaValue.Nil));
        }

        if (discarded && level > 1)
            level--;

        if (level < 0 || level >= state.GetCallStackFrames().Length)
        {
            context.ThrowBadArgument(1, "level out of range");
        }

        ref var local = ref FindLocal(state, level, index, out var name);
        if (name is null)
        {
            return new(context.Return(LuaValue.Nil));
        }

        local = value;
        return new(context.Return(name));
    }

    public ValueTask<int> GetUpValue(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var func = context.GetArgument<LuaFunction>(0);
        var index = context.GetIntegerArgument(1) - 1;
        if (func is not LuaClosure closure)
        {
            if (func is CSharpClosure csClosure)
            {
                var upValues = csClosure.UpValues;
                if (index < 0 || index >= upValues.Length)
                {
                    return new(context.Return());
                }

                return new(context.Return("", upValues[index]));
            }

            return new(context.Return());
        }

        {
            var upValues = closure.UpValues;
            var descriptions = closure.Proto.UpValues;
            if (index < 0 || index >= descriptions.Length)
            {
                return new(context.Return());
            }

            var description = descriptions[index];
            return new(context.Return(description.Name.ToString(), upValues[index].GetValue()));
        }
    }

    public ValueTask<int> SetUpValue(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var func = context.GetArgument<LuaFunction>(0);
        var index = context.GetIntegerArgument(1) - 1;
        var value = context.GetArgument(2);
        if (func is not LuaClosure closure)
        {
            if (func is CSharpClosure csClosure)
            {
                var upValues = csClosure.UpValues;
                if (index >= 0 && index < upValues.Length)
                {
                    upValues[index] = value;
                    return new(context.Return(""));
                }
            }

            return new(context.Return());
        }

        {
            var upValues = closure.UpValues;
            var descriptions = closure.Proto.UpValues;
            if (index < 0 || index >= descriptions.Length)
            {
                return new(context.Return());
            }

            var description = descriptions[index];
            upValues[index].SetValue(value);
            return new(context.Return(description.Name.ToString()));
        }
    }

    public ValueTask<int> GetMetatable(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var arg0 = context.GetArgument(0);

        if (context.GlobalState.TryGetMetatable(arg0, out var table))
        {
            return new(context.Return(table));
        }
        else
        {
            return new(context.Return(LuaValue.Nil));
        }
    }

    public ValueTask<int> SetMetatable(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var arg0 = context.GetArgument(0);
        var arg1 = context.GetArgument(1);

        if (arg1.Type is not (LuaValueType.Nil or LuaValueType.Table))
        {
            LuaRuntimeException.BadArgument(
                context.State,
                2,
                [LuaValueType.Nil, LuaValueType.Table],
                arg1.Type
            );
        }

        context.GlobalState.SetMetatable(arg0, arg1.UnsafeRead<LuaTable>());

        return new(context.Return(arg0));
    }

    public ValueTask<int> GetUserValue(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        if (!context.GetArgumentOrDefault(0).TryRead<ILuaUserData>(out var iUserData))
        {
            return new(context.Return(LuaValue.Nil));
        }

        var index = 1; // context.GetArgument<int>(1); //for lua 5.4
        var userValues = iUserData.UserValues;
        if (
            index > userValues.Length /* index < 1 ||  // for lua 5.4 */
        )
        {
            return new(context.Return(LuaValue.Nil));
        }

        return new(context.Return(userValues[index - 1]));
    }

    public ValueTask<int> SetUserValue(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var iUserData = context.GetArgument<ILuaUserData>(0);
        var value = context.GetArgument(1);
        var index = 1; // context.GetArgument<int>(2); // for lua 5.4
        var userValues = iUserData.UserValues;
        if (
            index > userValues.Length /* || index < 1 // for lua 5.4 */
        )
        {
            return new(context.Return(LuaValue.Nil));
        }

        userValues[index - 1] = value;
        return new(context.Return(new LuaValue(iUserData)));
    }

    public ValueTask<int> Traceback(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var state = GetLuaThread(context, out var argOffset);

        var message = context.GetArgumentOrDefault(argOffset);
        if (message.Type is not (LuaValueType.Nil or LuaValueType.String or LuaValueType.Number))
        {
            return new(context.Return(message));
        }

        var level =
            context.HasArgument(argOffset + 1) ? context.GetIntegerArgument(argOffset + 1)
            : state == context.State ? 1
            : 0;

        ReadOnlySpan<CallStackFrame> frames =
            level < 0 ? []
            : state is { IsCoroutine: true, LuaTraceback: { } traceback } ? traceback.StackFrames
            : state.GetCallStackFrames();
        return new(
            context.Return(
                Runtime.Traceback.CreateTracebackMessage(
                    context.GlobalState,
                    frames,
                    message,
                    level,
                    abbreviate: true
                )
            )
        );
    }

    public ValueTask<int> GetRegistry(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        return new(context.Return(context.GlobalState.Registry));
    }

    static LuaFunction CheckUpValue(
        in LuaFunctionExecutionContext context,
        int functionIndex,
        int upValueIndex,
        out int index
    )
    {
        index = context.GetIntegerArgument(upValueIndex);
        var function = context.GetArgument<LuaFunction>(functionIndex);
        var count = function switch
        {
            LuaClosure closure => closure.GetUpValuesSpan().Length,
            CSharpClosure closure => closure.UpValues.Length,
            _ => 0,
        };
        if (index <= 0 || index > count)
            context.ThrowBadArgument(upValueIndex + 1, "invalid upvalue index");
        return function;
    }

    public ValueTask<int> UpValueId(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var function = CheckUpValue(context, 0, 1, out var index);

        if (function is LuaClosure closure)
        {
            var identity = upValueIds.GetValue(
                closure.GetUpValuesSpan()[index - 1],
                static _ => new object()
            );
            return new(context.Return(LuaValue.FromObject(identity)));
        }

        var csClosure = (CSharpClosure)function;
        var identities = csUpValueIds.GetValue(
            csClosure.UpValues,
            static values =>
            {
                var result = new object[values.Length];
                for (var i = 0; i < result.Length; i++)
                    result[i] = new object();
                return result;
            }
        );
        return new(context.Return(LuaValue.FromObject(identities[index - 1])));
    }

    public ValueTask<int> UpValueJoin(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var f1 = CheckUpValue(context, 0, 1, out var n1);
        var f2 = CheckUpValue(context, 2, 3, out var n2);

        if (f1 is not LuaClosure)
            context.ThrowBadArgument(1, "Lua function expected");
        if (f2 is not LuaClosure)
            context.ThrowBadArgument(3, "Lua function expected");

        var upValues1 = ((LuaClosure)f1).GetUpValuesSpan();
        var upValues2 = ((LuaClosure)f2).GetUpValuesSpan();
        upValues1[n1 - 1] = upValues2[n2 - 1];
        return new(context.Return());
    }

    [AsyncMethodBuilder(typeof(LightAsyncValueTaskMethodBuilder<>))]
    public async ValueTask<int> SetHook(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var state = GetLuaThread(context, out var argOffset);
        var hook = context.GetArgumentOrDefault<LuaFunction?>(argOffset);
        if (hook is null)
        {
            state.SetHookFromLua(null, "", 0);
            return context.Return();
        }
        var mask = context.GetStringArgument(argOffset + 1);
        var count = context.HasArgument(argOffset + 2)
            ? context.GetIntegerArgument(argOffset + 2)
            : 0;
        state.SetHookFromLua(hook, mask, count);

        // Re-arming a hook must not dispatch another hook while its callback is running.
        if (state.IsReturnHookEnabled && context.State == state && !state.IsInHook)
        {
            var stack = state.Stack;
            var top = stack.Count;
            stack.Push("return");
            stack.Push(LuaValue.Nil);
            context.State.IsInHook = true;
            context.State.NonYieldableCallCount++;
            var framePushed = false;
            try
            {
                var frame = context.State.CreateCallStackFrame(hook, 2, top, 0);
                context.State.PushCallStackFrame(frame);
                framePushed = true;
                LuaFunctionExecutionContext funcContext = new()
                {
                    State = context.State,
                    ArgumentCount = stack.Count - frame.Base,
                    ReturnFrameBase = frame.ReturnBase,
                };
                await hook.Func(funcContext, cancellationToken);
            }
            finally
            {
                context.State.NonYieldableCallCount--;
                context.State.IsInHook = false;
                if (framePushed)
                {
                    context.State.PopCallStackFrameWithStackPop();
                }
            }
        }

        return context.Return();
    }

    public ValueTask<int> GetHook(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var state = GetLuaThread(context, out _);
        return new(
            context.Return(
                state.Hook is null ? LuaValue.Nil : state.Hook,
                ((state.HookMask & 1) != 0 ? "c" : "")
                    + ((state.HookMask & 2) != 0 ? "r" : "")
                    + ((state.HookMask & 4) != 0 ? "l" : ""),
                state.BaseHookCount
            )
        );
    }

    public ValueTask<int> GetInfo(
        LuaFunctionExecutionContext context,
        CancellationToken cancellationToken
    )
    {
        var state = GetLuaThread(context, out var argOffset);
        var what = context.GetArgumentOrDefault(argOffset + 1, "flnStu");
        var arg1 = context.GetArgument(argOffset);

        arg1.TryReadFunction(out var functionToInspect);
        Prototype? prototype = null;
        int? pc = null;
        var isTailCall = false;
        (Prototype? Prototype, int Pc) caller = default;

        if (functionToInspect == null)
        {
            if (!arg1.TryReadDouble(out _))
                context.ThrowBadArgument(argOffset, "function or level expected");
            var level = context.GetIntegerArgument(argOffset);

            var callStack = state.GetCallStackFrames();
            var tailCaller = state.GetNativeTailCaller();
            var discarded = tailCaller.Prototype != null;
            if ((uint)level >= (uint)(callStack.Length + (discarded ? 1 : 0)))
            {
                return new(context.Return(LuaValue.Nil));
            }

            // Level 1 uses the native frame's slot to describe the discarded Lua caller.
            var synthetic = discarded && level == 1;
            var index = callStack.Length - 1 - level;
            if (discarded && level > 0)
                index++;
            ref readonly var frame = ref callStack[index];
            if (synthetic)
            {
                prototype = tailCaller.Prototype;
                pc = tailCaller.Pc;
                isTailCall = tailCaller.IsTailCall;
            }
            else
            {
                functionToInspect = frame.Function;
                pc = index + 1 < callStack.Length ? callStack[index + 1].CallerInstructionIndex : 0;
                isTailCall = frame.IsTailCall;
            }

            if (discarded && level == 0)
            {
                caller = (tailCaller.Prototype, tailCaller.Pc);
                isTailCall = false;
            }
            else if (index > 0 && !(synthetic && isTailCall))
            {
                caller = (
                    (callStack[index - 1].Function as LuaClosure)?.Proto,
                    frame.CallerInstructionIndex
                );
            }
        }

        prototype ??= (functionToInspect as LuaClosure)?.Proto;
        using var debug = LuaDebug.Create(
            context.GlobalState,
            functionToInspect,
            prototype,
            pc,
            isTailCall,
            caller,
            what,
            out var isValid
        );
        if (!isValid)
        {
            context.ThrowBadArgument(argOffset + 2, "invalid option");
        }

        LuaTable table = new(0, 1);
        if (what.Contains('S'))
        {
            table["source"] = debug.Source ?? LuaValue.Nil;
            table["short_src"] = debug.ShortSource.ToString();
            table["linedefined"] = debug.LineDefined;
            table["lastlinedefined"] = debug.LastLineDefined;
            table["what"] = debug.What ?? LuaValue.Nil;
            ;
        }

        if (what.Contains('l'))
        {
            table["currentline"] = debug.CurrentLine;
        }

        if (what.Contains('u'))
        {
            table["nups"] = debug.UpValueCount;
            table["nparams"] = debug.ParameterCount;
            table["isvararg"] = debug.IsVarArg;
        }

        if (what.Contains('n'))
        {
            table["name"] = debug.Name ?? LuaValue.Nil;
            table["namewhat"] = debug.NameWhat ?? LuaValue.Nil;
        }

        if (what.Contains('t'))
        {
            table["istailcall"] = debug.IsTailCall;
        }

        if (what.Contains('f'))
        {
            // A discarded caller's closure is released, so only its prototype-based fields remain.
            table["func"] = functionToInspect is { } inspectedFunction
                ? inspectedFunction
                : LuaValue.Nil;
        }

        if (what.Contains('L') && prototype is not null)
        {
            LuaTable activeLines = new(0, 8);
            foreach (var line in prototype.LineInfo)
                activeLines[line] = true;
            table["activelines"] = activeLines;
        }

        return new(context.Return(table));
    }
}