using Antlr4.Runtime;
using Antlr4.Runtime.Misc;
using Antlr4.Runtime.Tree;
using QuickScript.Core;
using QuickScript.Core.VM.Function;
using QuickScript.Core.VM.Opcode;
using QuickScript.Core.VM.Value;
using QuickScript.Parser;
using VmValueTool = QuickScript.Core.VM.Value.ValueTool;
using VmValueType = QuickScript.Core.VM.Value.VMType;

namespace QuickScript.Compiler;

/// <summary>
/// QuickScript 字节码编译器 - 将 ANTLR parse tree 编译为 ScriptLib（Init 段 + Opcode 段）
/// </summary>
public partial class BytecodeCompiler
{
    private readonly ScriptLib _lib = new();
    private readonly List<OpcodeInstruction> _code = [];

    // 叽数管理：name → (index, type)
    private readonly Dictionary<string, (int Index, VmValueType Type, VmValueType ElementType)> _variables = new();
    private readonly List<string> _constStrings = new();
    private int _globalVarIndex;

    // 函数管理：name → (index, signature)
    private readonly Dictionary<string, FuncInfo> _functions = new();
    private readonly List<(string? Name, ScriptFunction Func)> _compiledFunctions = new();
    private int _funcIndex;

    // 函数体编译上下文
    private bool _inFunction;
    private int _funcParamCount;
    private VmValueType _currentReturnType;

    // 循环层次（用于 break/continue）
    private readonly Stack<LoopInfo> _loopStack = new();

    /// <summary>
    /// 初始化内置函数
    /// </summary>
    private void InitializeBuiltinFunctions()
    {
        // pow 不是真正的函数，只是语法糖，直接编译为 POW 指令
        // 因此不需要在 _functions 和 _compiledFunctions 中注册
    }

    private record LoopInfo(int ContinueOffset, List<int> BreakJumpIndices, List<int> ContinueJumpIndices);

    // switch 上下文：break 在 switch 中也需要跳到 end
    private readonly Stack<List<int>> _switchBreakStack = new();

    private record FuncInfo(int Index, FunctionSignature Signature);

    /// <summary>
    /// 编译 parse tree 为 ScriptLib
    /// </summary>
    public ScriptLib Compile(QuickScriptParser.ProgramContext ctx)
    {
        _lib.Inits.Clear();
        _lib.Opcodes.Clear();
        _code.Clear();
        _variables.Clear();
        _constStrings.Clear();
        _globalVarIndex = 0;
        _loopStack.Clear();
        _switchBreakStack.Clear();
        _functions.Clear();
        _compiledFunctions.Clear();
        _funcIndex = 0;
        _inFunction = false;
        _funcParamCount = 0;

        // 初始化内置函数
        InitializeBuiltinFunctions();

        // 第一遍：收集函数声明并编译函数体
        foreach (var stmt in ctx.statement())
        {
            if (stmt.funcDeclaration() is { } funcDecl)
                CompileFuncDeclaration(funcDecl);
        }

        // 第二遍：编译非函数语句
        foreach (var stmt in ctx.statement())
        {
            if (stmt.funcDeclaration() is null)
                CompileStatement(stmt);
        }

        // 确保 main block 末尾有 Rtn
        if (_code.Count == 0 || (_code[^1].Opcode != Opcode.Rtn && _code[^1].Opcode != Opcode.RtnValue))
        {
            Emit(OpcodeInstruction.CreateRtn());
        }

        // 生成 Init 段
        BuildInitSegment();

        _lib.Opcodes.AddRange(_code);
        return _lib;
    }

    // ===== Init 段生成 =====

    private void BuildInitSegment()
    {
        // 计算全局变量所需字节空间
        int globalSize = 0;
        foreach (var kv in _variables)
        {
            globalSize += VmValueTool.GetDataLengthWithType(kv.Value.Type);
        }

        // 额外空间用于函数调用中的 VarNewTmp 等动态分配
        int constSize = Math.Max(globalSize, 64);
        globalSize = Math.Max(globalSize + 256, 256);
        int tempSize = 1024;

        _lib.Inits.Add(new InitVarSpaceInstruction(constSize, globalSize, tempSize));

        // 字符串常量
        foreach (var str in _constStrings)
        {
            _lib.Inits.Add(new InitConstStringInstruction(str));
        }

        // 函数（无 public 修饰的函数 Name 为 null）
        foreach (var (name, func) in _compiledFunctions)
        {
            _lib.Inits.Add(new InitFuncInstruction(func, name));
        }

        _lib.Inits.Add(new InitEndInstruction());
    }

    // ===== 语句编译 =====

    private void CompileStatement(QuickScriptParser.StatementContext ctx)
    {
        if (ctx.variableDeclaration() is { } varDecl)
            CompileVariableDeclaration(varDecl);
        else if (ctx.constDeclaration() is { } constDecl)
            CompileConstDeclaration(constDecl);
        else if (ctx.funcDeclaration() is { })
        {
            // 函数声明已在第一遍处理，跳过
        }
        else if (ctx.ifStatement() is { } ifStmt)
            CompileIfStatement(ifStmt);
        else if (ctx.switchStatement() is { } switchStmt)
            CompileSwitchStatement(switchStmt);
        else if (ctx.whileStatement() is { } whileStmt)
            CompileWhileStatement(whileStmt);
        else if (ctx.forStatement() is { } forStmt)
            CompileForStatement(forStmt);
        else if (ctx.doStatement() is { } doStmt)
            CompileDoWhileStatement(doStmt);
        else if (ctx.returnStatement() is { } retStmt)
            CompileReturnStatement(retStmt);
        else if (ctx.blockStatement() is { } block)
            CompileBlockStatement(block);
        else if (ctx.expressionStmt() is { } exprStmt)
            CompileExpressionStmt(exprStmt);
        else if (ctx.breakStatement() is { })
            CompileBreakStatement();
        else if (ctx.continueStatement() is { })
            CompileContinueStatement();
        else if (ctx.emptyStatement() is { })
        {
            // 空语句，什么都不做
        }
    }

    private void CompileVariableDeclaration(QuickScriptParser.VariableDeclarationContext ctx)
    {
        var name = ctx.Identifier().GetText();
        bool isVar = ctx.VAR() is not null;
        VmValueType vmType;
        VmValueType elementType; // for arrays, the element type; same as vmType for non-arrays

        if (isVar)
        {
            vmType = InferExpressionType(ctx.expression());
            elementType = vmType;
        }
        else
        {
            var typeCtx = ctx.type();
            elementType = TypeToVmType(typeCtx);
            if (IsArrayType(typeCtx))
            {
                vmType = VmValueType.MemRef; // array variable holds a MemRef
            }
            else
            {
                vmType = elementType;
            }
        }

        if (_inFunction)
        {
            int varIndex = _funcParamCount + _globalVarIndex++;
            _variables[name] = (varIndex, vmType, elementType);

            if (vmType == VmValueType.MemRef)
            {
                // Array in function: compile new expression, store to temp
                if (ctx.expression() is { } expr)
                    CompileExpression(expr);
                else
                    EmitDefaultZero(vmType);
                EmitVarStoreFunc(varIndex);
            }
            else
            {
                // 先编译表达式（压入值），再 VarNewTmp（压入索引），最后 VarStTmp（弹出索引和值）
                if (ctx.expression() is { } expr)
                    CompileExpression(expr);
                else
                    EmitDefaultZero(vmType);
                Emit(OpcodeInstruction.CreateVarNewTmp(vmType));
                Emit(OpcodeInstruction.CreateVarStTmp());
            }
        }
        else
        {
            int varIndex = _globalVarIndex++;
            _variables[name] = (varIndex, vmType, elementType);

            if (vmType == VmValueType.MemRef)
            {
                // Array global: Var 和 Memory 共享 StoreGlobal MemoryArea，
                // 所以每个数组变量占用 2 个 slot：1 个 Var + 1 个 Memory 数据

                // 分配 Var slot
                Emit(OpcodeInstruction.CreateVarNewGlobal(vmType));

                if (ctx.expression() is { } expr)
                {
                    // 编译 new 表达式，栈: [index, MemRef]
                    CompileExpression(expr);
                }
                else
                {
                    // 无初始值：零长数组
                    EmitDefaultZero(vmType);
                }

                // 交换栈顶: [MemRef, index]
                Emit(OpcodeInstruction.CreateStackSwap());
                // VarStGlobal: pop index, pop MemRef, 存储
                Emit(OpcodeInstruction.CreateVarStGlobal());

                // MemoryNewGlobal 会在 StoreGlobal 中额外分配一个 Memory slot
                // 跳过该 slot，使后续变量的 index 与运行时一致
                _globalVarIndex++;
            }
            else
            {
                // 先编译表达式（压入值），再 VarNewGlobal（压入索引），最后 VarStGlobal（弹出索引和值）
                if (ctx.expression() is { } expr)
                    CompileExpression(expr);
                else
                    EmitDefaultZero(vmType);
                Emit(OpcodeInstruction.CreateVarNewGlobal(vmType));
                Emit(OpcodeInstruction.CreateVarStGlobal());
            }
        }
    }

    private void CompileConstDeclaration(QuickScriptParser.ConstDeclarationContext ctx)
    {
        var typeCtx = ctx.type();
        var name = ctx.Identifier().GetText();
        var vmType = TypeToVmType(typeCtx);

        int varIndex = _globalVarIndex++;
        _variables[name] = (varIndex, vmType, vmType);
        CompileExpression(ctx.expression());

        // VarNewGlobal + VarStGlobal 存储常量值
        Emit(OpcodeInstruction.CreateVarNewGlobal(vmType));
        Emit(OpcodeInstruction.CreateVarStGlobal());
    }

    private void CompileIfStatement(QuickScriptParser.IfStatementContext ctx)
    {
        CompileExpression(ctx.expression());

        var stmts = ctx.statement();
        bool hasElse = ctx.ELSE() is not null;

        // JmpFalse → else/end（占位）
        int jmpFalseIdx = _code.Count;
        Emit(OpcodeInstruction.CreateJmpFalse(0));

        // then 分支
        CompileStatement(stmts[0]);

        if (hasElse)
        {
            // then 分支结束后 Jmp → end（占位）
            int jmpEndIdx = _code.Count;
            Emit(OpcodeInstruction.CreateJmp(0));

            // else 分支
            PatchJump(jmpFalseIdx, _code.Count);
            CompileStatement(stmts[1]);

            // 修补 Jmp end
            PatchJump(jmpEndIdx, _code.Count);
        }
        else
        {
            // 无 else，直接修补 JmpFalse → end
            PatchJump(jmpFalseIdx, _code.Count);
        }
    }

    private void CompileWhileStatement(QuickScriptParser.WhileStatementContext ctx)
    {
        int loopStart = _code.Count;

        // 编译条件
        CompileExpression(ctx.expression());

        int jmpFalseIdx = _code.Count;
        Emit(OpcodeInstruction.CreateJmpFalse(0));

        // 循环体
        var breakJumps = new List<int>();
        var continueJumps = new List<int>();
        _loopStack.Push(new LoopInfo(loopStart, breakJumps, continueJumps));
        CompileStatement(ctx.statement());
        _loopStack.Pop();

        // 回跳到条件
        Emit(OpcodeInstruction.CreateJmp(ComputeByteOffset(loopStart)));

        // 修补退出
        PatchJump(jmpFalseIdx, _code.Count);

        // 修补所有 break 跳转
        foreach (int breakIdx in breakJumps)
        {
            PatchJump(breakIdx, _code.Count);
        }

        // 修补所有 continue 跳转到条件（while 的 continue 回到条件判断）
        foreach (int contIdx in continueJumps)
        {
            PatchJump(contIdx, loopStart);
        }
    }

    private void CompileDoWhileStatement(QuickScriptParser.DoStatementContext ctx)
    {
        int loopStart = _code.Count;

        // 循环体
        var breakJumps = new List<int>();
        var continueJumps = new List<int>();
        _loopStack.Push(new LoopInfo(loopStart, breakJumps, continueJumps));
        CompileStatement(ctx.statement());
        _loopStack.Pop();

        // 条件判断（continue 跳到这里）
        int condStart = _code.Count;
        CompileExpression(ctx.expression());
        Emit(OpcodeInstruction.CreateJmpTrue(ComputeByteOffset(loopStart)));

        // 修补所有 break 跳转
        foreach (int breakIdx in breakJumps)
        {
            PatchJump(breakIdx, _code.Count);
        }

        // 修补所有 continue 跳转到条件判断
        foreach (int contIdx in continueJumps)
        {
            PatchJump(contIdx, condStart);
        }
    }

    private void CompileForStatement(QuickScriptParser.ForStatementContext ctx)
    {
        var expressions = ctx.expression();
        bool hasVarDecl = ctx.varDecl() is not null;
        int condIndex = hasVarDecl ? 0 : 1;
        int incIndex = hasVarDecl ? 1 : 2;

        // init 部分
        if (ctx.varDecl() is { } varDecl)
        {
            var typeCtx = varDecl.type();
            var name = varDecl.Identifier().GetText();
            var vmType = TypeToVmType(typeCtx);

            int varIndex = _globalVarIndex++;
            _variables[name] = (varIndex, vmType, vmType);

            if (varDecl.expression() is { } initExpr)
            {
                CompileExpression(initExpr);
            }
            else
            {
                EmitDefaultZero(vmType);
            }
            Emit(OpcodeInstruction.CreateVarNewGlobal(vmType));
            Emit(OpcodeInstruction.CreateVarStGlobal());
        }
        else if (expressions.Length > 0)
        {
            CompileExpression(expressions[0]);
            Emit(OpcodeInstruction.CreateStackPop());
        }

        int loopStart = _code.Count;

        // condition 部分
        if (condIndex < expressions.Length)
        {
            CompileExpression(expressions[condIndex]);
        }
        else
        {
            // 无条件 → 永真
            Emit(OpcodeInstruction.CreateLdcBool(true));
        }

        int jmpFalseIdx = _code.Count;
        Emit(OpcodeInstruction.CreateJmpFalse(0));

        // 循环体
        var breakJumps = new List<int>();
        var continueJumps = new List<int>();
        _loopStack.Push(new LoopInfo(0, breakJumps, continueJumps)); // ContinueOffset 占位，后面修补
        CompileStatement(ctx.statement());
        _loopStack.Pop();

        // increment 部分（continue 应跳到这里）
        int incStart = _code.Count;
        if (incIndex < expressions.Length)
        {
            CompileExpression(expressions[incIndex]);
            Emit(OpcodeInstruction.CreateStackPop());
        }

        // 回跳
        Emit(OpcodeInstruction.CreateJmp(ComputeByteOffset(loopStart)));

        // 修补退出
        PatchJump(jmpFalseIdx, _code.Count);

        // 修补所有 continue 跳转到 increment 开始位置
        foreach (int contIdx in continueJumps)
        {
            PatchJump(contIdx, incStart);
        }

        // 修补所有 break 跳转到循环结束
        foreach (int breakIdx in breakJumps)
        {
            PatchJump(breakIdx, _code.Count);
        }
    }

    private void CompileReturnStatement(QuickScriptParser.ReturnStatementContext ctx)
    {
        if (ctx.expression() is { } expr)
        {
            CompileExpression(expr);
            // 函数内 String 返回值需要先复制到全局区域
            if (_inFunction && _currentReturnType == VmValueType.String)
                Emit(OpcodeInstruction.CreateStrToGlobal());
            Emit(OpcodeInstruction.CreateRtnValue());
        }
        else
        {
            Emit(OpcodeInstruction.CreateRtn());
        }
    }

    private void CompileBlockStatement(QuickScriptParser.BlockStatementContext ctx)
    {
        foreach (var stmt in ctx.statement())
        {
            CompileStatement(stmt);
        }
    }

    private void CompileExpressionStmt(QuickScriptParser.ExpressionStmtContext ctx)
    {
        if (ctx.expression() is { } expr)
        {
            var assignCtx = expr.assignmentExpression();
            bool isAssignment = assignCtx.Identifier() is not null && assignCtx.expression().Length > 0;

            // print() 编译为 Dmp 指令，Dmp 会自己弹出栈顶值，不需要额外的 StackPop
            bool isPrintCall = TryGetPrintCall(assignCtx);

            CompileExpression(expr);

            // 非赋值表达式弹出栈顶值（表达式结果丢弃）
            // 但 print() 已由 Dmp 弹出，不需要再弹
            if (!isAssignment && !isPrintCall)
                Emit(OpcodeInstruction.CreateStackPop());
        }
    }

    /// <summary>
    /// 检测表达式是否是 print() 调用
    /// </summary>
    private bool TryGetPrintCall(QuickScriptParser.AssignmentExpressionContext assignCtx)
    {
        var condExpr = assignCtx.conditionalExpression();
        if (condExpr == null) return false;
        var primary = WalkToPrimary(condExpr);
        if (primary == null) return false;

        if (primary.Identifier() is { } idNode && primary.argList() is not null)
        {
            return idNode.GetText() == "print";
        }
        return false;
    }

    /// <summary>
    /// 检测表达式是否是 void 函数调用
    /// </summary>
    private bool TryGetVoidFuncCall(QuickScriptParser.AssignmentExpressionContext assignCtx, out string? funcName)
    {
        funcName = null;
        // 路径: assignmentExpression → conditionalExpression → ... → primary → Identifier(argList)
        var condExpr = assignCtx.conditionalExpression();
        if (condExpr == null) return false;

        // 沿着条件表达式的逻辑或/与/位运算/比较/移位/加减/乘除/一元/后缀/primary 链路走到 primary
        var primary = WalkToPrimary(condExpr);
        if (primary == null) return false;

        if (primary.Identifier() is { } idNode && (primary.argList() is not null || primary.LeftParen() is not null))
        {
            funcName = idNode.GetText();
            if (funcName == "pow") return false; // pow 不是真函数

            if (_functions.TryGetValue(funcName, out var funcInfo) &&
                funcInfo.Signature.ReturnType == VmValueType.Void)
            {
                return true;
            }
        }
        return false;
    }

    /// <summary>
    /// 从条件表达式沿着一元链路走到 primary
    /// </summary>
    private static QuickScriptParser.PrimaryContext? WalkToPrimary(QuickScriptParser.ConditionalExpressionContext ctx)
    {
        var logicalOr = ctx.logicalOrExpression();
        if (logicalOr == null) return null;

        // 逐层展开：logicalOr → logicalAnd → bitwiseOr → ... → multiplicative → unary → postfix → primary
        var exprs = logicalOr.logicalAndExpression();
        if (exprs.Length != 1) return null;
        var bitwiseOrCtx = exprs[0].bitwiseOrExpression();
        if (bitwiseOrCtx.Length != 1) return null;
        var bitwiseXorCtx = bitwiseOrCtx[0].bitwiseXorExpression();
        if (bitwiseXorCtx.Length != 1) return null;
        var bitwiseAndCtx = bitwiseXorCtx[0].bitwiseAndExpression();
        if (bitwiseAndCtx.Length != 1) return null;
        var equalityCtx = bitwiseAndCtx[0].equalityExpression();
        if (equalityCtx.Length != 1) return null;
        var relationalCtx = equalityCtx[0].relationalExpression();
        if (relationalCtx.Length != 1) return null;
        var shiftCtx = relationalCtx[0].shiftExpression();
        if (shiftCtx.Length != 1) return null;
        var additiveCtx = shiftCtx[0].additiveExpression();
        if (additiveCtx.Length != 1) return null;
        var mulCtx = additiveCtx[0].multiplicativeExpression();
        if (mulCtx.Length != 1) return null;
        var unaryCtx = mulCtx[0].unaryExpression();
        if (unaryCtx.Length != 1) return null;
        if (unaryCtx[0].postfixExpression() is not { } postfix) return null;
        return postfix.primary();
    }

    private void CompileBreakStatement()
    {
        if (_loopStack.Count > 0)
        {
            var loop = _loopStack.Peek();
            int jmpIdx = _code.Count;
            Emit(OpcodeInstruction.CreateJmp(0));
            loop.BreakJumpIndices.Add(jmpIdx);
        }
        else if (_switchBreakStack.Count > 0)
        {
            // break in switch: pop switch value + jump to end
            Emit(OpcodeInstruction.CreateStackPop());
            int jmpIdx = _code.Count;
            Emit(OpcodeInstruction.CreateJmp(0));
            _switchBreakStack.Peek().Add(jmpIdx);
        }
        else
        {
            throw new InvalidOperationException("break 语句不在循环体或 switch 中");
        }
    }

    private void CompileContinueStatement()
    {
        if (_loopStack.Count == 0)
            throw new InvalidOperationException("continue 语句不在循环体内");

        var loop = _loopStack.Peek();
        int jmpIdx = _code.Count;
        Emit(OpcodeInstruction.CreateJmp(0)); // 占位，后续修补
        loop.ContinueJumpIndices.Add(jmpIdx);
    }

    private void CompileSwitchStatement(QuickScriptParser.SwitchStatementContext ctx)
    {
        // 编译 switch 表达式，结果留在栈上: [switchValue]
        CompileExpression(ctx.expression());

        var cases = ctx.switchCase();
        var caseJmpIndices = new int[cases.Length]; // 每个 case 的 JmpTrue 占位索引
        Array.Fill(caseJmpIndices, -1);
        var caseEndJumps = new List<int>();
        int defaultIdx = -1; // default 在 cases 数组中的索引
        int jmpToDefaultIdx = -1; // 跳转到 default 的 Jmp 占位

        // 阶段 1: 为每个非 default case 生成比较 + JmpTrue
        // 栈始终保持: [switchValue]
        // 每次比较: dup → [switchValue, switchValue]
        //           compile caseLiteral → [switchValue, switchValue, caseVal]
        //           LogicEq → [switchValue, bool]
        //           JmpTrue → [switchValue] (如果 false)
        for (int i = 0; i < cases.Length; i++)
        {
            var switchCase = cases[i];
            if (switchCase.DEFAULT() is not null)
            {
                defaultIdx = i;
                continue;
            }

            Emit(OpcodeInstruction.CreateStackDup());
            CompileLiteral(switchCase.literal());
            Emit(OpcodeInstruction.CreateLogicEq());
            caseJmpIndices[i] = _code.Count;
            Emit(OpcodeInstruction.CreateJmpTrue(0)); // 占位
        }

        // 所有 case 都不匹配时:
        // - 有 default: 无条件跳转到 default
        // - 无 default: pop switchValue 并跳到 end
        if (defaultIdx >= 0)
        {
            jmpToDefaultIdx = _code.Count;
            Emit(OpcodeInstruction.CreateJmp(0)); // 占位，跳转到 default body
        }
        else
        {
            Emit(OpcodeInstruction.CreateStackPop());
            caseEndJumps.Add(_code.Count);
            Emit(OpcodeInstruction.CreateJmp(0)); // 占位，跳转到 end
        }

        // 阶段 2: 编译每个 case body
        // 当跳转到 case body 时，栈上是: [switchValue]（从 dup 后 JmpTrue 消耗了 bool）
        var switchBreakJumps = new List<int>();
        _switchBreakStack.Push(switchBreakJumps);

        for (int i = 0; i < cases.Length; i++)
        {
            // 修补对应的跳转到当前位置
            if (cases[i].DEFAULT() is not null)
            {
                // default: 修补前面的无条件跳转
                if (jmpToDefaultIdx >= 0)
                    PatchJump(jmpToDefaultIdx, _code.Count);
            }
            else
            {
                PatchJump(caseJmpIndices[i], _code.Count);
            }

            // 编译 case body 语句
            foreach (var stmt in cases[i].statement())
            {
                CompileStatement(stmt);
            }

            // case body 结束后: pop switchValue，跳到 end
            // 除非是最后一个 case 且后面没有其他 case（fall-through 到结束）
            // 或者最后一条语句是 break/return
            bool lastStmtIsExit = false;
            var stmts = cases[i].statement();
            if (stmts.Length > 0)
            {
                var lastStmt = stmts[stmts.Length - 1];
                if (lastStmt.breakStatement() is not null || lastStmt.returnStatement() is not null)
                    lastStmtIsExit = true;
            }

            if (!lastStmtIsExit)
            {
                Emit(OpcodeInstruction.CreateStackPop());
                caseEndJumps.Add(_code.Count);
                Emit(OpcodeInstruction.CreateJmp(0)); // 占位，跳转到 end
            }
        }

        _switchBreakStack.Pop();

        // end:
        int endIdx = _code.Count;

        // 修补所有跳转到 end 的 Jmp（包括 caseEndJumps 和 switchBreakJumps）
        foreach (int jmpIdx in caseEndJumps.Concat(switchBreakJumps))
        {
            PatchJump(jmpIdx, endIdx);
        }
    }

}