gutr_parser.cpp

#include "gutr_parser.hpp"
#include <iostream>
#include <sstream>
#include <cctype>
#include <vector>
#include <memory>
#include <stdexcept>
#include <fstream>
#include <unordered_set>

namespace uzaleat {

enum class TokenType {
    END, IDENT, NUMBER, STRING, KEYWORD,
    LPAREN, RPAREN, LBRACE, RBRACE, LBRACK, RBRACK,
    COMMA, SEMICOLON, COLON, DOT,
    PLUS, MINUS, STAR, SLASH, EQ, NE, LT, GT, LE, GE, BANG,
    AND, OR, ASSIGN, ARROW
};

struct Token {
    TokenType type = TokenType::END;
    std::string text;
    int line = 0, col = 0;
    Token() = default;
    Token(TokenType t, const std::string& txt, int l, int c) : type(t), text(txt), line(l), col(c) {}
};

class Lexer {
    std::string src;
    size_t pos = 0;
    int line = 1, col = 1;

    void advance() {
        if (pos < src.size()) {
            if (src[pos] == '\n') { line++; col = 1; } else { col++; }
            pos++;
        }
    }
    char peek() const { return pos < src.size() ? src[pos] : '\0'; }

    void skipWhitespace() { while (pos < src.size() && std::isspace(src[pos])) advance(); }
    void skipComment() {
        if (pos + 1 >= src.size()) return;
        if (src[pos] == '/' && src[pos+1] == '/') { while (pos < src.size() && src[pos] != '\n') advance(); if (pos < src.size()) advance(); }
        else if (src[pos] == '/' && src[pos+1] == '*') { advance(); advance(); while (pos + 1 < src.size() && !(src[pos] == '*' && src[pos+1] == '/')) advance(); if (pos + 1 < src.size()) { advance(); advance(); } }
    }

    Token readString() {
        int l = line, c = col; advance();
        std::string str;
        while (pos < src.size()) {
            if (src[pos] == '"') { advance(); return Token(TokenType::STRING, str, l, c); }
            if (src[pos] == '\\') {
                advance();
                if (pos < src.size()) {
                    switch (src[pos]) { case 'n': str += '\n'; break; case 't': str += '\t'; break; case 'r': str += '\r'; break; case '\\': str += '\\'; break; case '"': str += '"'; break; default: str += src[pos]; break; }
                    advance();
                }
            } else { str += src[pos]; advance(); }
        }
        throw std::runtime_error("unclosed string");
    }

    Token readNumber() {
        int l = line, c = col; size_t start = pos; bool has_dot = false;
        while (pos < src.size() && (std::isdigit(src[pos]) || (src[pos] == '.' && !has_dot))) { if (src[pos] == '.') has_dot = true; advance(); }
        if (pos < src.size() && (src[pos] == 'e' || src[pos] == 'E')) { advance(); if (pos < src.size() && (src[pos] == '+' || src[pos] == '-')) advance(); while (pos < src.size() && std::isdigit(src[pos])) advance(); }
        return Token(TokenType::NUMBER, src.substr(start, pos - start), l, c);
    }

    Token readIdent() {
        int l = line, c = col; size_t start = pos;
        while (pos < src.size() && (std::isalnum(src[pos]) || src[pos] == '_')) advance();
        std::string ident = src.substr(start, pos - start);
        static const std::unordered_set<std::string> keywords = {"fn", "sect", "export", "if", "else", "while", "for", "return", "break", "continue", "true", "false", "nil", "u64", "f64", "bool", "and", "or", "not", "let", "var"};
        if (keywords.find(ident) != keywords.end()) return Token(TokenType::KEYWORD, ident, l, c);
        return Token(TokenType::IDENT, ident, l, c);
    }

public:
    Lexer(const std::string& s) : src(s) { if (s.empty()) throw std::runtime_error("Empty source"); }
    Token next() {
        while (true) {
            skipWhitespace();
            if (pos + 1 < src.size() && src[pos] == '/' && (src[pos+1] == '/' || src[pos+1] == '*')) { skipComment(); continue; }
            if (pos >= src.size()) return Token(TokenType::END, "", line, col);
            char c = src[pos];
            if (c == '=' && pos+1 < src.size() && src[pos+1]=='=') { advance(); advance(); return Token(TokenType::EQ, "==", line, col-2); }
            if (c == '!' && pos+1 < src.size() && src[pos+1]=='=') { advance(); advance(); return Token(TokenType::NE, "!=", line, col-2); }
            if (c == '<' && pos+1 < src.size() && src[pos+1]=='=') { advance(); advance(); return Token(TokenType::LE, "<=", line, col-2); }
            if (c == '>' && pos+1 < src.size() && src[pos+1]=='=') { advance(); advance(); return Token(TokenType::GE, ">=", line, col-2); }
            if (c == '&' && pos+1 < src.size() && src[pos+1]=='&') { advance(); advance(); return Token(TokenType::AND, "&&", line, col-2); }
            if (c == '|' && pos+1 < src.size() && src[pos+1]=='|') { advance(); advance(); return Token(TokenType::OR, "||", line, col-2); }
            if (c == '-' && pos+1 < src.size() && src[pos+1]=='>') { advance(); advance(); return Token(TokenType::ARROW, "->", line, col-2); }
            switch (c) {
                case '(': advance(); return Token(TokenType::LPAREN, "(", line, col-1);
                case ')': advance(); return Token(TokenType::RPAREN, ")", line, col-1);
                case '{': advance(); return Token(TokenType::LBRACE, "{", line, col-1);
                case '}': advance(); return Token(TokenType::RBRACE, "}", line, col-1);
                case '[': advance(); return Token(TokenType::LBRACK, "[", line, col-1);
                case ']': advance(); return Token(TokenType::RBRACK, "]", line, col-1);
                case ',': advance(); return Token(TokenType::COMMA, ",", line, col-1);
                case ';': advance(); return Token(TokenType::SEMICOLON, ";", line, col-1);
                case ':': advance(); return Token(TokenType::COLON, ":", line, col-1);
                case '.': advance(); return Token(TokenType::DOT, ".", line, col-1);
                case '+': advance(); return Token(TokenType::PLUS, "+", line, col-1);
                case '-': advance(); return Token(TokenType::MINUS, "-", line, col-1);
                case '*': advance(); return Token(TokenType::STAR, "*", line, col-1);
                case '/': advance(); return Token(TokenType::SLASH, "/", line, col-1);
                case '=': advance(); return Token(TokenType::ASSIGN, "=", line, col-1);
                case '<': advance(); return Token(TokenType::LT, "<", line, col-1);
                case '>': advance(); return Token(TokenType::GT, ">", line, col-1);
                case '!': advance(); return Token(TokenType::BANG, "!", line, col-1);
                case '"': return readString();
            }
            if (std::isdigit(c) || (c == '.' && pos+1 < src.size() && std::isdigit(src[pos+1]))) return readNumber();
            if (std::isalpha(c) || c == '_') return readIdent();
            throw std::runtime_error(std::string("unexpected char: '") + c + "'");
        }
    }
};

class Parser {
    Lexer lexer;
    Token current;
    bool has_peek = false;
    Token peek_token;

    Token peek() { if (!has_peek) { peek_token = lexer.next(); has_peek = true; } return peek_token; }
    void next() { if (has_peek) { current = peek_token; has_peek = false; } else { current = lexer.next(); } }
    bool match(TokenType t) { if (current.type == t) { next(); return true; } return false; }
    void expect(TokenType t) { if (current.type != t) throw std::runtime_error("unexpected token"); next(); }

public:
    Parser(const std::string& src) : lexer(src) { next(); }

    std::shared_ptr<GUTRBlockStmt> parse() {
        auto prog = std::make_shared<GUTRBlockStmt>();
        while (current.type != TokenType::END) {
            auto stmt = parseStatement();
            if (stmt) prog->statements.push_back(stmt);
            if (current.type == TokenType::RBRACE) break;
        }
        return prog;
    }

private:
    std::shared_ptr<GUTRExpr> parseStatement() {
        if (current.type == TokenType::RBRACE || current.type == TokenType::END) return nullptr;
        if (current.type == TokenType::KEYWORD) {
            std::string kw = current.text;
            if (kw == "fn") return parseFunctionDef();
            if (kw == "sect") return parseSectDef();
            if (kw == "export") { next(); return parseStatement(); }
            if (kw == "if") return parseIf();
            if (kw == "while") return parseWhile();
            if (kw == "for") return parseFor();
            if (kw == "return") return parseReturn();
            if (kw == "break") { next(); return std::make_shared<GUTRBreakStmt>(); }
            if (kw == "continue") { next(); return std::make_shared<GUTRContinueStmt>(); }
            if (kw == "let" || kw == "var") next();
        }
        if (current.type == TokenType::LBRACE) {
            next();
            auto block = std::make_shared<GUTRBlockStmt>();
            while (current.type != TokenType::RBRACE && current.type != TokenType::END) {
                auto stmt = parseStatement();
                if (stmt) block->statements.push_back(stmt);
            }
            expect(TokenType::RBRACE);
            return block;
        }
        if (current.type == TokenType::IDENT) {
            Token save = current; next();
            if (current.type == TokenType::ASSIGN) {
                next();
                auto expr = parseExpression();
                auto assign = std::make_shared<GUTRAssignExpr>();
                assign->var_name = save.text;
                assign->value = expr;
                if (current.type == TokenType::SEMICOLON) next();
                return assign;
            }
            current = save;
        }
        auto expr = parseExpression();
        if (current.type == TokenType::SEMICOLON) next();
        return expr;
    }

    std::shared_ptr<GUTRExpr> parseExpression() { return parseBinary(0); }

    int precedence(TokenType t) {
        switch (t) {
            case TokenType::OR: return 1;
            case TokenType::AND: return 2;
            case TokenType::EQ: case TokenType::NE: return 3;
            case TokenType::LT: case TokenType::GT: case TokenType::LE: case TokenType::GE: return 4;
            case TokenType::PLUS: case TokenType::MINUS: return 5;
            case TokenType::STAR: case TokenType::SLASH: return 6;
            default: return 0;
        }
    }

    std::shared_ptr<GUTRExpr> parseBinary(int min_prec) {
        auto left = parseUnary();
        if (!left) return nullptr;
        while (true) {
            if (current.type == TokenType::RBRACE || current.type == TokenType::END || current.type == TokenType::SEMICOLON || current.type == TokenType::COMMA || current.type == TokenType::RPAREN || current.type == TokenType::RBRACK) break;
            int prec = precedence(current.type);
            if (prec < min_prec) break;
            TokenType op = current.type; next();
            auto right = parseBinary(prec + 1);
            if (!right) throw std::runtime_error("Missing right operand");
            auto bin = std::make_shared<GUTRBinaryOpExpr>();
            switch (op) {
                case TokenType::PLUS: bin->op = GUTRBinaryOpExpr::ADD; break;
                case TokenType::MINUS: bin->op = GUTRBinaryOpExpr::SUB; break;
                case TokenType::STAR: bin->op = GUTRBinaryOpExpr::MUL; break;
                case TokenType::SLASH: bin->op = GUTRBinaryOpExpr::DIV; break;
                case TokenType::LT: bin->op = GUTRBinaryOpExpr::LT; break;
                case TokenType::GT: bin->op = GUTRBinaryOpExpr::GT; break;
                case TokenType::EQ: bin->op = GUTRBinaryOpExpr::EQ; break;
                case TokenType::NE: bin->op = GUTRBinaryOpExpr::NE; break;
                case TokenType::LE: bin->op = GUTRBinaryOpExpr::LE; break;
                case TokenType::GE: bin->op = GUTRBinaryOpExpr::GE; break;
                case TokenType::AND: bin->op = GUTRBinaryOpExpr::AND; break;
                case TokenType::OR: bin->op = GUTRBinaryOpExpr::OR; break;
                default: break;
            }
            bin->left = left; bin->right = right;
            left = bin;
        }
        return left;
    }

    std::shared_ptr<GUTRExpr> parseUnary() {
        if (match(TokenType::MINUS)) { auto u = std::make_shared<GUTRUnaryOpExpr>(); u->op = GUTRUnaryOpExpr::NEG; u->arg = parsePrimary(); return u; }
        if (match(TokenType::BANG)) { auto u = std::make_shared<GUTRUnaryOpExpr>(); u->op = GUTRUnaryOpExpr::NOT; u->arg = parsePrimary(); return u; }
        return parsePrimary();
    }

    std::shared_ptr<GUTRExpr> parsePrimary() {
        if (current.type == TokenType::RBRACE || current.type == TokenType::END || current.type == TokenType::SEMICOLON || current.type == TokenType::COMMA || current.type == TokenType::RPAREN || current.type == TokenType::RBRACK) return nullptr;
        if (match(TokenType::LPAREN)) { auto expr = parseExpression(); expect(TokenType::RPAREN); return expr; }
        if (current.type == TokenType::NUMBER) { std::string num = current.text; next(); return std::make_shared<GUTRLiteralExpr>(num.find('.') != std::string::npos ? GUTRValue::real(std::stod(num)) : GUTRValue::integer(std::stoll(num))); }
        if (current.type == TokenType::STRING) { std::string val = current.text; next(); return std::make_shared<GUTRLiteralExpr>(GUTRValue::string(val)); }
        if (current.type == TokenType::KEYWORD) {
            std::string kw = current.text;
            if (kw == "true") { next(); return std::make_shared<GUTRLiteralExpr>(GUTRValue::boolean(true)); }
            if (kw == "false") { next(); return std::make_shared<GUTRLiteralExpr>(GUTRValue::boolean(false)); }
            if (kw == "nil") { next(); return std::make_shared<GUTRLiteralExpr>(GUTRValue::nil()); }
            if (kw == "list" || kw == "tensor" || kw == "print") { next(); if (current.type == TokenType::LPAREN) return parseCall(std::make_shared<GUTRVarExpr>(kw)); }
        }
        if (current.type == TokenType::IDENT) {
            std::string id = current.text; Token save = current; next();
            if (current.type == TokenType::ASSIGN) { next(); auto expr = parseExpression(); auto assign = std::make_shared<GUTRAssignExpr>(); assign->var_name = id; assign->value = expr; if (current.type == TokenType::SEMICOLON) next(); return assign; }
            if (current.type == TokenType::LPAREN) return parseCall(std::make_shared<GUTRVarExpr>(id));
            if (current.type == TokenType::LBRACK) { auto idx = std::make_shared<GUTRIndexExpr>(); idx->object = std::make_shared<GUTRVarExpr>(id); next(); idx->index = parseExpression(); expect(TokenType::RBRACK); return idx; }
            return std::make_shared<GUTRVarExpr>(id);
        }
        if (match(TokenType::LBRACK)) {
            auto arr = std::make_shared<GUTRArrayLiteralExpr>();
            if (current.type != TokenType::RBRACK) {
                while (true) { auto elem = parseExpression(); if (elem) arr->elements.push_back(elem); if (current.type == TokenType::COMMA) next(); else break; }
            }
            expect(TokenType::RBRACK);
            return arr;
        }
        throw std::runtime_error("Unexpected token");
    }

    std::shared_ptr<GUTRExpr> parseCall(std::shared_ptr<GUTRExpr> callee) {
        auto call = std::make_shared<GUTRCallExpr>(); call->callee = callee; next();
        while (current.type != TokenType::RPAREN && current.type != TokenType::END) {
            auto arg = parseExpression(); if (arg) call->args.push_back(arg);
            if (current.type == TokenType::COMMA) next(); else if (current.type != TokenType::RPAREN) throw std::runtime_error("Expected ',' or ')'");
        }
        expect(TokenType::RPAREN);
        return call;
    }

    std::shared_ptr<GUTRFunctionDefExpr> parseFunctionDef() {
        next();
        if (current.type != TokenType::IDENT) throw std::runtime_error("Expected function name");
        std::string name = current.text; next();
        expect(TokenType::LPAREN);
        std::vector<std::string> params;
        if (current.type != TokenType::RPAREN) {
            while (true) {
                if (current.type == TokenType::KEYWORD) next();
                if (current.type != TokenType::IDENT) throw std::runtime_error("Expected parameter name");
                params.push_back(current.text); next();
                if (current.type == TokenType::COMMA) next(); else break;
            }
        }
        expect(TokenType::RPAREN);
        if (current.type == TokenType::ARROW) { next(); if (current.type == TokenType::LPAREN) { next(); int d=1; while(current.type != TokenType::END && d>0) { if(current.type==TokenType::LPAREN)d++; else if(current.type==TokenType::RPAREN)d--; next(); } } else next(); }
        auto body = std::dynamic_pointer_cast<GUTRBlockStmt>(parseStatement());
        if (!body) body = std::make_shared<GUTRBlockStmt>();
        auto def = std::make_shared<GUTRFunctionDefExpr>(); def->name = name; def->params = params; def->body = body;
        return def;
    }

    std::shared_ptr<GUTRSectDefExpr> parseSectDef() {
        next();
        if (current.type != TokenType::IDENT) throw std::runtime_error("Expected sect name");
        std::string name = current.text; next();
        expect(TokenType::LBRACE);
        std::vector<std::pair<std::string, std::shared_ptr<GUTRExpr>>> fields;
        while (current.type != TokenType::RBRACE && current.type != TokenType::END) {
            if (current.type != TokenType::IDENT) throw std::runtime_error("Expected field name");
            std::string fname = current.text; next();
            expect(TokenType::ASSIGN);
            auto expr = parseExpression();
            if (expr) fields.emplace_back(fname, expr);
            if (current.type == TokenType::COMMA || current.type == TokenType::SEMICOLON) next();
        }
        expect(TokenType::RBRACE);
        auto sect = std::make_shared<GUTRSectDefExpr>(); sect->name = name; sect->field_inits = fields;
        return sect;
    }

    std::shared_ptr<GUTRIfStmt> parseIf() {
        next();
        auto cond = parseExpression();
        auto then = parseStatement();
        std::shared_ptr<GUTRExpr> els;
        if (current.type == TokenType::KEYWORD && current.text == "else") { next(); els = parseStatement(); }
        auto stmt = std::make_shared<GUTRIfStmt>(); stmt->condition = cond; stmt->then_branch = then; stmt->else_branch = els;
        return stmt;
    }

    std::shared_ptr<GUTRWhileStmt> parseWhile() {
        next();
        auto cond = parseExpression();
        auto body = parseStatement();
        auto stmt = std::make_shared<GUTRWhileStmt>(); stmt->condition = cond; stmt->body = body;
        return stmt;
    }

    std::shared_ptr<GUTRForStmt> parseFor() {
        next();
        if (current.type != TokenType::IDENT) throw std::runtime_error("Expected loop var");
        std::string var = current.text; next();
        expect(TokenType::ASSIGN);
        auto start = parseExpression();
        expect(TokenType::COMMA);
        auto end = parseExpression();
        std::shared_ptr<GUTRExpr> step;
        if (current.type == TokenType::COMMA) { next(); step = parseExpression(); }
        auto body = parseStatement();
        auto stmt = std::make_shared<GUTRForStmt>(); stmt->var_name = var; stmt->start_expr = start; stmt->end_expr = end; stmt->step_expr = step ? step : std::make_shared<GUTRLiteralExpr>(GUTRValue::integer(1)); stmt->body = body;
        return stmt;
    }

    std::shared_ptr<GUTRReturnStmt> parseReturn() {
        next();
        auto stmt = std::make_shared<GUTRReturnStmt>();
        if (current.type != TokenType::RBRACE && current.type != TokenType::END && current.type != TokenType::SEMICOLON) stmt->value = parseExpression();
        return stmt;
    }
};

bool GUTRParser::parse(const std::string& source, GUTRProgram& program) {
    if (source.empty()) return false;
    try {
        Parser parser(source);
        auto ast = parser.parse();
        if (!ast) return false;

        Compiler compiler;
        ast->compile(compiler);
        compiler.emit(OpCode::LOAD_CONST);
        compiler.emit_int(compiler.add_constant(GUTRValue::nil()));
        compiler.emit(OpCode::RETURN);

        VM vm(program.global_ctx);
        vm.run(reinterpret_cast<GUTRUserFunction*>(&compiler), {});

        auto get_func = [&](const std::string& name) -> std::shared_ptr<GUTRUserFunction> {
            try {
                GUTRValue v = program.global_ctx->get(name);
                if (v.type != ValueType::FUNCTION) throw std::runtime_error("not a function");
                return std::dynamic_pointer_cast<GUTRUserFunction>(v.func_val);
            } catch (...) { throw std::runtime_error("Required function '" + name + "' not found"); }
        };

        program.init_func = get_func("init");
        program.train_step_func = get_func("train_step");
        program.save_func = get_func("save");
        program.load_func = get_func("load");
        program.generate_func = get_func("generate");
        return true;
    } catch (const std::exception& e) {
        std::cerr << "GUTR Parse error: " << e.what() << std::endl;
        return false;
    }
}

} // namespace uzaleat