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();
            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();
    
    class Symbol {
        required public string Name;
        public VarType Type;
        public int Index;
    }
    
    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)
            syms[v.Name]=new Symbol{Name=v.Name,Type=v.Type,Index=idx++};

        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_S}, {(TokType.GT, VarType.SINT), Bytecode.OpCodes.GT_S},
        {(TokType.LE, VarType.SINT), Bytecode.OpCodes.LE_S}, {(TokType.GE, VarType.SINT), Bytecode.OpCodes.GE_S},
        {(TokType.EQ, VarType.SINT), Bytecode.OpCodes.EQ_S}, {(TokType.NEQ, VarType.SINT), Bytecode.OpCodes.NEQ_S},
        {(TokType.LT, VarType.INT), Bytecode.OpCodes.LT_S}, {(TokType.GT, VarType.INT), Bytecode.OpCodes.GT_S},
        {(TokType.LE, VarType.INT), Bytecode.OpCodes.LE_S}, {(TokType.GE, VarType.INT), Bytecode.OpCodes.GE_S},
        {(TokType.EQ, VarType.INT), Bytecode.OpCodes.EQ_S}, {(TokType.NEQ, VarType.INT), Bytecode.OpCodes.NEQ_S},
        {(TokType.LT, VarType.DINT), Bytecode.OpCodes.LT_S}, {(TokType.GT, VarType.DINT), Bytecode.OpCodes.GT_S},
        {(TokType.LE, VarType.DINT), Bytecode.OpCodes.LE_S}, {(TokType.GE, VarType.DINT), Bytecode.OpCodes.GE_S},
        {(TokType.EQ, VarType.DINT), Bytecode.OpCodes.EQ_S}, {(TokType.NEQ, VarType.DINT), Bytecode.OpCodes.NEQ_S},
        {(TokType.LT, VarType.LINT), Bytecode.OpCodes.LT_S}, {(TokType.GT, VarType.LINT), Bytecode.OpCodes.GT_S},
        {(TokType.LE, VarType.LINT), Bytecode.OpCodes.LE_S}, {(TokType.GE, VarType.LINT), Bytecode.OpCodes.GE_S},
        {(TokType.EQ, VarType.LINT), Bytecode.OpCodes.EQ_S}, {(TokType.NEQ, VarType.LINT), Bytecode.OpCodes.NEQ_S},

        // Comparisons unsigned
        {(TokType.LT, VarType.USINT), Bytecode.OpCodes.LT_U}, {(TokType.GT, VarType.USINT), Bytecode.OpCodes.GT_U},
        {(TokType.LE, VarType.USINT), Bytecode.OpCodes.LE_U}, {(TokType.GE, VarType.USINT), Bytecode.OpCodes.GE_U},
        {(TokType.EQ, VarType.USINT), Bytecode.OpCodes.EQ_U}, {(TokType.NEQ, VarType.USINT), Bytecode.OpCodes.NEQ_U},
        {(TokType.LT, VarType.UINT), Bytecode.OpCodes.LT_U}, {(TokType.GT, VarType.UINT), Bytecode.OpCodes.GT_U},
        {(TokType.LE, VarType.UINT), Bytecode.OpCodes.LE_U}, {(TokType.GE, VarType.UINT), Bytecode.OpCodes.GE_U},
        {(TokType.EQ, VarType.UINT), Bytecode.OpCodes.EQ_U}, {(TokType.NEQ, VarType.UINT), Bytecode.OpCodes.NEQ_U},
        {(TokType.LT, VarType.UDINT), Bytecode.OpCodes.LT_U}, {(TokType.GT, VarType.UDINT), Bytecode.OpCodes.GT_U},
        {(TokType.LE, VarType.UDINT), Bytecode.OpCodes.LE_U}, {(TokType.GE, VarType.UDINT), Bytecode.OpCodes.GE_U},
        {(TokType.EQ, VarType.UDINT), Bytecode.OpCodes.EQ_U}, {(TokType.NEQ, VarType.UDINT), Bytecode.OpCodes.NEQ_U},
        {(TokType.LT, VarType.ULINT), Bytecode.OpCodes.LT_U}, {(TokType.GT, VarType.ULINT), Bytecode.OpCodes.GT_U},
        {(TokType.LE, VarType.ULINT), Bytecode.OpCodes.LE_U}, {(TokType.GE, VarType.ULINT), Bytecode.OpCodes.GE_U},
        {(TokType.EQ, VarType.ULINT), Bytecode.OpCodes.EQ_U}, {(TokType.NEQ, VarType.ULINT), Bytecode.OpCodes.NEQ_U},

        // Comparisons floating
        {(TokType.LT, VarType.REAL), Bytecode.OpCodes.LT_F}, {(TokType.GT, VarType.REAL), Bytecode.OpCodes.GT_F},
        {(TokType.LE, VarType.REAL), Bytecode.OpCodes.LE_F}, {(TokType.GE, VarType.REAL), Bytecode.OpCodes.GE_F},
        {(TokType.EQ, VarType.REAL), Bytecode.OpCodes.EQ_F}, {(TokType.NEQ, VarType.REAL), Bytecode.OpCodes.NEQ_F},
        {(TokType.LT, VarType.LREAL), Bytecode.OpCodes.LT_F}, {(TokType.GT, VarType.LREAL), Bytecode.OpCodes.GT_F},
        {(TokType.LE, VarType.LREAL), Bytecode.OpCodes.LE_F}, {(TokType.GE, VarType.LREAL), Bytecode.OpCodes.GE_F},
        {(TokType.EQ, VarType.LREAL), Bytecode.OpCodes.EQ_F}, {(TokType.NEQ, VarType.LREAL), Bytecode.OpCodes.NEQ_F}
    };

    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}'");
                }
                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);
                // removed emitting the type byte into code stream; type is stored with constant table
                break;
                
            case RealExpr re:
                int cri = AddConst(re.Value);
                EmitByte(Bytecode.OpCodes.PUSH_REAL_CONST); // PUSH_REAL_CONST
                EmitU16((ushort)cri);
                // removed emitting the type byte into code stream; type is stored with constant table
                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 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}");
                }
                
                // Wenn nötig, Typenkonvertierung vor der Operation
                // Conversion opcodes removed - runtime will handle type differences using constant type markers and operation opcodes
                // if (be.L.Type != be.Type)
                //     EmitByte((byte)(0x50 + (int)be.Type)); // CONVERT_TO_*
                // if (be.R.Type != be.Type)
                //     EmitByte((byte)(0x50 + (int)be.Type)); // CONVERT_TO_*
                    
                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);

        // 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()}");
            }
        }
        
        // Variablen
        w.Write((ushort)syms.Count);
        var types = new byte[syms.Count];
        foreach(var kv in syms)
            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();
    }
}