MicroST-Compiler/STCompiler.Compiler/Parser.cs

namespace STCompiler.Compiler;

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

public record Symbol {
    required public string Name;
    public VarType Type;
    public int Index;
}

public class StParser {
    StLexer lex; 
    Token cur;
    Dictionary<string,Symbol> syms = new();
    public Dictionary<string,ArrayType> arrayTypes = new(); // Add dictionary for array types
    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;
    }

    VarType? ParseType(string varName) {
        if (cur.Type == TokType.ARRAY) {
            Next(); // Skip ARRAY
            if (!Expect(TokType.LBRACKET)) return null;
            var start = ParseExpr();
            if (start == null || start is not IntExpr startInt) {
                AddError("Array start index must be an integer constant");
                return null;
            }
            if (!Expect(TokType.DOTS)) return null;
            var end = ParseExpr();
            if (end == null || end is not IntExpr endInt) {
                AddError("Array end index must be an integer constant");
                return null;
            }
            if (!Expect(TokType.RBRACKET)) return null;
            if (!Expect(TokType.OF)) return null;
            if (cur.Type != TokType.IDENT) {
                AddError("Expected element type name after OF");
                return null;
            }
            string elemTypeName = cur.Text.ToLowerInvariant();
            Next();

            VarType? elementType = elemTypeName switch {
                "bool" => VarType.BOOL,
                "byte" => VarType.BYTE,
                "word" => VarType.WORD,
                "dword" => VarType.DWORD,
                "lword" => VarType.LWORD,
                "sint" => VarType.SINT,
                "int" => VarType.INT,
                "dint" => VarType.DINT,
                "lint" => VarType.LINT,
                "usint" => VarType.USINT,
                "uint" => VarType.UINT,
                "udint" => VarType.UDINT,
                "ulint" => VarType.ULINT,
                "real" => VarType.REAL,
                "lreal" => VarType.LREAL,
                _ => null
            };

            if (elementType == null) {
                AddError($"Unknown element type '{elemTypeName}'");
                return null;
            }

            arrayTypes[varName] = new ArrayType(elementType.Value, (int)startInt.Value, (int)endInt.Value);
            return VarType.ARRAY;
        }
        else if (cur.Type == TokType.IDENT) {
            string typeName = cur.Text.ToLowerInvariant();
            Next();

            return typeName switch {
                "bool" => VarType.BOOL,
                "byte" => VarType.BYTE,
                "word" => VarType.WORD,
                "dword" => VarType.DWORD,
                "lword" => VarType.LWORD,
                "sint" => VarType.SINT,
                "int" => VarType.INT,
                "dint" => VarType.DINT,
                "lint" => VarType.LINT,
                "usint" => VarType.USINT,
                "uint" => VarType.UINT,
                "udint" => VarType.UDINT,
                "ulint" => VarType.ULINT,
                "real" => VarType.REAL,
                "lreal" => VarType.LREAL,
                _ => null
            };
        }
        
        AddError("Expected type name or ARRAY");
        return null;
    }

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

        var vt = ParseType(name);
        if (vt == null) {
            AddError($"Invalid type for variable '{name}'");
            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;
        
        // Füge Variable zur Symboltabelle hinzu
        syms[name] = new Symbol { Name = name, Type = vt.Value };
        
        return new VarDecl{Name=name,Type=vt.Value,Init=init};
    }

    Stmt? ParseAssign() {
        string target = cur.Text.ToUpperInvariant();
        Next(); // consume identifier

        // Check if this is an array assignment
        Expr? index = null;
        if (cur.Type == TokType.LBRACKET) {
            Symbol? sym;
            if (!syms.TryGetValue(target, out sym)) {
                AddError($"Undeclared variable '{target}'");
                return null;
            }
            if (sym.Type != VarType.ARRAY) {
                AddError($"Cannot use array indexing on non-array variable '{target}'");
                return null;
            }

            Next(); // consume [
            index = ParseExpr();
            if (index == null) return null;
            if (!Expect(TokType.RBRACKET)) return null;
        }

        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, Index=index, 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.ToUpperInvariant(); // Variablennamen in Großbuchstaben
        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;
        // Support unary + and -
        if (cur.Type == TokType.PLUS || cur.Type == TokType.MINUS) {
            var sign = cur.Type;
            Next();
            var p = ParsePrimary();
            if (p == null) return null;
            Expr zero;
            if (p.Type == VarType.REAL || p.Type == VarType.LREAL) {
                zero = new RealExpr(0.0, p.Type);
            } else {
                zero = new IntExpr(0, p.Type);
            }
            var op = sign == TokType.PLUS ? TokType.PLUS : TokType.MINUS;
            return new BinaryExpr(zero, op, p);
        }

        switch(cur.Type) {
            case TokType.INT:
                if (!long.TryParse(cur.Text, out var v)) {
                    AddError($"Invalid integer literal '{cur.Text}'");
                    return null;
                }
                Next();
                return new IntExpr(v, VarType.DINT);
                
            case TokType.REAL:
                if (!double.TryParse(cur.Text, out var d)) {
                    AddError($"Invalid floating point literal '{cur.Text}'");
                    return null;
                }
                Next();
                return new RealExpr(d);
                
            case TokType.IDENT:
                string n = cur.Text.ToUpperInvariant(); // Variablennamen in Großbuchstaben
                Next();
                Symbol? sym;
                if (!syms.TryGetValue(n, out sym)) {
                    AddError($"Undeclared variable '{n}'");
                    return null;
                }
                
                // Check for array access
                if (cur.Type == TokType.LBRACKET) {
                    if (sym.Type != VarType.ARRAY) {
                        AddError($"Cannot index non-array variable '{n}'");
                        return null;
                    }
                    
                    Next(); // consume [
                    var idx = ParseExpr();
                    if (idx == null) return null;
                    
                    if (!Expect(TokType.RBRACKET)) return null;
                    
                    ArrayType? arrayType;
                    if (!arrayTypes.TryGetValue(n, out arrayType)) {
                        AddError($"Internal error: Array type information missing for '{n}'");
                        return null;
                    }
                    
                    return new ArrayAccessExpr(n, idx, arrayType.ElementType);
                }
                
                return new VarExpr(n, sym.Type);
                
            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;
        }
    }

    public ArrayType? GetArrayType(string name) {
        ArrayType? type;
        return arrayTypes.TryGetValue(name, out type) ? type : null;
    }
}