using QuickScript.Core;
using QuickScript.Core.VM.Function;
using QuickScript.Core.VM.Opcode;
using QuickScript.Core.VM.Value;
using QuickScript.Parser;
using VmValueType = QuickScript.Core.VM.Value.VMType;

namespace QuickScript.Compiler;

/// <summary>
/// BytecodeCompiler 函数编译部分
/// </summary>
public partial class BytecodeCompiler
{
    // ===== 函数声明编译 =====

    private void CompileFuncDeclaration(QuickScriptParser.FuncDeclarationContext ctx)
    {
        var name = ctx.Identifier().GetText();

        // 检测 public 修饰符
        bool isPublic = ctx.PUBLIC() is not null;

        // 解析返回类型
        VmValueType returnType = TypeToVmType(ctx.type()); // type 必填

        // 解析参数
        var paramTypes = new List<VmValueType>();
        var paramNames = new List<string>();
        if (ctx.paramList() is { } paramListCtx)
        {
            foreach (var param in paramListCtx.param())
            {
                paramTypes.Add(TypeToVmType(param.type()));
                paramNames.Add(param.Identifier().GetText());
            }
        }

        // 创建函数签名
        var signature = new FunctionSignature(paramTypes.ToArray(), returnType);

        // 注册函数
        int funcIndex = _funcIndex++;
        _functions[name] = new FuncInfo(funcIndex, signature);

        // 保存并切换编译上下文
        var savedCode = new List<OpcodeInstruction>(_code);
        var savedVariables = new Dictionary<string, (int Index, VmValueType Type, VmValueType ElementType)>(_variables);
        var savedConstStrings = new List<string>(_constStrings);
        int savedGlobalVarIndex = _globalVarIndex;
        var savedLoopStack = new Stack<LoopInfo>(_loopStack.Reverse());
        bool savedInFunction = _inFunction;
        int savedFuncParamCount = _funcParamCount;
        var savedReturnType = _currentReturnType;

        _code.Clear();
        _variables.Clear();
        _constStrings.Clear();
        _globalVarIndex = 0;
        _loopStack.Clear();

        // 设置函数编译上下文
        _inFunction = true;
        _funcParamCount = paramNames.Count;
        _currentReturnType = returnType;

        // 注册参数为变量
        for (int i = 0; i < paramNames.Count; i++)
        {
            _variables[paramNames[i]] = (i, paramTypes[i], paramTypes[i]);
        }

        // 编译函数体
        CompileBlockStatement(ctx.blockStatement());

        // 如果函数没有 return 语句，确保末尾有 Rtn
        if (_code.Count == 0 || (_code[^1].Opcode != Opcode.Rtn && _code[^1].Opcode != Opcode.RtnValue))
        {
            Emit(OpcodeInstruction.CreateRtn());
        }

        // 序列化函数指令为字节数组
        byte[] instrBytes = SerializeFunctionCode(_code);

        // 构建 ScriptFunction 反序列化所需的完整数据: [signature][instDataLen:4B][instructions]
        byte[] signatureBytes = signature.Serialize();
        byte[] fullData = new byte[signatureBytes.Length + 4 + instrBytes.Length];
        Array.Copy(signatureBytes, 0, fullData, 0, signatureBytes.Length);
        BitConverter.TryWriteBytes(fullData.AsSpan(signatureBytes.Length, 4), instrBytes.Length);
        Array.Copy(instrBytes, 0, fullData, signatureBytes.Length + 4, instrBytes.Length);

        // 使用安全的反序列化方法创建 ScriptFunction
        var scriptFunc = ScriptFunction.Deserialize(fullData);
        _compiledFunctions.Add((isPublic ? name : null, scriptFunc));

        // 恢复编译上下文
        _code.Clear();
        _code.AddRange(savedCode);
        _variables.Clear();
        foreach (var kv in savedVariables) _variables[kv.Key] = kv.Value;
        _constStrings.Clear();
        _constStrings.AddRange(savedConstStrings);
        _globalVarIndex = savedGlobalVarIndex;
        _loopStack.Clear();
        foreach (var loop in savedLoopStack) _loopStack.Push(loop);
        _inFunction = savedInFunction;
        _funcParamCount = savedFuncParamCount;
        _currentReturnType = savedReturnType;
    }

    /// <summary>
    /// 将函数体的指令列表序列化为字节数组
    /// </summary>
    private static byte[] SerializeFunctionCode(List<OpcodeInstruction> instructions)
    {
        // 计算总字节数
        int totalSize = 0;
        foreach (var instr in instructions)
        {
            totalSize += 1 + (instr.Data?.Length ?? 0);
        }

        byte[] result = new byte[totalSize];
        int offset = 0;
        foreach (var instr in instructions)
        {
            result[offset++] = (byte)instr.Opcode;
            if (instr.Data is not null)
            {
                Array.Copy(instr.Data, 0, result, offset, instr.Data.Length);
                offset += instr.Data.Length;
            }
        }

        return result;
    }

    /// <summary>
    /// 在函数体内编译变量加载（参数用 LoadArg，局部变量用 VarLdTmp）
    /// </summary>
    private void EmitVarLoadFunc(int index)
    {
        if (index < _funcParamCount)
        {
            // 参数使用 LoadArg
            Emit(OpcodeInstruction.CreateLoadArg((byte)index));
        }
        else
        {
            // 局部变量使用 VarLdTmp
            Emit(OpcodeInstruction.CreateLdcI32(index - _funcParamCount));
            Emit(OpcodeInstruction.CreateVarLdTmp());
        }
    }

    /// <summary>
    /// 在函数体内编译变量存储（局部变量用 VarStTmp）
    /// </summary>
    private void EmitVarStoreFunc(int index)
    {
        Emit(OpcodeInstruction.CreateLdcI32(index - _funcParamCount));
        Emit(OpcodeInstruction.CreateVarStTmp());
    }
}