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);
 }
 }

}