ast.h

#ifndef AST_H
#define AST_H

#include "token.h"
#include <memory>
#include <vector>
#include <string>

namespace AR437 {

// Forward declarations
class ASTVisitor;

// Base AST Node
class ASTNode {
public:
    virtual ~ASTNode() = default;
    virtual void accept(ASTVisitor& visitor) = 0;
    virtual std::string toString() const = 0;
};

// Expression nodes
class Expression : public ASTNode {};

class NumberLiteral : public Expression {
public:
    double value;
    
    NumberLiteral(double val) : value(val) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class StringLiteral : public Expression {
public:
    std::string value;
    
    StringLiteral(const std::string& val) : value(val) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class BooleanLiteral : public Expression {
public:
    bool value;
    
    BooleanLiteral(bool val) : value(val) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class NullLiteral : public Expression {
public:
    NullLiteral() {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class Identifier : public Expression {
public:
    std::string name;
    
    Identifier(const std::string& n) : name(n) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class BinaryOperation : public Expression {
public:
    std::unique_ptr<Expression> left;
    TokenType operator_;
    std::unique_ptr<Expression> right;
    
    BinaryOperation(std::unique_ptr<Expression> l, TokenType op, std::unique_ptr<Expression> r)
        : left(std::move(l)), operator_(op), right(std::move(r)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class UnaryOperation : public Expression {
public:
    TokenType operator_;
    std::unique_ptr<Expression> operand;
    
    UnaryOperation(TokenType op, std::unique_ptr<Expression> expr)
        : operator_(op), operand(std::move(expr)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class Assignment : public Expression {
public:
    std::string variable;
    std::unique_ptr<Expression> value;
    
    Assignment(const std::string& var, std::unique_ptr<Expression> val)
        : variable(var), value(std::move(val)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class FunctionCall : public Expression {
public:
    std::string name;
    std::vector<std::unique_ptr<Expression>> arguments;
    
    FunctionCall(const std::string& n) : name(n) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

// Statement nodes
class Statement : public ASTNode {};

class ExpressionStatement : public Statement {
public:
    std::unique_ptr<Expression> expression;
    
    ExpressionStatement(std::unique_ptr<Expression> expr) : expression(std::move(expr)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class VariableDeclaration : public Statement {
public:
    std::string name;
    std::unique_ptr<Expression> initializer;
    
    VariableDeclaration(const std::string& n, std::unique_ptr<Expression> init = nullptr)
        : name(n), initializer(std::move(init)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class Block : public Statement {
public:
    std::vector<std::unique_ptr<Statement>> statements;
    
    Block() {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class IfStatement : public Statement {
public:
    std::unique_ptr<Expression> condition;
    std::unique_ptr<Statement> thenBranch;
    std::unique_ptr<Statement> elseBranch;
    
    IfStatement(std::unique_ptr<Expression> cond, std::unique_ptr<Statement> then,
                std::unique_ptr<Statement> else_ = nullptr)
        : condition(std::move(cond)), thenBranch(std::move(then)), elseBranch(std::move(else_)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class WhileStatement : public Statement {
public:
    std::unique_ptr<Expression> condition;
    std::unique_ptr<Statement> body;
    
    WhileStatement(std::unique_ptr<Expression> cond, std::unique_ptr<Statement> b)
        : condition(std::move(cond)), body(std::move(b)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class ReturnStatement : public Statement {
public:
    std::unique_ptr<Expression> value;
    
    ReturnStatement(std::unique_ptr<Expression> val = nullptr) : value(std::move(val)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class DisplayStatement : public Statement {
public:
    std::unique_ptr<Expression> expression;
    
    DisplayStatement(std::unique_ptr<Expression> expr) : expression(std::move(expr)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class ReadStatement : public Statement {
public:
    std::string variableName;
    
    ReadStatement(const std::string& varName) : variableName(varName) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class FunctionDeclaration : public Statement {
public:
    std::string name;
    std::vector<std::string> parameters;
    std::unique_ptr<Block> body;
    
    FunctionDeclaration(const std::string& n, const std::vector<std::string>& params,
                       std::unique_ptr<Block> b)
        : name(n), parameters(params), body(std::move(b)) {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

class Program : public ASTNode {
public:
    std::vector<std::unique_ptr<Statement>> statements;
    
    Program() {}
    void accept(ASTVisitor& visitor) override;
    std::string toString() const override;
};

// Visitor pattern for AST traversal
class ASTVisitor {
public:
    virtual ~ASTVisitor() = default;
    
    virtual void visit(NumberLiteral& node) = 0;
    virtual void visit(StringLiteral& node) = 0;
    virtual void visit(BooleanLiteral& node) = 0;
    virtual void visit(NullLiteral& node) = 0;
    virtual void visit(Identifier& node) = 0;
    virtual void visit(BinaryOperation& node) = 0;
    virtual void visit(UnaryOperation& node) = 0;
    virtual void visit(Assignment& node) = 0;
    virtual void visit(FunctionCall& node) = 0;
    virtual void visit(ExpressionStatement& node) = 0;
    virtual void visit(VariableDeclaration& node) = 0;
    virtual void visit(Block& node) = 0;
    virtual void visit(IfStatement& node) = 0;
    virtual void visit(WhileStatement& node) = 0;
    virtual void visit(ReturnStatement& node) = 0;
    virtual void visit(DisplayStatement& node) = 0;
    virtual void visit(ReadStatement& node) = 0;
    virtual void visit(FunctionDeclaration& node) = 0;
    virtual void visit(Program& node) = 0;
};

} // namespace AR437

#endif // AST_H