namespace STCompiler.Compiler; 

using System;
using System.IO;
using System.Text;
using System.Collections.Generic;
using STCompiler.Common;

// === ENTRY POINT ===
class Program {
    static int Main(string[] args) {
        if (args.Length < 2) {
            Console.WriteLine("Usage: StEmitter <input.st> <output.bin>");
            return 1;
        }

        var input = File.ReadAllText(args[0]);
        var parser = new StParser(input);
        var prog = parser.ParseProgram();
        
        if (parser.HasErrors) {
            Console.WriteLine($"Compilation failed with {parser.Errors.Count} errors:");
            foreach (var error in parser.Errors) {
                Console.WriteLine($"Error at line {error.Line}: {error.Message}");
            }
            return 1;
        }

        if (prog == null) {
            Console.WriteLine("Compilation failed: invalid program structure");
            return 1;
        }

        try {
            var emitter = new BytecodeEmitter();
            // Register array types from parser
            foreach (var (name, arrayType) in parser.arrayTypes) {
                emitter.RegisterArrayType(name, arrayType);
            }
            emitter.Compile(prog);
            File.WriteAllBytes(args[1], emitter.BuildBinary());
            Console.WriteLine($"Wrote {args[1]}: consts={emitter.ConstantsCount}, vars={emitter.VarCount}, code={emitter.CodeLength}");
            return 0;
        } catch (Exception ex) {
            Console.WriteLine($"Internal compiler error: {ex.Message}");
            return 1;
        }
    }
}

// Parser, Lexer und AST wurden ausgelagert nach Parser.cs, Lexer.cs und Ast.cs
// Die Implementierungen dort werden vom Entry-Point genutzt.

// === BYTECODE ===
public class BytecodeEmitter {
    List<object> consts = new(); // Kann nun int, long, float oder double speichern
    Dictionary<string,Symbol> syms = new();
    List<byte> code = new();
    Dictionary<string,ArrayType> arrayTypes = new();
    
    class Symbol {
        required public string Name;
        public VarType Type;
        public int Index;
        public ArrayType? ArrayInfo; // Null for non-array types
    }
    
    public void RegisterArrayType(string name, ArrayType type) {
        arrayTypes[name] = type;
    }
    
    public int ConstantsCount => consts.Count;
    public int VarCount => syms.Count;
    public int CodeLength => code.Count;

    public void Compile(ProgramNode p){
        int idx=0;
        foreach(var v in p.Vars) {
            var symbol = new Symbol{Name=v.Name, Type=v.Type, Index=idx};
            if (v.Type == VarType.ARRAY) {
                var arrayInfo = arrayTypes[v.Name];
                symbol.ArrayInfo = arrayInfo;
                // For arrays, increment the index by array size
                idx += arrayInfo.Length;
            } else {
                // For regular variables, increment by 1
                idx++;
            }
            syms[v.Name] = symbol;
        }

        foreach(var v in p.Vars)
            if(v.Init!=null){EmitExpr(v.Init);EmitByte(Bytecode.OpCodes.STORE_VAR);EmitU16((ushort)syms[v.Name].Index);} 

        foreach(var s in p.Stmts)
            EmitStmt(s);

        EmitByte(Bytecode.OpCodes.HALT); // Program End
    }

    // Operationscodes für verschiedene Datentypen
    Dictionary<(TokType, VarType), byte> opCodes = new() {
        // Signed integer arithmetic
        {(TokType.PLUS, VarType.SINT), Bytecode.OpCodes.ADD_SINT}, {(TokType.MINUS, VarType.SINT), Bytecode.OpCodes.SUB_SINT},
        {(TokType.MUL, VarType.SINT), Bytecode.OpCodes.MUL_SINT}, {(TokType.DIV, VarType.SINT), Bytecode.OpCodes.DIV_SINT},
        {(TokType.PLUS, VarType.INT), Bytecode.OpCodes.ADD_INT}, {(TokType.MINUS, VarType.INT), Bytecode.OpCodes.SUB_INT},
        {(TokType.MUL, VarType.INT), Bytecode.OpCodes.MUL_INT}, {(TokType.DIV, VarType.INT), Bytecode.OpCodes.DIV_INT},
        {(TokType.PLUS, VarType.DINT), Bytecode.OpCodes.ADD_DINT}, {(TokType.MINUS, VarType.DINT), Bytecode.OpCodes.SUB_DINT},
        {(TokType.MUL, VarType.DINT), Bytecode.OpCodes.MUL_DINT}, {(TokType.DIV, VarType.DINT), Bytecode.OpCodes.DIV_DINT},
        {(TokType.PLUS, VarType.LINT), Bytecode.OpCodes.ADD_LINT}, {(TokType.MINUS, VarType.LINT), Bytecode.OpCodes.SUB_LINT},
        {(TokType.MUL, VarType.LINT), Bytecode.OpCodes.MUL_LINT}, {(TokType.DIV, VarType.LINT), Bytecode.OpCodes.DIV_LINT},

        // Unsigned integer arithmetic
        {(TokType.PLUS, VarType.USINT), Bytecode.OpCodes.ADD_USINT}, {(TokType.MINUS, VarType.USINT), Bytecode.OpCodes.SUB_USINT},
        {(TokType.MUL, VarType.USINT), Bytecode.OpCodes.MUL_USINT}, {(TokType.DIV, VarType.USINT), Bytecode.OpCodes.DIV_USINT},
        {(TokType.PLUS, VarType.UINT), Bytecode.OpCodes.ADD_UINT}, {(TokType.MINUS, VarType.UINT), Bytecode.OpCodes.SUB_UINT},
        {(TokType.MUL, VarType.UINT), Bytecode.OpCodes.MUL_UINT}, {(TokType.DIV, VarType.UINT), Bytecode.OpCodes.DIV_UINT},
        {(TokType.PLUS, VarType.UDINT), Bytecode.OpCodes.ADD_UDINT}, {(TokType.MINUS, VarType.UDINT), Bytecode.OpCodes.SUB_UDINT},
        {(TokType.MUL, VarType.UDINT), Bytecode.OpCodes.MUL_UDINT}, {(TokType.DIV, VarType.UDINT), Bytecode.OpCodes.DIV_UDINT},
        {(TokType.PLUS, VarType.ULINT), Bytecode.OpCodes.ADD_ULINT}, {(TokType.MINUS, VarType.ULINT), Bytecode.OpCodes.SUB_ULINT},
        {(TokType.MUL, VarType.ULINT), Bytecode.OpCodes.MUL_ULINT}, {(TokType.DIV, VarType.ULINT), Bytecode.OpCodes.DIV_ULINT},

        // Floating point arithmetic
        {(TokType.PLUS, VarType.REAL), Bytecode.OpCodes.ADD_REAL}, {(TokType.MINUS, VarType.REAL), Bytecode.OpCodes.SUB_REAL},
        {(TokType.MUL, VarType.REAL), Bytecode.OpCodes.MUL_REAL}, {(TokType.DIV, VarType.REAL), Bytecode.OpCodes.DIV_REAL},
        {(TokType.PLUS, VarType.LREAL), Bytecode.OpCodes.ADD_LREAL}, {(TokType.MINUS, VarType.LREAL), Bytecode.OpCodes.SUB_LREAL},
        {(TokType.MUL, VarType.LREAL), Bytecode.OpCodes.MUL_LREAL}, {(TokType.DIV, VarType.LREAL), Bytecode.OpCodes.DIV_LREAL},

        // Comparisons signed
        {(TokType.LT, VarType.SINT), Bytecode.OpCodes.LT_SINT}, {(TokType.GT, VarType.SINT), Bytecode.OpCodes.GT_SINT},
        {(TokType.LE, VarType.SINT), Bytecode.OpCodes.LE_SINT}, {(TokType.GE, VarType.SINT), Bytecode.OpCodes.GE_SINT},
        {(TokType.EQ, VarType.SINT), Bytecode.OpCodes.EQ_SINT}, {(TokType.NEQ, VarType.SINT), Bytecode.OpCodes.NEQ_SINT},
        {(TokType.LT, VarType.INT), Bytecode.OpCodes.LT_INT}, {(TokType.GT, VarType.INT), Bytecode.OpCodes.GT_INT},
        {(TokType.LE, VarType.INT), Bytecode.OpCodes.LE_INT}, {(TokType.GE, VarType.INT), Bytecode.OpCodes.GE_INT},
        {(TokType.EQ, VarType.INT), Bytecode.OpCodes.EQ_INT}, {(TokType.NEQ, VarType.INT), Bytecode.OpCodes.NEQ_INT},
        {(TokType.LT, VarType.DINT), Bytecode.OpCodes.LT_DINT}, {(TokType.GT, VarType.DINT), Bytecode.OpCodes.GT_DINT},
        {(TokType.LE, VarType.DINT), Bytecode.OpCodes.LE_DINT}, {(TokType.GE, VarType.DINT), Bytecode.OpCodes.GE_DINT},
        {(TokType.EQ, VarType.DINT), Bytecode.OpCodes.EQ_DINT}, {(TokType.NEQ, VarType.DINT), Bytecode.OpCodes.NEQ_DINT},
        {(TokType.LT, VarType.LINT), Bytecode.OpCodes.LT_LINT}, {(TokType.GT, VarType.LINT), Bytecode.OpCodes.GT_LINT},
        {(TokType.LE, VarType.LINT), Bytecode.OpCodes.LE_LINT}, {(TokType.GE, VarType.LINT), Bytecode.OpCodes.GE_LINT},
        {(TokType.EQ, VarType.LINT), Bytecode.OpCodes.EQ_LINT}, {(TokType.NEQ, VarType.LINT), Bytecode.OpCodes.NEQ_LINT},

        // Comparisons unsigned
        {(TokType.LT, VarType.USINT), Bytecode.OpCodes.LT_USINT}, {(TokType.GT, VarType.USINT), Bytecode.OpCodes.GT_USINT},
        {(TokType.LE, VarType.USINT), Bytecode.OpCodes.LE_USINT}, {(TokType.GE, VarType.USINT), Bytecode.OpCodes.GE_USINT},
        {(TokType.EQ, VarType.USINT), Bytecode.OpCodes.EQ_USINT}, {(TokType.NEQ, VarType.USINT), Bytecode.OpCodes.NEQ_USINT},
        {(TokType.LT, VarType.UINT), Bytecode.OpCodes.LT_UINT}, {(TokType.GT, VarType.UINT), Bytecode.OpCodes.GT_UINT},
        {(TokType.LE, VarType.UINT), Bytecode.OpCodes.LE_UINT}, {(TokType.GE, VarType.UINT), Bytecode.OpCodes.GE_UINT},
        {(TokType.EQ, VarType.UINT), Bytecode.OpCodes.EQ_UINT}, {(TokType.NEQ, VarType.UINT), Bytecode.OpCodes.NEQ_UINT},
        {(TokType.LT, VarType.UDINT), Bytecode.OpCodes.LT_UDINT}, {(TokType.GT, VarType.UDINT), Bytecode.OpCodes.GT_UDINT},
        {(TokType.LE, VarType.UDINT), Bytecode.OpCodes.LE_UDINT}, {(TokType.GE, VarType.UDINT), Bytecode.OpCodes.GE_UDINT},
        {(TokType.EQ, VarType.UDINT), Bytecode.OpCodes.EQ_UDINT}, {(TokType.NEQ, VarType.UDINT), Bytecode.OpCodes.NEQ_UDINT},
        {(TokType.LT, VarType.ULINT), Bytecode.OpCodes.LT_ULINT}, {(TokType.GT, VarType.ULINT), Bytecode.OpCodes.GT_ULINT},
        {(TokType.LE, VarType.ULINT), Bytecode.OpCodes.LE_ULINT}, {(TokType.GE, VarType.ULINT), Bytecode.OpCodes.GE_ULINT},
        {(TokType.EQ, VarType.ULINT), Bytecode.OpCodes.EQ_ULINT}, {(TokType.NEQ, VarType.ULINT), Bytecode.OpCodes.NEQ_ULINT},

        // Comparisons floating
        {(TokType.LT, VarType.REAL), Bytecode.OpCodes.LT_REAL}, {(TokType.GT, VarType.REAL), Bytecode.OpCodes.GT_REAL},
        {(TokType.LE, VarType.REAL), Bytecode.OpCodes.LE_REAL}, {(TokType.GE, VarType.REAL), Bytecode.OpCodes.GE_REAL},
        {(TokType.EQ, VarType.REAL), Bytecode.OpCodes.EQ_REAL}, {(TokType.NEQ, VarType.REAL), Bytecode.OpCodes.NEQ_REAL},
        {(TokType.LT, VarType.LREAL), Bytecode.OpCodes.LT_LREAL}, {(TokType.GT, VarType.LREAL), Bytecode.OpCodes.GT_LREAL},
        {(TokType.LE, VarType.LREAL), Bytecode.OpCodes.LE_LREAL}, {(TokType.GE, VarType.LREAL), Bytecode.OpCodes.GE_LREAL},
        {(TokType.EQ, VarType.LREAL), Bytecode.OpCodes.EQ_LREAL}, {(TokType.NEQ, VarType.LREAL), Bytecode.OpCodes.NEQ_LREAL}
    };

    void EmitStmt(Stmt s){
        switch(s){
            case AssignStmt a:
                Symbol? symbol = null;
                if (!syms.TryGetValue(a.Target, out symbol)) {
                    throw new Exception($"Undeclared variable '{a.Target}'");
                }
                
                // If this is an array assignment, we need to evaluate the index and adjust the target address
                if (symbol.Type == VarType.ARRAY && a.Index != null) {
                    // Evaluate array index
                    EmitExpr(a.Index);
                    
                    // Bounds check
                    if (symbol.ArrayInfo == null) 
                        throw new Exception($"Internal error: Array info missing for '{a.Target}'");
                    
                    EmitByte(Bytecode.OpCodes.ARRAY_BOUNDS_CHECK);
                    EmitU16((ushort)symbol.ArrayInfo.Start);
                    EmitU16((ushort)symbol.ArrayInfo.End);
                    
                    // Convert to zero-based index
                    if (symbol.ArrayInfo.Start != 0) {
                        EmitByte(Bytecode.OpCodes.PUSH_CONST);
                        EmitU16((ushort)AddConst(symbol.ArrayInfo.Start));
                        EmitByte(Bytecode.OpCodes.SUB_INT);
                    }
                    
                    // Add base address
                    EmitByte(Bytecode.OpCodes.PUSH_CONST);
                    EmitU16((ushort)symbol.Index);
                    EmitByte(Bytecode.OpCodes.ADD_INT);
                    
                    // Now evaluate and store the value
                    EmitExpr(a.Expr);
                    EmitByte(Bytecode.OpCodes.STORE_VAR);
                    EmitU16(0); // Index is already on stack
                }
                else {
                    // Regular variable assignment
                    EmitExpr(a.Expr);
                    EmitByte(Bytecode.OpCodes.STORE_VAR);
                    EmitU16((ushort)symbol.Index);
                }
                break;

            case IfStmt iff:
                EmitExpr(iff.Cond);
                EmitByte(Bytecode.OpCodes.JZ); int jz=code.Count; EmitU16(0);
                foreach(var st in iff.ThenStmts) EmitStmt(st);
                if(iff.ElseStmts.Count>0){
                    EmitByte(Bytecode.OpCodes.JMP); int jmp=code.Count; EmitU16(0);
                    PatchJump(jz,code.Count);
                    foreach(var st in iff.ElseStmts) EmitStmt(st);
                    PatchJump(jmp,code.Count);
                } else PatchJump(jz,code.Count);
                break;

            case WhileStmt w:
                int loopStart=code.Count;
                EmitExpr(w.Cond);
                EmitByte(Bytecode.OpCodes.JZ); int jzpos=code.Count; EmitU16(0);
                foreach(var st in w.Body) EmitStmt(st);
                EmitByte(Bytecode.OpCodes.JMP); EmitU16((ushort)loopStart);
                PatchJump(jzpos,code.Count);
                break;

            case ForStmt f:
                Symbol? forSymbol = null;
                if (!syms.TryGetValue(f.Var, out forSymbol)) {
                    throw new Exception($"Undeclared variable '{f.Var}'");
                }
                // Initialisierung: var := start
                EmitExpr(f.Start);
                EmitByte(Bytecode.OpCodes.STORE_VAR); EmitU16((ushort)forSymbol.Index);

                int cmpPos = code.Count;  // Position des Vergleichs
                EmitExpr(new VarExpr(f.Var, forSymbol.Type));
                EmitExpr(f.End);
                var key = (TokType.LE, forSymbol.Type);
                if (!opCodes.ContainsKey(key)) {
                    throw new Exception($"Comparison not supported for type {forSymbol.Type}");
                }
                EmitByte(opCodes[key]);
                EmitByte(Bytecode.OpCodes.JZ); int jzFor = code.Count; EmitU16(0); // Jump zum Ende

                // Body
                foreach(var st in f.Body) EmitStmt(st);

                // Inkrement: i := i + step
                EmitExpr(new VarExpr(f.Var, forSymbol.Type));
                EmitExpr(f.Step);
                var addKey = (TokType.PLUS, forSymbol.Type);
                if (!opCodes.ContainsKey(addKey)) {
                    throw new Exception($"Addition not supported for type {forSymbol.Type}");
                }
                EmitByte(opCodes[addKey]);
                EmitByte(Bytecode.OpCodes.STORE_VAR); EmitU16((ushort)forSymbol.Index);

                // Vergleich erneut, aber wir merken uns den Sprung zum Vergleich
                EmitByte(Bytecode.OpCodes.JMP); EmitU16((ushort)cmpPos);

                // Patch Jump: Ende der Schleife springt hierhin
                PatchJump(jzFor, code.Count);

                // Korrektur: Schleifenvariable auf Endwert, falls sie überläuft
                EmitExpr(f.End);
                EmitByte(Bytecode.OpCodes.STORE_VAR); EmitU16((ushort)syms[f.Var].Index);
                break;
        }
    }

    void PatchJump(int pos,int target){code[pos]=(byte)(target&0xFF);code[pos+1]=(byte)(target>>8);}    
    void EmitExpr(Expr e){
        switch(e){
            case IntExpr ie:
                int ci = AddConst(ie.Value);
                EmitByte(Bytecode.OpCodes.PUSH_CONST); // PUSH_CONST
                EmitU16((ushort)ci);
                break;
                
            case RealExpr re:
                int cri = AddConst(re.Value);
                EmitByte(Bytecode.OpCodes.PUSH_REAL_CONST); // PUSH_REAL_CONST
                EmitU16((ushort)cri);
                break;
                
            case VarExpr ve:
                Symbol? symbol = null;
                if (!syms.TryGetValue(ve.Name, out symbol)) {
                    throw new Exception($"Undeclared variable '{ve.Name}'");
                }
                EmitByte(Bytecode.OpCodes.LOAD_VAR); // LOAD_VAR
                EmitU16((ushort)symbol.Index);
                break;

            case ArrayAccessExpr ae:
                Symbol? arraySymbol = null;
                if (!syms.TryGetValue(ae.ArrayName, out arraySymbol)) {
                    throw new Exception($"Undeclared array '{ae.ArrayName}'");
                }
                if (arraySymbol.Type != VarType.ARRAY || arraySymbol.ArrayInfo == null) {
                    throw new Exception($"Variable '{ae.ArrayName}' is not an array");
                }

                // Emit the index expression first
                EmitExpr(ae.Index);

                // Emit array bounds check
                int start = arraySymbol.ArrayInfo.Start;
                int end = arraySymbol.ArrayInfo.End;
                EmitByte(Bytecode.OpCodes.ARRAY_BOUNDS_CHECK);
                EmitU16((ushort)start);
                EmitU16((ushort)end);

                // Convert index to zero-based
                if (start != 0) {
                    EmitByte(Bytecode.OpCodes.PUSH_CONST);
                    EmitU16((ushort)AddConst(start));
                    EmitByte(Bytecode.OpCodes.SUB_INT);
                }

                // Load array base address
                EmitByte(Bytecode.OpCodes.PUSH_CONST);
                EmitU16((ushort)arraySymbol.Index);
                EmitByte(Bytecode.OpCodes.ADD_INT);

                // Load array element
                EmitByte(Bytecode.OpCodes.LOAD_VAR);
                EmitU16(0); // Index is already on stack
                break;
                
            case BinaryExpr be:
                EmitExpr(be.L);
                EmitExpr(be.R);
                
                var key = (be.Op, be.Type);
                if (!opCodes.ContainsKey(key)) {
                    throw new Exception($"Unknown operator '{be.Op}' for type {be.Type}");
                }
                EmitByte(opCodes[key]);
                break;
        }
    }

    int AddConst(object v) {
        int i = consts.FindIndex(c => c.Equals(v));
        if (i >= 0) return i;
        consts.Add(v);
        return consts.Count - 1;
    }
    
    void EmitByte(byte b) => code.Add(b);
    void EmitU16(ushort v) {code.Add((byte)(v&0xFF));code.Add((byte)(v>>8));}
    
    public byte[] BuildBinary() {
        using var ms = new MemoryStream();
        var w = new BinaryWriter(ms);
        
        // Header
        w.Write(Encoding.ASCII.GetBytes(Bytecode.Magic));
        w.Write(Bytecode.Version);
        w.Write(Bytecode.DefaultCycleTime);

        // Konstanten
        w.Write((ushort)consts.Count);
        foreach(var c in consts) {
            if (c is long l) {
                w.Write((byte)1); // Long type marker
                w.Write(l);
            }
            else if (c is double d) {
                w.Write((byte)2); // Double type marker
                w.Write(d);
            }
            else if (c is float f) {
                w.Write((byte)3); // Float type marker
                w.Write(f);
            }
            else if (c is int i) {
                w.Write((byte)4); // Int type marker
                w.Write(i);
            }
            else {
                throw new Exception($"Unsupported constant type: {c.GetType()}");
            }
        }
        
        // Calculate total variable space needed (accounting for arrays)
        int totalVarSpace = 0;
        foreach(var kv in syms) {
            if (kv.Value.Type == VarType.ARRAY && kv.Value.ArrayInfo != null) {
                totalVarSpace = Math.Max(totalVarSpace, kv.Value.Index + kv.Value.ArrayInfo.Length);
            } else {
                totalVarSpace = Math.Max(totalVarSpace, kv.Value.Index + 1);
            }
        }

        // Variablen
        w.Write((ushort)totalVarSpace);
        var types = new byte[totalVarSpace];
        foreach(var kv in syms) {
            if (kv.Value.Type == VarType.ARRAY && kv.Value.ArrayInfo != null) {
                // Fill array elements with their element type
                for (int i = 0; i < kv.Value.ArrayInfo.Length; i++) {
                    types[kv.Value.Index + i] = (byte)kv.Value.ArrayInfo.ElementType;
                }
            } else {
                types[kv.Value.Index] = (byte)kv.Value.Type;
            }
        }
        foreach(var b in types)
            w.Write(b);
            
        // Code
        w.Write((ushort)code.Count);
        w.Write(code.ToArray());
        
        return ms.ToArray();
    }
}