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

// === 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;
        }
    }
}

// === AST ===
public enum VarType { Int8=1, Int16=2, Int32=3, Byte=4, Bool=5 }
public abstract class StNode{}
public class ProgramNode:StNode{ public List<VarDecl> Vars=new(); public List<Stmt> Stmts=new(); }
public class VarDecl:StNode{ 
    required public string Name; 
    public VarType Type; 
    public Expr? Init; 
}

public abstract class Stmt:StNode{}
public class AssignStmt:Stmt{ 
    required public string Target; 
    required public Expr Expr; 
}
public class IfStmt:Stmt{ 
    required public Expr Cond; 
    public List<Stmt> ThenStmts=new(); 
    public List<Stmt> ElseStmts=new(); 
}
public class WhileStmt:Stmt{ 
    required public Expr Cond; 
    public List<Stmt> Body=new(); 
}
public class ForStmt:Stmt{ 
    required public string Var; 
    required public Expr Start; 
    required public Expr End; 
    public Expr Step = new IntExpr(1); 
    public List<Stmt> Body=new(); 
}

public abstract class Expr:StNode{}
public class IntExpr:Expr{ public int Value; public IntExpr(int v){Value=v;} }
public class VarExpr:Expr{ public string Name; public VarExpr(string n){Name=n;} }
public class BinaryExpr:Expr{ public Expr L; public Expr R; public TokType Op; public BinaryExpr(Expr l,TokType op,Expr r){L=l;Op=op;R=r;} }

// === TOKENIZER ===
public enum TokType {
    IDENT, INT, ASSIGN, SEMI, LPAREN, RPAREN,
    PLUS, MINUS, MUL, DIV,
    LT, GT, LE, GE, EQ, NEQ,
    IF, THEN, ELSE, END_IF,
    WHILE, DO, END_WHILE,
    FOR, TO, BY, END_FOR,
    PROGRAM, VAR, END_VAR, END_PROGRAM,
    EOF
}
public class Token{ 
    public TokType Type; 
    public string Text; 
    public int Line;
    public Token(TokType t, string s, int line) { Type=t; Text=s; Line=line; }
}

public class CompileError {
    public int Line;
    public string Message;
    public CompileError(int line, string msg) { Line=line; Message=msg; }
}

public class StLexer {
    private readonly string src;
    private int i;
    private int currentLine = 1;
    public List<CompileError> Errors = new();
    public StLexer(string s){src=s;}
    char Peek()=> i<src.Length?src[i]:'\0';
    char Next(){ 
        if (i >= src.Length) return '\0';
        char c = src[i++];
        if (c == '\n') currentLine++;
        return c;
    }

    void AddError(string msg) => Errors.Add(new CompileError(currentLine, msg));

    public Token NextToken() {
        while (char.IsWhiteSpace(Peek())) Next();
        if (Peek()=='\0') return new Token(TokType.EOF,"", currentLine);

        if (char.IsLetter(Peek())||Peek()=='_'){
            var sb=new StringBuilder();
            int startLine = currentLine;
            while (char.IsLetterOrDigit(Peek())||Peek()=='_') sb.Append(Next());
            var s=sb.ToString().ToUpperInvariant();
            return s switch {
                "PROGRAM"=>new Token(TokType.PROGRAM,s,startLine),
                "VAR"=>new Token(TokType.VAR,s,startLine),
                "END_VAR"=>new Token(TokType.END_VAR,s,startLine),
                "END_PROGRAM"=>new Token(TokType.END_PROGRAM,s,startLine),
                "IF"=>new Token(TokType.IF,s,startLine),
                "THEN"=>new Token(TokType.THEN,s,startLine),
                "ELSE"=>new Token(TokType.ELSE,s,startLine),
                "END_IF"=>new Token(TokType.END_IF,s,startLine),
                "WHILE"=>new Token(TokType.WHILE,s,startLine),
                "DO"=>new Token(TokType.DO,s,startLine),
                "END_WHILE"=>new Token(TokType.END_WHILE,s,startLine),
                "FOR"=>new Token(TokType.FOR,s,startLine),
                "TO"=>new Token(TokType.TO,s,startLine),
                "BY"=>new Token(TokType.BY,s,startLine),
                "END_FOR"=>new Token(TokType.END_FOR,s,startLine),
                _=>new Token(TokType.IDENT,s,startLine)
            };
        }

        if (char.IsDigit(Peek())) {
            var sb=new StringBuilder();
            int startLine = currentLine;
            while(char.IsDigit(Peek())) sb.Append(Next());
            return new Token(TokType.INT,sb.ToString(),startLine);
        }

        int tokenLine = currentLine;
        if (Peek()==':'){ 
            Next(); 
            if(Peek()=='='){
                Next(); 
                return new Token(TokType.ASSIGN,":=",tokenLine);
            } 
            AddError("Expected '=' after ':' for assignment");
            // Bei einem einzelnen ':' geben wir EOF zurück und stoppen das Parsen
            i--; // Gehen einen Schritt zurück, damit der fehlerhafte ':' Token beim nächsten Mal neu gelesen wird
            return new Token(TokType.EOF,"",tokenLine);
        }
        if (Peek()=='<'){
            Next();
            if (Peek()=='='){Next(); return new Token(TokType.LE,"<=",tokenLine);}
            if (Peek()=='>'){Next(); return new Token(TokType.NEQ,"<>",tokenLine);}
            return new Token(TokType.LT,"<",tokenLine);
        }
        if (Peek()=='>'){
            Next();
            if (Peek()=='='){Next(); return new Token(TokType.GE,">=",tokenLine);}
            return new Token(TokType.GT,">",tokenLine);
        }
        if (Peek()=='='){Next();return new Token(TokType.EQ,"=",tokenLine);}

        char c=Next();
        if (c == ';') return new Token(TokType.SEMI,";",tokenLine);
        if (c == '(') return new Token(TokType.LPAREN,"(",tokenLine);
        if (c == ')') return new Token(TokType.RPAREN,")",tokenLine);
        if (c == '+') return new Token(TokType.PLUS,"+",tokenLine);
        if (c == '-') return new Token(TokType.MINUS,"-",tokenLine);
        if (c == '*') return new Token(TokType.MUL,"*",tokenLine);
        if (c == '/') return new Token(TokType.DIV,"/",tokenLine);
        
        AddError($"Unexpected character '{c}'");
        return new Token(TokType.EOF,"",tokenLine); // Skip invalid character
    }
}

// === PARSER ===
public class StParser {
    StLexer lex; 
    Token cur;
    public List<CompileError> Errors => lex.Errors;
    public bool HasErrors => Errors.Count > 0;

    public StParser(string s){
        lex=new StLexer(s);
        cur=lex.NextToken();
    }

    void Next()=>cur=lex.NextToken();

    void AddError(string msg) => lex.Errors.Add(new CompileError(cur.Line, msg));

    bool Expect(TokType t){
        if(cur.Type!=t) {
            AddError($"Expected {t}, got {cur.Type}");
            return false;
        }
        Next();
        return true;
    }

    public ProgramNode? ParseProgram(){
        var p=new ProgramNode();
        if (!Expect(TokType.PROGRAM)) return null;
        
        if(cur.Type==TokType.IDENT) Next();
        
        if(cur.Type==TokType.VAR){
            Next();
            while(cur.Type==TokType.IDENT) {
                var varDecl = ParseVarDecl();
                if (varDecl == null) return null;
                p.Vars.Add(varDecl);
            }
            if (!Expect(TokType.END_VAR)) return null;
        }
        
        while(cur.Type!=TokType.END_PROGRAM && cur.Type!=TokType.EOF) {
            var stmt = ParseStmt();
            if (stmt == null) return null;
            p.Stmts.Add(stmt);
        }
        
        if (!Expect(TokType.END_PROGRAM)) return null;
        return p;
    }

    VarDecl? ParseVarDecl(){
        if (cur.Type != TokType.IDENT) {
            AddError("Expected identifier for variable declaration");
            return null;
        }
        string name=cur.Text; 
        Next();

        if (cur.Type != TokType.IDENT) {
            AddError("Expected type name");
            return null;
        }
        string tname=cur.Text.ToLowerInvariant(); 
        Next();

        VarType? vt = tname switch {
            "int8" => VarType.Int8,
            "int16" => VarType.Int16,
            "int32" => VarType.Int32,
            "byte" => VarType.Byte,
            "bool" => VarType.Bool,
            _ => null
        };
        
        if (vt == null) {
            AddError($"Unknown type '{tname}'");
            return null;
        }

        Expr? init=null;
        if(cur.Type==TokType.ASSIGN){
            Next(); // consume :=
            init=ParseExpr();
            if (init == null) {
                AddError($"Expected expression after ':=' in variable declaration");
                return null;
            }
        }
        
        if (!Expect(TokType.SEMI)) return null;
        
        return new VarDecl{Name=name,Type=vt.Value,Init=init};
    }

    Stmt? ParseAssign() {
        // Der Aufrufer hat bereits geprüft, dass wir bei einem IDENT sind
        string target = cur.Text;
        Next(); // consume identifier
        if (cur.Type != TokType.ASSIGN) {
            AddError($"Expected ':=' after identifier '{target}'");
            return null;
        }
        Next(); // consume :=
        var e = ParseExpr();
        if (e == null) return null;
        if (!Expect(TokType.SEMI)) return null;
        return new AssignStmt{Target=target, Expr=e};
    }

    Stmt? ParseStmt(){
        switch(cur.Type) {
            case TokType.IF: return ParseIf();
            case TokType.WHILE: return ParseWhile();
            case TokType.FOR: return ParseFor();
            case TokType.IDENT:
                return ParseAssign();
            default:
                AddError($"Unexpected token {cur.Type} in statement");
                return null;
        }
    }

    IfStmt? ParseIf(){
        Next(); // IF
        var cond=ParseExpr(); 
        if (cond == null) return null;
        if (!Expect(TokType.THEN)) return null;

        var node=new IfStmt{Cond=cond};
        while(cur.Type!=TokType.ELSE && cur.Type!=TokType.END_IF && cur.Type!=TokType.EOF) {
            var stmt = ParseStmt();
            if (stmt == null) return null;
            node.ThenStmts.Add(stmt);
        }

        if(cur.Type==TokType.ELSE){
            Next();
            while(cur.Type!=TokType.END_IF && cur.Type!=TokType.EOF) {
                var stmt = ParseStmt();
                if (stmt == null) return null;
                node.ElseStmts.Add(stmt);
            }
        }

        if (!Expect(TokType.END_IF)) return null;
        if (!Expect(TokType.SEMI)) return null;
        return node;
    }

    WhileStmt? ParseWhile(){
        Next(); // WHILE
        var cond=ParseExpr(); 
        if (cond == null) return null;
        if (!Expect(TokType.DO)) return null;

        var ws=new WhileStmt{Cond=cond};
        while(cur.Type!=TokType.END_WHILE && cur.Type!=TokType.EOF) {
            var stmt = ParseStmt();
            if (stmt == null) return null;
            ws.Body.Add(stmt);
        }

        if (!Expect(TokType.END_WHILE)) return null;
        if (!Expect(TokType.SEMI)) return null;
        return ws;
    }

    ForStmt? ParseFor(){
        Next(); // FOR
        if (cur.Type != TokType.IDENT) {
            AddError("Expected identifier for FOR loop variable");
            return null;
        }
        string varName = cur.Text;
        Next(); // consume identifier
        
        if (cur.Type != TokType.ASSIGN) {
            AddError($"Expected ':=' after identifier '{varName}'");
            return null;
        }
        Next(); // consume :=
        var start = ParseExpr();
        if (start == null) return null;

        if (!Expect(TokType.TO)) return null;
        var end = ParseExpr();
        if (end == null) return null;

        Expr step = new IntExpr(1);
        if(cur.Type==TokType.BY){
            Next();
            step = ParseExpr() ?? step;
        }

        if (!Expect(TokType.DO)) return null;

        var fs = new ForStmt{Var=varName, Start=start, End=end, Step=step};
        while(cur.Type!=TokType.END_FOR && cur.Type!=TokType.EOF) {
            var stmt = ParseStmt();
            if (stmt == null) return null;
            fs.Body.Add(stmt);
        }

        if (!Expect(TokType.END_FOR)) return null;
        if (!Expect(TokType.SEMI)) return null;
        return fs;
    }

    Expr? ParseExpr() => ParseCompare();
    
    Expr? ParseCompare(){
        var l = ParseAddSub();
        if (l == null) return null;
        
        while(cur.Type is TokType.LT or TokType.GT or TokType.LE or TokType.GE or TokType.EQ or TokType.NEQ){
            var op=cur.Type; 
            Next();
            var r=ParseAddSub();
            if (r == null) return null;
            l=new BinaryExpr(l,op,r);
        }
        return l;
    }
    
    Expr? ParseAddSub(){
        var l=ParseMulDiv();
        if (l == null) return null;
        
        while(cur.Type==TokType.PLUS||cur.Type==TokType.MINUS){
            var op=cur.Type;
            Next();
            var r=ParseMulDiv();
            if (r == null) return null;
            l=new BinaryExpr(l,op,r);
        }
        return l;
    }
    
    Expr? ParseMulDiv(){
        var l=ParsePrimary();
        if (l == null) return null;
        
        while(cur.Type==TokType.MUL||cur.Type==TokType.DIV){
            var op=cur.Type;
            Next();
            var r=ParsePrimary();
            if (r == null) return null;
            l=new BinaryExpr(l,op,r);
        }
        return l;
    }
    
    Expr? ParsePrimary(){
        int startLine = cur.Line;
        switch(cur.Type) {
            case TokType.INT:
                if (!int.TryParse(cur.Text, out var v)) {
                    AddError($"Invalid integer literal '{cur.Text}'");
                    return null;
                }
                Next();
                return new IntExpr(v);
                
            case TokType.IDENT:
                string n=cur.Text;
                Next();
                return new VarExpr(n);
                
            case TokType.LPAREN:
                Next();
                var e=ParseExpr();
                if (e == null) return null;
                if (!Expect(TokType.RPAREN)) return null;
                return e;
                
            default:
                AddError($"Unexpected token {cur.Type} in expression");
                return null;
        }
    }
    }


// === BYTECODE ===
public class BytecodeEmitter {
    List<int> consts=new(); 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(0x03);EmitU16((ushort)syms[v.Name].Index);}

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

        EmitByte(0xF0); // Program End
    }

    Dictionary<TokType, byte> opCodes = new() {
        {TokType.PLUS, 0x10}, {TokType.MINUS, 0x11}, {TokType.MUL, 0x12}, {TokType.DIV, 0x13},
        {TokType.LT, 0x14}, {TokType.GT, 0x15}, {TokType.LE, 0x16}, {TokType.GE, 0x17},
        {TokType.EQ, 0x18}, {TokType.NEQ, 0x19}
    };

    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(0x03); EmitU16((ushort)symbol.Index);
                break;

            case IfStmt iff:
                EmitExpr(iff.Cond);
                EmitByte(0x20); int jz=code.Count; EmitU16(0);
                foreach(var st in iff.ThenStmts) EmitStmt(st);
                if(iff.ElseStmts.Count>0){
                    EmitByte(0x21); 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(0x20); int jzpos=code.Count; EmitU16(0);
                foreach(var st in w.Body) EmitStmt(st);
                EmitByte(0x21); 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(0x03); EmitU16((ushort)forSymbol.Index);

                int cmpPos = code.Count;  // Position des Vergleichs
                EmitExpr(new VarExpr(f.Var));
                EmitExpr(f.End);
                EmitByte(0x16); // LE
                EmitByte(0x20); 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));
                EmitExpr(f.Step);
                EmitByte(0x10); // PLUS
                EmitByte(0x03); EmitU16((ushort)syms[f.Var].Index);

                // Vergleich erneut, aber wir merken uns den Sprung zum Vergleich
                EmitByte(0x21); 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(0x03); 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(0x01);EmitU16((ushort)ci);break;
            case VarExpr ve:
                Symbol? symbol = null;
                if (!syms.TryGetValue(ve.Name, out symbol)) {
                    throw new Exception($"Undeclared variable '{ve.Name}'");
                }
                EmitByte(0x02);EmitU16((ushort)symbol.Index);break;
            case BinaryExpr be:
                if (!opCodes.ContainsKey(be.Op)) {
                    throw new Exception($"Unknown operator '{be.Op}'");
                }
                EmitExpr(be.L);EmitExpr(be.R);
                EmitByte(opCodes[be.Op]);
                break;
        }
    }

    int AddConst(int v){int i=consts.IndexOf(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);
        w.Write(Encoding.ASCII.GetBytes("STBC"));
        w.Write((ushort)1);
        w.Write((ushort)consts.Count);foreach(var c in consts)w.Write(c);
        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);
        w.Write((ushort)code.Count);w.Write(code.ToArray());
        return ms.ToArray();
    }
}